Update app_quant_latent.py

app_quant_latent.py

@@ -245,7 +245,88 @@ log_system_stats("AFTER PIPELINE BUILD")
 
 
 
+@spaces.GPU
+def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
+    """
+    Returns 3 outputs for Gradio:
+    - image: PIL.Image
+    - gallery: optional latents or empty list
+    - logs: list of messages
+    """
+
+    LOGS = []
+    image = None
+    latents = None
+    gallery = []
+
+    try:
+        generator = torch.Generator(device).manual_seed(int(seed))
+
+        # -------------------------------
+        # Attempt advanced latent generation
+        # -------------------------------
+        try:
+            batch_size = 1
+            num_channels_latents = getattr(pipe.unet, "in_channels", None)
+            if num_channels_latents is None:
+                raise AttributeError("pipe.unet.in_channels not found, fallback to standard pipeline")
+
+            latents = pipe.prepare_latents(
+                batch_size=batch_size,
+                num_channels=num_channels_latents,
+                height=height,
+                width=width,
+                dtype=torch.float32,
+                device=device,
+                generator=generator
+            )
+            LOGS.append(f"✅ Latents prepared: {latents.shape}")
+
+            output = pipe(
+                prompt=prompt,
+                height=height,
+                width=width,
+                num_inference_steps=steps,
+                guidance_scale=guidance_scale,
+                generator=generator,
+                latents=latents
+            )
+            image = output.images[0]
+            gallery = [image]  # You can also return other latent previews if needed
+            LOGS.append("✅ Advanced latent generation succeeded.")
+
+        # -------------------------------
+        # Fallback to standard pipeline
+        # -------------------------------
+        except Exception as e_latent:
+            LOGS.append(f"⚠️ Advanced latent generation failed: {e_latent}")
+            LOGS.append("🔁 Falling back to standard pipeline...")
+            try:
+                output = pipe(
+                    prompt=prompt,
+                    height=height,
+                    width=width,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance_scale,
+                    generator=generator
+                )
+                image = output.images[0]
+                gallery = [image]  # still return at least the final image
+                LOGS.append("✅ Standard pipeline generation succeeded.")
+            except Exception as e_standard:
+                LOGS.append(f"❌ Standard pipeline generation failed: {e_standard}")
+                # Fallback: return None image and empty gallery
+                image = None
+                gallery = []
+
+        # -------------------------------
+        # Return all 3 outputs
+        # -------------------------------
+        return image, gallery, LOGS
 
+    except Exception as e:
+        LOGS.append(f"❌ Inference failed entirely: {e}")
+        return None, [], LOGS
 
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):