Update app_quant_latent1.py

app_quant_latent1.py

@@ -250,53 +250,6 @@ log_system_stats("AFTER PIPELINE BUILD")
 from PIL import Image
 import torch
 
-def safe_generate_with_latents(
- transformer,
- vae,
- text_encoder,
- tokenizer,
- scheduler,
- pipe,
- prompt,
- height,
- width,
- steps,
- guidance_scale,
- negative_prompt,
- num_images_per_prompt,
- generator,
- cfg_normalization,
- cfg_truncation,
- max_sequence_length,
- ):
-
-  try:
-
-     latents_or_images = generate(
-     transformer=transformer,
-     vae=vae,
-     text_encoder=text_encoder,
-     tokenizer=tokenizer,
-     scheduler=scheduler,
-     prompt=prompt,
-     height=height,
-     width=width,
-     num_inference_steps=steps,
-     guidance_scale=guidance_scale,
-     negative_prompt=negative_prompt,
-     num_images_per_prompt=num_images_per_prompt,
-     generator=generator,
-     cfg_normalization=cfg_normalization,
-     cfg_truncation=cfg_truncation,
-     max_sequence_length=max_sequence_length,
-     output_type="latent",  # IMPORTANT
-    )
-     return latents_or_images, None
-
-  except Exception as e:
-    return None, e
-
-
 
 
 
@@ -364,22 +317,84 @@ def safe_get_latents(pipe, height, width, generator, device, LOGS):
 # --------------------------
 @spaces.GPU
 def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
-    LOGS = []
-    latents = None
-    image = None
-    gallery = []
-
-    # placeholder image if all fails
-    placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
-    print(prompt)
-
-    latents, latent_err = safe_generate_with_latents( transformer=transformer, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, scheduler=scheduler,
-                                                     pipe=pipe, prompt=prompt, height=height, width=width, steps=steps, guidance_scale=guidance_scale, negative_prompt="", num_images_per_prompt=1, generator=generator, cfg_normalization=False, cfg_truncation=1.0, max_sequence_length=4096, ) 
-    if latent_err is None: log("✅ Latent generator succeeded.") 
-        try: # Decode latents to image shift_factor = getattr(vae.config, "shift_factor", 0.0) or 0.0 dec = (latents.to(vae.dtype) / vae.config.scaling_factor) +
-            shift_factor image = vae.decode(dec, return_dict=False)[0] image = (image / 2 + 0.5).clamp(0, 1) image = image.cpu().permute(0, 2, 3, 1).numpy() image = (image * 255).round().astype("uint8") from PIL import Image image = Image.fromarray(image[0]) log("🟢 Final image decoded from latent generator.") return image, latents, LOGS except Exception as decode_error: log(f"⚠️ Latent decode failed: {decode_error}") log("🔁 Falling back to standard pipeline...") else: log(f"⚠️ Latent generator failed: {latent_err}") log("🔁 Switching to standard pipeline...") # ========================================================== # 🟩 STANDARD PIPELINE FALLBACK (Never fails) # ========================================================== try: output = pipe( prompt=prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=guidance_scale, generator=generator, ) image = output.images[0] log("🟢 Standard pipeline succeeded.") return image, None, LOGS except Exception as e: log(f"❌ Standard pipeline failed: {e}") return None, None, LOGS
-
+LOGS = []
+def log(msg):
+ LOGS.append(msg)
+ print(msg)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+generator = torch.Generator(device).manual_seed(int(seed))
+
+log("🎨 START IMAGE GENERATION")
+
+
+# ==========================================================
+# 🧪 TRY ADVANCED LATENT GENERATOR (Your original generate())
+# ==========================================================
+latents, latent_err = safe_generate_with_latents(
+    transformer=transformer,
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    scheduler=scheduler,
+    pipe=pipe,
+    prompt=prompt,
+    height=height,
+    width=width,
+    steps=steps,
+    guidance_scale=guidance_scale,
+    negative_prompt="",
+    num_images_per_prompt=1,
+    generator=generator,
+    cfg_normalization=False,
+    cfg_truncation=1.0,
+    max_sequence_length=4096,
+)
+
+if latent_err is None:
+    log("✅ Latent generator succeeded.")
+    try:
+        # Decode latents to image
+        shift_factor = getattr(vae.config, "shift_factor", 0.0) or 0.0
+        dec = (latents.to(vae.dtype) / vae.config.scaling_factor) + shift_factor
+        image = vae.decode(dec, return_dict=False)[0]
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).numpy()
+        image = (image * 255).round().astype("uint8")
+        from PIL import Image
+        image = Image.fromarray(image[0])
+
+        log("🟢 Final image decoded from latent generator.")
+        return image, latents, LOGS
+
+    except Exception as decode_error:
+        log(f"⚠️ Latent decode failed: {decode_error}")
+        log("🔁 Falling back to standard pipeline...")
+
+else:
+    log(f"⚠️ Latent generator failed: {latent_err}")
+    log("🔁 Switching to standard pipeline...")
+
+# ==========================================================
+# 🟩 STANDARD PIPELINE FALLBACK (Never fails)
+# ==========================================================
+try:
+    output = pipe(
+        prompt=prompt,
+        height=height,
+        width=width,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        generator=generator,
+    )
+    image = output.images[0]
+    log("🟢 Standard pipeline succeeded.")
+    return image, None, LOGS
 
+except Exception as e:
+    log(f"❌ Standard pipeline failed: {e}")
+    return None, None, LOGS
 
 # --------------------------
 # Helper: Safe latent extractor