Update app_lora1.py

app_lora1.py

@@ -146,40 +146,142 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale):
 # ============================================================
 # GRADIO UI
 # ============================================================
+css = """
+#container {
+    max-width: 1200px;
+    margin: auto;
+}
+
+.section {
+    margin-top: 16px;
+    margin-bottom: 8px;
+}
+
+.generate-btn {
+    background: linear-gradient(90deg, #4b6cb7, #182848) !important;
+    color: white !important;
+    border-radius: 8px !important;
+    font-weight: 600;
+    height: 44px;
+}
+
+.secondary-btn {
+    border-radius: 8px !important;
+    height: 44px;
+}
+
+input, textarea {
+    border-radius: 8px !important;
+}
+"""
+
+with gr.Blocks(
+    title="Z-Image-Turbo (LoRA Safe)",
+    css=css,
+) as demo:
+
+    gr.Markdown(
+        """
+        # 🎨 Z-Image-Turbo  
+        **Runtime LoRA · Safe Mode · Production UI**
+        """,
+        elem_classes="section"
+    )
 
-with gr.Blocks(title="Z-Image-Turbo (LoRA Safe)") as demo:
-    gr.Markdown("# 🎨 Z-Image-Turbo — Runtime LoRA (SAFE MODE)")
+    with gr.Column(elem_id="container"):
+
+        # ===============================
+        # PROMPT
+        # ===============================
+        with gr.Column(elem_classes="section"):
+            prompt = gr.Textbox(
+                label="Prompt",
+                value="boat in ocean",
+                lines=3,
+                placeholder="Describe what you want to generate…",
+            )
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Textbox(label="Prompt", value="boat in ocean")
-            height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
+        # ===============================
+        # PARAMETERS
+        # ===============================
+        with gr.Row(elem_classes="section"):
             width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
+            height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
+
+        with gr.Row():
             steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
             guidance = gr.Slider(0, 10, value=0.0, step=0.5, label="Guidance Scale")
-            seed = gr.Number(value=42, label="Seed")
-            run_btn = gr.Button("🚀 Generate")
+            seed = gr.Number(value=42, label="Seed", precision=0)
+
+        # ===============================
+        # ACTION BUTTON
+        # ===============================
+        with gr.Row(elem_classes="section"):
+            run_btn = gr.Button("🚀 Generate Image", elem_classes="generate-btn")
+
+        # ===============================
+        # OUTPUT
+        # ===============================
+        with gr.Row(elem_classes="section"):
+            final_image = gr.Image(
+                label="Final Image",
+                height=420,
+            )
 
-        with gr.Column(scale=1):
-            final_image = gr.Image(label="Final Image")
-            gallery = gr.Gallery(label="Steps", columns=4, height=256)
+        gallery = gr.Gallery(
+            label="Generation Steps",
+            columns=4,
+            height=260,
+        )
 
-    gr.Markdown("## 🧩 LoRA Controls")
+        # ===============================
+        # LORA CONTROLS
+        # ===============================
+        gr.Markdown("## 🧩 LoRA Controls", elem_classes="section")
 
-    with gr.Row():
-        lora_repo = gr.Textbox(label="LoRA Repo", value=DEFAULT_LORA_REPO)
-        lora_dropdown = gr.Dropdown(label="LoRA File", choices=[])
-        refresh_btn = gr.Button("🔄 Refresh")
+        with gr.Row():
+            lora_repo = gr.Textbox(
+                label="LoRA Repository",
+                value=DEFAULT_LORA_REPO,
+                placeholder="username/repo",
+            )
+
+        with gr.Row():
+            lora_dropdown = gr.Dropdown(
+                label="Available LoRA Files",
+                choices=[],
+                interactive=True,
+            )
 
-    logs_box = gr.Textbox(label="Logs", lines=18)
+        with gr.Row():
+            refresh_btn = gr.Button("🔄 Refresh LoRA List", elem_classes="secondary-btn")
+            clear_lora_btn = gr.Button("❌ Clear LoRA", elem_classes="secondary-btn")
+
+        # ===============================
+        # LOGS
+        # ===============================
+        logs_box = gr.Textbox(
+            label="Logs",
+            lines=14,
+            interactive=False,
+        )
 
