22

app.py

@@ -68,14 +68,16 @@ def generate_lora(
     ulid = str(ULID()).lower()[:12]
     print(f"ulid: {ulid}")
 
+    if not images:
+        print("images is empty.")
     # Load images 
-    images = [
-        Image.open("examples/style/1/0.jpg"),
-        Image.open("examples/style/1/1.jpg"),
-        Image.open("examples/style/1/2.jpg"),
-        Image.open("examples/style/1/3.jpg"),
-        Image.open("examples/style/1/4.jpg"),
-    ]
+    # images = [
+    #     Image.open("examples/style/1/0.jpg"),
+    #     Image.open("examples/style/1/1.jpg"),
+    #     Image.open("examples/style/1/2.jpg"),
+    #     Image.open("examples/style/1/3.jpg"),
+    #     Image.open("examples/style/1/4.jpg"),
+    # ]
 
 
     # Model inference
@@ -88,15 +90,18 @@ def generate_lora(
 
     save_file(lora, lora_path)
 
-    return lora_path, gr.update(visible=True, value=lora_path), gr.update(visible=True)
+    return lora_name, gr.update(visible=True, value=lora_path), gr.update(visible=True)
 
 @spaces.GPU
 def generate_image(
+    lora_name,
     prompt,
-    negative_prompt,
+    negative_prompt="blurry ugly bad",
+    width=1024,
+    height=1024,
     seed=42,
     randomize_seed=True,
-    guidance_scale=1.5,
+    guidance_scale=3.5,
     num_inference_steps=8,
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -119,6 +124,9 @@ h3{
     text-align: center;
     display:block;
 }
+#imagen-container {
+    padding: 12px;
+}
 """
 
 
@@ -131,7 +139,8 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 input_images = gr.Gallery(
-                    label="Generated images", 
+                    label="Input images", 
+                    file_types=["image"],
                     show_label=False, 
                     elem_id="gallery", 
                     columns=2, 
@@ -141,9 +150,9 @@ with gr.Blocks() as demo:
                 lora_button = gr.Button("Generate LoRA", variant="primary")
 
             with gr.Column():
-                lora_path = gr.Textbox(label="Generated LoRA path",lines=2, interactive=False)
+                lora_name = gr.Textbox(label="Generated LoRA path",lines=2, interactive=False)
                 lora_download = gr.DownloadButton(label=f"Download LoRA", visible=False)
-        with gr.Column(visible=False) as imagen_container:
+        with gr.Column(visible=False, elem_classes='imagen-container') as imagen_container:
             gr.Markdown("### After your LoRA is ready, you can try generate image here.")
             with gr.Row():
                 with gr.Column():
@@ -158,47 +167,52 @@ with gr.Blocks() as demo:
                     
                     imagen_button = gr.Button("Generate Image", variant="primary")    
                     with gr.Accordion("Advanced Settings", open=False):
-                        
                         negative_prompt = gr.Textbox(
                             label="Negative prompt",
                             lines=2,
                             container=False,
                             placeholder="Enter your negative prompt",
-                            value="blurry, ugly, bad"
+                            value="blurry ugly bad"
+                        )
+                        num_inference_steps = gr.Slider(
+                            label="Steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=20,
                         )
                         with gr.Row():
-                            num_inference_steps = gr.Slider(
-                                label="Steps",
-                                minimum=1,
-                                maximum=30,
-                                step=1,
-                                value=9,
-                            )
-                            control_context_scale = gr.Slider(
-                                label="Context scale",
-                                minimum=0.0,
-                                maximum=1.0,
-                                step=0.01,
-                                value=0.75,
+                            width = gr.Slider(
+                                label="Width",
+                                minimum=512,
+                                maximum=1280,
+                                step=32,
+                                value=768, 
                             )
 
+                            height = gr.Slider(
+                                label="Height",
+                                minimum=512,
+                                maximum=1280,
+                                step=32,
+                                value=1024,
+                            )
                         with gr.Row():
+                            seed = gr.Slider(
+                                label="Seed",
+                                minimum=0,
+                                maximum=MAX_SEED,
+                                step=1,
+                                value=42,
+                            )
                             guidance_scale = gr.Slider(
                                 label="Guidance scale",
                                 minimum=0.0,
                                 maximum=10.0,
                                 step=0.1,
-                                value=1.0,
+                                value=3.5,
                             )
-
-                        seed = gr.Slider(
-                            label="Seed",
-                            minimum=0,
-                            maximum=MAX_SEED,
-                            step=1,
-                            value=42,
-                        )
-                        randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
+                            randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
 
                 with gr.Column():
                     output_image = gr.Image(label="Generated image", show_label=False)
@@ -211,7 +225,22 @@ with gr.Blocks() as demo:
         inputs=[
             input_images
         ],
-        outputs=[lora_path, lora_download, imagen_container],
+        outputs=[lora_name, lora_download, imagen_container],
+    )
+    imagen_button.click(
+        fn=generate_image,
+        inputs=[
+            lora_name,
+            prompt,
+            negative_prompt,
+            width,
+            height,
+            seed,
+            randomize_seed,
+            guidance_scale,
+            num_inference_steps,
+        ],
+        outputs=[lora_name, lora_download, imagen_container],
     )