Tesring latent

app_quant_latent.py

@@ -560,71 +560,91 @@ def generate_image(prompt, height, width, steps, seed, guidance_scale=0.0):
     device = "cuda"
     generator = torch.Generator(device).manual_seed(int(seed))
 
-    # placeholders
     placeholder = Image.new("RGB", (width, height), color=(255, 255, 255))
     latent_gallery = []
     final_gallery = []
 
     try:
-        # --- Try advanced latent mode ---
+        # ==========================================================
+        # ADVANCED LATENT MODE
+        # ==========================================================
         try:
             latents = safe_get_latents(pipe, height, width, generator, device, LOGS)
 
             for i, t in enumerate(pipe.scheduler.timesteps):
-                # Step-wise denoising
                 with torch.no_grad():
-                    noise_pred = pipe.unet(latents, t, encoder_hidden_states=pipe.get_text_embeddings(prompt))["sample"]
+                    noise_pred = pipe.unet(
+                        latents, t,
+                        encoder_hidden_states=pipe.get_text_embeddings(prompt)
+                    )["sample"]
+
                     latents = pipe.scheduler.step(noise_pred, t, latents)["prev_sample"]
 
-                # Convert latent to preview image
+                # Convert latents to image
                 try:
-                    latent_img = latent_to_image(latents)  # returns single PIL image
+                    latent_img = latent_to_image(latents)
                 except Exception:
                     latent_img = placeholder
 
                 latent_gallery.append(latent_img)
 
-                # Yield intermediate update: final gallery empty for now
-                yield None, latent_gallery, final_gallery, LOGS
+                # 🔥 STREAM UPDATE: (final_image=None for now)
+                yield {
+                    "final_image": None,
+                    "latent_gallery": latent_gallery,
+                    "logs_box": "\n".join(LOGS),
+                }
 
-            # decode final image after all timesteps
+            # ----------------------
+            # Final decode
+            # ----------------------
             final_img = pipe.decode_latents(latents)[0]
             final_gallery.append(final_img)
             LOGS.append("✅ Advanced latent pipeline succeeded.")
-            yield final_img, latent_gallery, final_gallery, LOGS
 
+            # 🔥 Final output
+            yield {
+                "final_image": final_img,
+                "latent_gallery": latent_gallery,
+                "logs_box": "\n".join(LOGS),
+            }
+
+        # ==========================================================
+        # STANDARD PIPELINE (fallback)
+        # ==========================================================
         except Exception as e:
             LOGS.append(f"⚠️ Advanced latent mode failed: {e}")
             LOGS.append("🔁 Switching to standard pipeline...")
 
-            # Standard pipeline fallback
-            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)  # optionally show in latent gallery as last step
-                LOGS.append("✅ Standard pipeline succeeded.")
-                yield final_img, latent_gallery, final_gallery, LOGS
-
-            except Exception as e2:
-                LOGS.append(f"❌ Standard pipeline failed: {e2}")
-                final_gallery.append(placeholder)
-                latent_gallery.append(placeholder)
-                yield placeholder, latent_gallery, final_gallery, LOGS
+            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_image": final_img,
+                "latent_gallery": latent_gallery,
+                "logs_box": "\n".join(LOGS),
+            }
 
     except Exception as e:
         LOGS.append(f"❌ Total failure: {e}")
-        final_gallery.append(placeholder)
-        latent_gallery.append(placeholder)
-        yield placeholder, latent_gallery, final_gallery, LOGS
+        placeholder_img = placeholder
 
+        yield {
+            "final_image": placeholder_img,
+            "latent_gallery": [placeholder_img],
+            "logs_box": "\n".join(LOGS),
+        }
 
 
 
@@ -941,9 +961,13 @@ with gr.Blocks(title="Z-Image-Turbo") as demo:
 
     # Wire the button AFTER all components exist
     run_btn.click(
-        generate_image,
-        inputs=[prompt, height, width, steps, seed],
-        outputs=[final_image, latent_gallery, logs_box]
+    fn=generate_image,
+    inputs=[prompt, height, width, steps, seed],
+    outputs={
+        "final_image": final_image,
+        "latent_gallery": latent_gallery,
+        "logs_box": logs_box,
+         }
     )
 
 demo.launch()