Update app_lora.py

app_lora.py

@@ -688,69 +688,70 @@ def generate_image_all_latents(prompt, height, width, steps, seed, guidance_scal
 
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=7.5):
+    """
+    Generate image using ZImagePipeline with optional LoRA adapter.
+    Shows step previews and final image.
+    """
     LOGS = []
-    device = "cuda"
-    generator = torch.Generator(device).manual_seed(int(seed))
-    
+    generator = torch.Generator("cuda").manual_seed(int(seed))
+
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
     latent_gallery = []
+    final_gallery = []
+
+    # Determine active LoRA adapter
+    active_adapter = None
+    active_strength = 1.0
+    if loaded_loras:
+        active_adapter = list(loaded_loras.keys())[-1]
+        active_strength = loaded_loras[active_adapter + "_strength"] if loaded_loras.get(active_adapter + "_strength") else 1.0
+        pipe.set_adapters([active_adapter], [active_strength])
+        LOGS.append(f"🧩 Using LoRA adapter: {active_adapter} (strength={active_strength})")
+    else:
+        pipe.set_adapters([], [])
+        LOGS.append("⚡ No LoRA applied")
 
     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),
+        # Generate small preview steps
+        num_preview_steps = min(5, steps)
+        for i in range(num_preview_steps):
+            step = i + 1
+            try:
+                preview_output = pipe(
+                    prompt=prompt,
+                    height=height // 4,  # small preview
+                    width=width // 4,
+                    num_inference_steps=step,
+                    guidance_scale=guidance_scale,
+                    generator=generator,
+                )
+                img = preview_output.images[0].resize((width, height))
+                latent_gallery.append(img)
+            except Exception as e:
+                LOGS.append(f"⚠️ Preview step {step} failed: {e}")
+                latent_gallery.append(placeholder)
+
+        # --- Final image ---
+        output = pipe(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
             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)
+        final_img = output.images[0]
+        final_gallery.append(final_img)
         latent_gallery.append(final_img)
-        LOGS.append("✅ Pipeline generation completed successfully.")
-        yield final_img, latent_gallery[-5:] + [final_img], LOGS
+        LOGS.append("✅ Image generation completed.")
+
+        yield final_img, latent_gallery, LOGS
 
     except Exception as e:
         LOGS.append(f"❌ Generation failed: {e}")
         latent_gallery.append(placeholder)
-        yield placeholder, latent_gallery[-5:] + [placeholder], LOGS
+        final_gallery.append(placeholder)
+        yield placeholder, latent_gallery, LOGS
 
 
 def generate_image1(prompt, height, width, steps, seed, guidance_scale=0.0):