Update app_lora.py

app_lora.py

@@ -973,6 +973,15 @@ import os
 # -------------------------
 # Helper: Recursive LoRA listing
 # -------------------------
+from huggingface_hub import list_repo_files
+import gradio as gr
+from PIL import Image
+
+
+
+# ----------------------------
+# LIST LoRA FILES HELPER
+# ----------------------------
 def list_loras_from_repo(repo_id: str):
     """
     List all .safetensors files in a Hugging Face repo, including subfolders.
@@ -986,9 +995,9 @@ def list_loras_from_repo(repo_id: str):
         log(f"❌ Failed to list repo files: {e}")
         return []
 
-# -------------------------
-# Gradio UI
-# -------------------------
+# ----------------------------
+# GRADIO UI
+# ----------------------------
 with gr.Blocks(title="Z-Image-Turbo") as demo:
     gr.Markdown("# 🎨 Z-Image-Turbo (LoRA-enabled UI)")
 
@@ -996,8 +1005,6 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
     # MAIN TABS
     # =========================
     with gr.Tabs():
-
-        # -------- Image Tab --------
         with gr.TabItem("Image & Latents"):
             with gr.Row():
                 with gr.Column(scale=1):
@@ -1012,7 +1019,6 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
                     final_image = gr.Image(label="Final Image")
                     latent_gallery = gr.Gallery(label="Latent Steps", columns=4, height=256, preview=True)
 
-        # -------- Logs Tab --------
         with gr.TabItem("Logs"):
             logs_box = gr.Textbox(label="Logs", lines=25, interactive=False)
 
@@ -1020,7 +1026,6 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
     # LoRA CONTROLS
     # =========================
     gr.Markdown("## 🧩 LoRA Controls")
-
     with gr.Row():
         lora_repo = gr.Textbox(label="LoRA Repo (HF)", value="rahul7star/ZImageLora", placeholder="username/repo")
         lora_file = gr.Dropdown(label="LoRA file (.safetensors)", choices=[])
@@ -1034,9 +1039,6 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
     # =========================
     # CALLBACKS
     # =========================
-     ------------------------
-    # CALLBACKS
-    # ------------------------
     def refresh_lora_list(repo_name):
         files = list_loras_from_repo(repo_name)
         if not files:
@@ -1047,6 +1049,8 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
 
     refresh_lora_btn.click(refresh_lora_list, inputs=[lora_repo], outputs=[lora_file])
 
+    loaded_loras = {}  # keep track of loaded adapters
+
     def apply_lora(repo_name, lora_filename, strength):
         global pipe, loaded_loras
         if pipe is None:
@@ -1069,6 +1073,8 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
             log(f"❌ Failed to apply LoRA: {e}")
             return f"Failed: {e}"
 
+    apply_lora_btn.click(apply_lora, inputs=[lora_repo, lora_file, lora_strength], outputs=[logs_box])
+
     def clear_lora():
         global pipe
         if pipe is None:
@@ -1093,6 +1099,4 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
     )
 
 
-
-
 demo.launch()