Update app_quant_latent.py

app_quant_latent.py

@@ -246,26 +246,31 @@ log_system_stats("AFTER PIPELINE BUILD")
 
 
 
+from PIL import Image
+
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0, return_latents=False):
     """
-    Robust dual pipeline image generation:
-    - Tries advanced latent method first
-    - Falls back to standard pipeline if latent fails
-    - Always returns an image
-    - Optionally returns latents
-    - Logs progress and errors
+    Robust dual pipeline:
+    - Advanced latent generation first
+    - Fallback to standard pipeline if latent fails
+    - Always returns final image
+    - Returns gallery (latents or final image) and logs
     """
 
     LOGS = []
     image = None
     latents = None
+    gallery = []
+
+    # Keep a placeholder original image (white) in case everything fails
+    original_image = Image.new("RGB", (width, height), color=(255, 255, 255))
 
     try:
         generator = torch.Generator(device).manual_seed(int(seed))
 
         # -------------------------------
-        # Attempt advanced latent generation
+        # Try advanced latent generation
         # -------------------------------
         try:
             batch_size = 1
@@ -294,6 +299,8 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0, retur
                 latents=latents
             )
             image = output.images[0]
+            gallery = [image] if image else []
+
             LOGS.append("✅ Advanced latent generation succeeded.")
 
         # -------------------------------
@@ -302,6 +309,7 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0, retur
         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,
@@ -312,23 +320,21 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0, retur
                     generator=generator
                 )
                 image = output.images[0]
+                gallery = [image] if image else []
                 LOGS.append("✅ Standard pipeline generation succeeded.")
             except Exception as e_standard:
                 LOGS.append(f"❌ Standard pipeline generation failed: {e_standard}")
-                # If everything fails, image remains None
+                image = original_image  # Always return some image
+                gallery = [image]
 
         # -------------------------------
-        # Return results
+        # Return all 3 outputs
         # -------------------------------
-        if return_latents and latents is not None:
-            return image, latents, LOGS
-        else:
-            return image, image, LOGS
+        return image, gallery, LOGS
 
     except Exception as e:
         LOGS.append(f"❌ Inference failed entirely: {e}")
-        return None, None,LOGS
-
+        return original_image, [original_image], LOGS
 
 # ============================================================
 # UI