Update app.py

app.py

@@ -17,8 +17,15 @@ def load_model(custom_model):
     pipe.to("cuda")
     return "Model loaded!"
 
-def infer (prompt):
-    image = pipe(prompt=prompt, num_inference_steps=50).images[0]
+def infer (prompt, inf_steps, guidance_scale, seed, lora_weigth, progress=gr.Progress(track_tqdm=True)):
+    generator = torch.Generator(device="cuda").manual_seed(seed)
+    image = pipe(
+        prompt=prompt, 
+        num_inference_steps=inf_steps,
+        guidance_scale = float(guidance_scale),
+        generator=generator,
+        cross_attention_kwargs={"scale": float(lora_weight)}
+    ).images[0]
     return image
 
 css = """
@@ -37,6 +44,34 @@ with gr.Blocks(css=css) as demo:
             load_model_btn = gr.Button("Load my model")
         
         prompt_in = gr.Textbox(label="Prompt")   
+        inf_steps = gr.Slider(
+            label="Inference steps",
+            minimum=12,
+            maximum=50,
+            step=1,
+            value=25
+        )
+        guidance_scale = gr.Slider(
+            label="Guidance scale",
+            minimum=0.1,
+            maximum=0.9,
+            step=0.1,
+            value=7.5
+        )
+        seed = gr.Slider(
+            label="Seed",
+            minimum=0,
+            maximum=500000,
+            step=1,
+            value=42
+        )
+        lora_weight = gr.Slider(
+            label="LoRa weigth",
+            minimum=0.0,
+            maximum=10.0,
+            step=0.01,
+            value=0.9
+        )
         submit_btn = gr.Button("Submit")
         image_out = gr.Image(label="Image output")
 
@@ -47,7 +82,7 @@ with gr.Blocks(css=css) as demo:
     )
     submit_btn.click(
         fn = infer,
-        inputs = [prompt_in],
+        inputs = [prompt_in, inf_steps, guidance_scale, seed, lora_weight],
         outputs = [image_out]
     )