Update app_lora.py

app_lora.py

@@ -754,90 +754,52 @@ def generate_imagenegative(prompt, height, width, steps, seed, guidance_scale=7.
         yield placeholder, latent_gallery, LOGS
 
 
+@spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
-    LOGS = []
-    device = "cuda"
-    cpu_device = "cpu"
-    generator = torch.Generator(device).manual_seed(int(seed))
+   
+    print(prompt)
+  
 
-    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
-    latent_gallery = []
-    final_gallery = []
-
-    last_latents = []  # store last 5 preview latents on CPU
+    if randomize_seed:
+        seed = torch.randint(0, 2**32 - 1, (1,)).item()
+        log(f"Randomized Seed → {seed}")
+    else:
+        log(f"Seed: {seed}")
+
+   
+
+    # Debug pipe info
+    log(pipeline_debug_info(pipe))
+
+    generator = torch.Generator().manual_seed(seed)
+
+    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",
+    )
 
+    # Correct latent diagnostics (Z-Image uses VAE + Transformer)
     try:
-        # --- Initial latents ---
-        latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
-        latents = latents.float().to(cpu_device)  # move to CPU
-
-        num_previews = min(10, steps)
-        preview_indices = torch.linspace(0, steps - 1, num_previews).long()
-
-        for i, step_idx in enumerate(preview_indices):
-            try:
-                with torch.no_grad():
-                    # --- Z-Image Turbo-style denoise simulation ---
-                    t = 1.0 - (i / num_previews)  # linear decay [1.0 -> 0.0]
-                    noise_scale = t ** 0.5  # reduce noise over steps (sqrt for smoother)
-                    denoise_latent = latents * t + torch.randn_like(latents) * noise_scale
-
-                    # Move to VAE device & dtype
-                    denoise_latent = denoise_latent.to(pipe.vae.device).to(pipe.vae.dtype)
-
-                    # Decode latent to image
-                    decoded = pipe.vae.decode(denoise_latent, return_dict=False)[0]
-                    decoded = (decoded / 2 + 0.5).clamp(0, 1)
-                    decoded = decoded.cpu().permute(0, 2, 3, 1).float().numpy()
-                    decoded = (decoded * 255).round().astype("uint8")
-                    latent_img = Image.fromarray(decoded[0])
-
-            except Exception as e:
-                LOGS.append(f"⚠️ Latent preview decode failed: {e}")
-                latent_img = placeholder
-
-            latent_gallery.append(latent_img)
-
-            # Keep last 5 latents only
-            last_latents.append(denoise_latent.cpu().clone())
-            if len(last_latents) > 5:
-                last_latents.pop(0)
-
-            # Show only last 5 previews in UI
-            yield None, latent_gallery[-5:], LOGS
-
-        # Optionally: upload last 5 latents
-        # latent_dict = {"latents": last_latents, "prompt": prompt, "seed": seed}
-        # hf_url = upload_latents_to_hf(latent_dict, filename=f"latents_last5_{seed}.pt")
-        # LOGS.append(f"🔹 Last 5 latents uploaded: {hf_url}")
+        log(f"VAE latent channels: {pipe.vae.config.latent_channels}")
+        log(f"VAE scaling factor: {pipe.vae.config.scaling_factor}")
+        log(f"Transformer latent size: {pipe.transformer.config.sample_size}")
+        log(latent_shape_info(height, width, pipe))
 
     except Exception as e:
-        LOGS.append(f"⚠️ Latent generation failed: {e}")
-        latent_gallery.append(placeholder)
-        yield None, latent_gallery[-5:], LOGS
+        log(f"Latent diagnostics error: {e}")
 
-    # --- Final image on GPU ---
-    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)
-        LOGS.append("✅ Standard pipeline succeeded.")
-        yield final_img, latent_gallery[-5:] + [final_img], LOGS  # last 5 previews + final
-
-    except Exception as e2:
-        LOGS.append(f"❌ Standard pipeline failed: {e2}")
-        final_gallery.append(placeholder)
-        latent_gallery.append(placeholder)
-        yield placeholder, latent_gallery[-5:] + [placeholder], LOGS
+    log("Pipeline finished.")
+    log("Returning images...")
 
+    return result.images, seed, log_buffer.getvalue()
 
 
 # this is astable vesopn tha can gen final and a noise to latent