Update app_lora.py

app_lora.py

@@ -687,7 +687,73 @@ def generate_image_all_latents(prompt, height, width, steps, seed, guidance_scal
         yield placeholder, latent_gallery, LOGS
 
 @spaces.GPU
-def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
+def generate_image(prompt, height, width, steps, seed, guidance_scale=7.5):
+    LOGS = []
+    device = "cuda"
+    generator = torch.Generator(device).manual_seed(int(seed))
+    
+    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
+    latent_gallery = []
+
+    try:
+        # -------------------
+        # Prepare latents
+        # -------------------
+        batch_size = 1
+        if hasattr(pipe, "vae") and hasattr(pipe.vae, "config"):
+            num_channels = pipe.vae.config.latent_channels
+        else:
+            num_channels = 4
+
+        latents = torch.randn(
+            (batch_size, num_channels, height // 8, width // 8),
+            generator=generator,
+            device=device,
+        )
+
+        # -------------------
+        # Encode prompt
+        # -------------------
+        text_embeddings = pipe._encode_prompt(prompt)
+
+        # -------------------
+        # Scheduler loop
+        # -------------------
+        num_previews = min(10, steps)
+        preview_indices = torch.linspace(0, steps - 1, num_previews).long()
+
+        for t_idx in range(steps):
+            t = pipe.scheduler.timesteps[t_idx]
+
+            with torch.no_grad():
+                latents = pipe.unet(latents, t, encoder_hidden_states=text_embeddings).sample
+                latents = pipe.scheduler.step(latents, t, latents).prev_sample
+
+            if t_idx in preview_indices:
+                try:
+                    decoded = pipe.decode_latents(latents)
+                    latent_gallery.append(decoded)
+                except Exception as e:
+                    LOGS.append(f"⚠️ Preview decode failed: {e}")
+                    latent_gallery.append(placeholder)
+
+            yield None, latent_gallery[-5:], LOGS
+
+        # -------------------
+        # Final image
+        # -------------------
+        final_img = pipe.decode_latents(latents)
+        latent_gallery.append(final_img)
+        LOGS.append("✅ Pipeline generation completed successfully.")
+        yield final_img, latent_gallery[-5:] + [final_img], LOGS
+
+    except Exception as e:
+        LOGS.append(f"❌ Generation failed: {e}")
+        latent_gallery.append(placeholder)
+        yield placeholder, latent_gallery[-5:] + [placeholder], LOGS
+
+
+def generate_image1(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
     device = "cuda"
     cpu_device = "cpu"
@@ -975,12 +1041,12 @@ loaded_loras = {}
 with gr.Blocks(title="Z-Image-Turbo") as demo:
     gr.Markdown("# 🎨 Z-Image-Turbo (LoRA-enabled UI)")
 
-    # =========================
-    # MAIN TABS
-    # =========================
+    # -------------------------
+    # Tabs
+    # -------------------------
     with gr.Tabs():
 
-        # -------- Image Tab --------
+        # -------- Image & Latents --------
         with gr.TabItem("Image & Latents"):
             with gr.Row():
                 with gr.Column(scale=1):
@@ -995,17 +1061,17 @@ 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 --------
+        # -------- Logs --------
         with gr.TabItem("Logs"):
             logs_box = gr.Textbox(label="Logs", lines=25, interactive=False)
 
-    # =========================
-    # LoRA CONTROLS
-    # =========================
+    # -------------------------
+    # LoRA Controls
+    # -------------------------
     gr.Markdown("## 🧩 LoRA Controls")
 
     with gr.Row():
-        lora_repo = gr.Textbox(label="LoRA Repo (HF)", value="rahul7star/ZImageLora", placeholder="username/repo")
+        lora_repo = gr.Textbox(label="LoRA Repo (HF)", value="rahul7star/ZImageLora")
         lora_file = gr.Dropdown(label="LoRA file (.safetensors)", choices=[])
         lora_strength = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="LoRA strength")
 
@@ -1014,9 +1080,9 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
         apply_lora_btn = gr.Button("✅ Apply LoRA")
         clear_lora_btn = gr.Button("❌ Clear LoRA")
 
-    # =========================
-    # CALLBACKS
-    # =========================
+    # -------------------------
+    # Callbacks
+    # -------------------------
     def refresh_lora_list(repo_name):
         files = list_loras_from_repo(repo_name)
         if not files:
@@ -1034,17 +1100,13 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
         if not lora_filename:
             return "⚠️ No LoRA file selected"
 
-        # Remove .safetensors and slashes for adapter name (safe for PEFT)
-        adapter_name = f"ui_lora_{lora_filename.replace('/', '_').replace('.safetensors', '')}"
-
+        adapter_name = f"ui_lora_{lora_filename.replace('/', '_').replace('.', '_')}"
         try:
-            # Only load if not already loaded
             if adapter_name not in loaded_loras:
                 pipe.load_lora_weights(repo_name, weight_name=lora_filename, adapter_name=adapter_name)
                 loaded_loras[adapter_name] = lora_filename
                 log(f"📥 Loaded LoRA: {lora_filename}")
 
-            # Activate adapter with given strength
             pipe.set_adapters([adapter_name], [strength])
             log(f"✅ Applied LoRA adapter: {adapter_name} (strength={strength})")
             return f"LoRA applied: {lora_filename}"
@@ -1069,9 +1131,9 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
 
     clear_lora_btn.click(clear_lora, outputs=[logs_box])
 
-    # =========================
-    # GENERATION
-    # =========================
+    # -------------------------
+    # Run Generation
+    # -------------------------
     run_btn.click(
         generate_image,
         inputs=[prompt, height, width, steps, seed],
@@ -1079,4 +1141,5 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
     )
 
 
+
 demo.launch()