Update app_quant_latent.py

app_quant_latent.py

@@ -246,66 +246,89 @@ log_system_stats("AFTER PIPELINE BUILD")
 
 
 
+
 @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
+    """
+
+    LOGS = []
+    image = None
+    latents = None
 
- try:
-    generator = torch.Generator(device).manual_seed(int(seed))
-
-    # Try advanced latent preparation
     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
-        )
-        log(f"✅ Latents prepared: {latents.shape}")
-
-        # Generate image using prepared latents
-        output = pipe(
-            prompt=prompt,
-            height=height,
-            width=width,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale,
-            generator=generator,
-            latents=latents
-        )
-
-    except Exception as e_inner:
-        # If advanced method fails, fallback to standard pipeline
-        log(f"⚠️ Advanced latent method failed: {e_inner}")
-        log("🔁 Falling back to standard pipeline...")
-        output = pipe(
-            prompt=prompt,
-            height=height,
-            width=width,
-            num_inference_steps=steps,
-            guidance_scale=guidance_scale,
-            generator=generator
-        )
-
-    image = output.images[0]
-    log("✅ Inference finished successfully.")
-
-    if return_latents and 'latents' in locals():
-        return image, latents, LOGS
-    else:
-        return image, LOGS
-
- except Exception as e:
-    log(f"❌ Inference failed entirely: {e}")
-    return None, LOGS
+        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]
+            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]
+                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
+
+        # -------------------------------
+        # Return results
+        # -------------------------------
+        if return_latents and latents is not None:
+            return image, latents, LOGS
+        else:
+            return image, LOGS
 
+    except Exception as e:
+        LOGS.append(f"❌ Inference failed entirely: {e}")
+        return None, LOGS
 
 
 # ============================================================