Update app_lora.py

app_lora.py

@@ -756,36 +756,78 @@ def generate_imagenegative(prompt, height, width, steps, seed, guidance_scale=7.
 
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
-   
-    print(prompt)
-  
+    LOGS = []
+    device = "cuda"
+    generator = torch.Generator(device).manual_seed(int(seed))
 
-    
+    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
+    latent_gallery = []
+    final_gallery = []
 
-   
+    # --- Generate latent previews in a loop ---
+    try:
+        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
-    # Debug pipe info
-    
+        # Convert latents to float32 if necessary
+        if latents.dtype != torch.float32:
+            latents = latents.float()
 
-    generator = torch.Generator().manual_seed(int(seed))
+        # Loop for multiple previews before final image
+        num_previews = min(10, steps)  # show ~10 previews
+        preview_steps = torch.linspace(0, 1, num_previews)
 
-    log("Running pipeline forward()...")
-    result = pipe(
-        prompt=prompt,
-        height=int(height),
-        width=int(width),
-        num_inference_steps=int(num_inference_steps),
-        guidance_scale=0.0,
-        generator=generator,
-        max_sequence_length=1024,
-        num_images_per_prompt=2,
-        output_type="pil",
-    )
+        for i, alpha in enumerate(preview_steps):
+            try:
+                with torch.no_grad():
+                    # Simple noise interpolation for preview (simulate denoising progress)
+                    preview_latent = latents * alpha + torch.randn_like(latents) * (1 - alpha)
+                    # Decode to PIL
+                    latent_img_tensor = pipe.vae.decode(preview_latent).sample  # [1,3,H,W]
+                    latent_img_tensor = (latent_img_tensor / 2 + 0.5).clamp(0, 1)
+                    latent_img_tensor = latent_img_tensor.cpu().permute(0, 2, 3, 1)[0]
+                    latent_img = Image.fromarray((latent_img_tensor.numpy() * 255).astype('uint8'))
+            except Exception as e:
+                LOGS.append(f"⚠️ Latent preview decode failed: {e}")
+                latent_img = placeholder
 
-    # Correct latent diagnostics (Z-Image uses VAE + Transformer)
-    
+            latent_gallery.append(latent_img)
+            yield None, latent_gallery, LOGS  # update Gradio with intermediate preview
+
+        # Save final latents to HF
+        latent_dict = {"latents": latents.cpu(), "prompt": prompt, "seed": seed}
+        try:
+            hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_{seed}.pt")
+            LOGS.append(f"🔹 Latents uploaded: {hf_url}")
+        except Exception as e:
+            LOGS.append(f"⚠️ Failed to upload latents: {e}")
+
+    except Exception as e:
+        LOGS.append(f"⚠️ Latent generation failed: {e}")
+        latent_gallery.append(placeholder)
+        yield None, latent_gallery, LOGS
+
+    # --- Final image: untouched standard pipeline ---
+    try:
+        output = pipe(
+            prompt=prompt,
+            height=height,
+            width=width,
+            num_inference_steps=steps,
+            guidance_scale=guidance_scale,
+            generator=generator,
+        )
+        final_img = output.images[0]
+        final_gallery.append(final_img)
+        latent_gallery.append(final_img)  # fallback preview if needed
+        LOGS.append("✅ Standard pipeline succeeded.")
+        yield final_img, latent_gallery, LOGS
+
+    except Exception as e2:
+        LOGS.append(f"❌ Standard pipeline failed: {e2}")
+        final_gallery.append(placeholder)
+        latent_gallery.append(placeholder)
+        yield placeholder, latent_gallery, LOGS
 
-    return result.images, seed, log_buffer.getvalue()
 
 
 # this is astable vesopn tha can gen final and a noise to latent
@@ -793,7 +835,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
 def generate_image_verygood_realnoise(prompt, height, width, steps, seed, guidance_scale=0.0):
     LOGS = []
     device = "cuda"
-    generator = torch.Generator(device).manual_seed(int(seed))
+    generator = torch.Generator().manual_seed(int(seed))
 
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
     latent_gallery = []