-    # ---- Callbacks ----
+    # ==================================================
+    # CALLBACKS
+    # ==================================================
 
     def refresh_loras(repo):
         files = list_loras_from_repo(repo)
-        return gr.update(choices=files, value=files[0] if files else None)
+        return gr.update(
+            choices=files,
+            value=files[0] if files else None
+        )
 
-    refresh_btn.click(refresh_loras, inputs=[lora_repo], outputs=[lora_dropdown])
+    refresh_btn.click(
+        refresh_loras,
+        inputs=[lora_repo],
+        outputs=[lora_dropdown],
+    )
 
     def select_lora(lora_file, repo):
         global CURRENT_LORA_FILE, CURRENT_LORA_REPO
@@ -193,6 +295,24 @@ with gr.Blocks(title="Z-Image-Turbo (LoRA Safe)") as demo:
         outputs=[logs_box],
     )
 
+    def clear_lora():
+        global CURRENT_LORA_FILE, CURRENT_LORA_REPO
+        CURRENT_LORA_FILE = None
+        CURRENT_LORA_REPO = None
+        try:
+            pipe.unload_lora_weights()
+        except Exception:
+            pass
+        return (
+            gr.update(value=None),
+            "🧹 LoRA cleared. Next generation will be base model only."
+        )
+
+    clear_lora_btn.click(
+        clear_lora,
+        outputs=[lora_dropdown, logs_box],
+    )
+
     run_btn.click(
         generate_image,
         inputs=[prompt, height, width, steps, seed, guidance],
@@ -200,3 +320,57 @@ with gr.Blocks(title="Z-Image-Turbo (LoRA Safe)") as demo:
     )
 
 demo.launch()
+
+# with gr.Blocks(title="Z-Image-Turbo (LoRA Safe)") as demo:
+#     gr.Markdown("# 🎨 Z-Image-Turbo — Runtime LoRA (SAFE MODE)")
+
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             prompt = gr.Textbox(label="Prompt", value="boat in ocean")
+#             height = gr.Slider(256, 2048, value=1024, step=8, label="Height")
+#             width = gr.Slider(256, 2048, value=1024, step=8, label="Width")
+#             steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
+#             guidance = gr.Slider(0, 10, value=0.0, step=0.5, label="Guidance Scale")
+#             seed = gr.Number(value=42, label="Seed")
+#             run_btn = gr.Button("🚀 Generate")
+
+#         with gr.Column(scale=1):
+#             final_image = gr.Image(label="Final Image")
+#             gallery = gr.Gallery(label="Steps", columns=4, height=256)
+
+#     gr.Markdown("## 🧩 LoRA Controls")
+
+#     with gr.Row():
+#         lora_repo = gr.Textbox(label="LoRA Repo", value=DEFAULT_LORA_REPO)
+#         lora_dropdown = gr.Dropdown(label="LoRA File", choices=[])
+#         refresh_btn = gr.Button("🔄 Refresh")
+
+#     logs_box = gr.Textbox(label="Logs", lines=18)
+
+#     # ---- Callbacks ----
+
+#     def refresh_loras(repo):
+#         files = list_loras_from_repo(repo)
+#         return gr.update(choices=files, value=files[0] if files else None)
+
+#     refresh_btn.click(refresh_loras, inputs=[lora_repo], outputs=[lora_dropdown])
+
+#     def select_lora(lora_file, repo):
+#         global CURRENT_LORA_FILE, CURRENT_LORA_REPO
+#         CURRENT_LORA_FILE = lora_file
+#         CURRENT_LORA_REPO = repo
+#         return f"🧩 Selected LoRA: {lora_file}"
+
+#     lora_dropdown.change(
+#         select_lora,
+#         inputs=[lora_dropdown, lora_repo],
+#         outputs=[logs_box],
+#     )
+
+#     run_btn.click(
+#         generate_image,
+#         inputs=[prompt, height, width, steps, seed, guidance],
+#         outputs=[final_image, gallery, logs_box],
+#     )
+
+demo.launch()