Update app.py

app.py

@@ -2,41 +2,88 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
-
-import spaces
-from dataclasses import dataclass
-import json
-import logging
-import os
-import random
-import re
+import diffusers
+import io
 import sys
-import warnings
+import logging
+
+# ------------------------
+# GLOBAL LOG BUFFER
+# ------------------------
+log_buffer = io.StringIO()
+
+def log(msg):
+    print(msg)
+    log_buffer.write(msg + "\n")
+
+# Enable diffusers debug logs
+diffusers.utils.logging.set_verbosity_info()
+
+log("Loading Z-Image-Turbo pipeline...")
 
-print("Loading Z-Image-Turbo pipeline...")
 pipe = DiffusionPipeline.from_pretrained(
-    "Tongyi-MAI/Z-Image-Turbo",#"T5B/Z-Image-Turbo-FP8",
+    "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
     attn_implementation="kernels-community/vllm-flash-attn3",
 )
 
-#pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-#spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
-
 pipe.to("cuda")
 
+
+#pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"] #spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+
+# ------------------------
+# ATTENTION + PIPE INFO
+# ------------------------
+def pipeline_debug_info():
+    info = []
+    info.append("=== PIPELINE DEBUG INFO ===")
+    info.append(f"UNet attention backend: {pipe.unet.config.attn_implementation}")
+    info.append(f"Transformer attention backend: {pipe.transformer.config.attn_implementation}")
+    
+    # Processor classes
+    try:
+        info.append(f"UNet mid-block processor: {type(pipe.unet.mid_block.attentions[0].processor)}")
+    except:
+        info.append("UNet mid-block processor: <not found>")
+
+    try:
+        info.append(f"Transformer block processor: {type(pipe.transformer.blocks[0].attn.processor)}")
+    except:
+        info.append("Transformer block processor: <not found>")
+
+    return "\n".join(info)
+
+# ------------------------
+# IMAGE GENERATOR
+# ------------------------
 @spaces.GPU
 def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, num_images):
-    """Generate multiple images from the given prompt."""
-    print(prompt)
+    log_buffer.truncate(0)
+    log_buffer.seek(0)
+
+    log("=== NEW GENERATION REQUEST ===")
+    log(f"Prompt: {prompt}")
+    log(f"Height: {height}, Width: {width}")
+    log(f"Inference Steps: {num_inference_steps}")
+    log(f"Num Images: {num_images}")
+
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
+        log(f"Randomized Seed → {seed}")
+    else:
+        log(f"Seed: {seed}")
 
-    # Clamp num_images to max 3
+    # Clamp images
     num_images = min(max(1, int(num_images)), 3)
-    
+
+    # Debug pipe info
+    log(pipeline_debug_info())
+
     generator = torch.Generator("cuda").manual_seed(int(seed))
+
+    log("Running pipeline forward()...")
     result = pipe(
         prompt=prompt,
         height=int(height),
@@ -45,117 +92,73 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         guidance_scale=0.0,
         generator=generator,
         max_sequence_length=1024,
-        num_images_per_prompt=num_images
+        num_images_per_prompt=num_images,
+        output_type="pil",
     )
-    
-    return result.images, seed
 
-# Example prompts
+    # Tensor diagnostics (shapes only)
+    try:
+        latent_shape = pipe.unet.config.sample_size
+        log(f"UNet latent resolution (approx): {latent_shape}")
+    except:
+        pass
+
+    log("Pipeline finished.")
+    log("Returning images...")
+
+    return result.images, seed, log_buffer.getvalue()
+
+# ------------------------
+# GRADIO UI
+# ------------------------
 examples = [
-    ["Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp, bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda, blurred colorful distant lights."],
-    ["A majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, detailed fantasy art style"],
-    ["Cozy coffee shop interior, warm lighting, rain on windows, plants on shelves, vintage aesthetic, photorealistic"],
-    ["Astronaut riding a horse on Mars, cinematic lighting, sci-fi concept art, highly detailed"],
-    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
+    ["Young Chinese woman in red Hanfu, intricate embroidery..."],
+    ["A majestic dragon soaring through clouds at sunset..."],
+    ["Cozy coffee shop interior, warm lighting, rain on windows..."],
+    ["Astronaut riding a horse on Mars, cinematic lighting..."],
+    ["Portrait of a wise old wizard..."],
 ]
 
-# Build the Gradio interface
-with gr.Blocks(title="Z-Image-Turbo Demo") as demo:
-    gr.Markdown(
-        """
-        # 🎨 Z-Image-Turbo  Multi Image Demo
-        
-        Generate high-quality images using the [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) model.
-        This turbo model generates images in just 8 inference steps!
-        """
-    )
-    
+with gr.Blocks(title="Z-Image-Turbo Debug Demo") as demo:
+    gr.Markdown("# 🎨 Z-Image-Turbo — Multi Image + Full Debug Logs")
+
     with gr.Row():
         with gr.Column(scale=1):
-            prompt = gr.Textbox(
-                label="Prompt",
-                placeholder="Enter your image description...",
-                lines=4,
-            )
-            
-            with gr.Row():
-                height = gr.Slider(
-                    minimum=512,
-                    maximum=2048,
-                    value=1024,
-                    step=64,
-                    label="Height",
-                )
-                width = gr.Slider(
-                    minimum=512,
-                    maximum=2048,
-                    value=1024,
-                    step=64,
-                    label="Width",
-                )
-            with gr.Row():
-                 num_images = gr.Slider(
-                 minimum=1,
-                maximum=3,
-                value=2,
-                  step=1,
-                 label="Number of Images",
-                          )
+            prompt = gr.Textbox(label="Prompt", lines=4)
 
             with gr.Row():
-                num_inference_steps = gr.Slider(
-                    minimum=1,
-                    maximum=20,
-                    value=9,
-                    step=1,
-                    label="Inference Steps",
-                    info="9 steps results in 8 DiT forwards",
-                )
-            
+                height = gr.Slider(512, 2048, 1024, step=64, label="Height")
+                width = gr.Slider(512, 2048, 1024, step=64, label="Width")
+
+            num_images = gr.Slider(1, 3, 2, step=1, label="Number of Images")
+
+            num_inference_steps = gr.Slider(
+                1, 20, 9, step=1, label="Inference Steps",
+                info="9 steps = 8 DiT forward passes",
+            )
+
             with gr.Row():
-                seed = gr.Number(
-                    label="Seed",
-                    value=42,
-                    precision=0,
-                )
-                randomize_seed = gr.Checkbox(
-                    label="Randomize Seed",
-                    value=False,
-                )
-            
-            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
-        
+                seed = gr.Number(label="Seed", value=42, precision=0)
+                randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
+
+            generate_btn = gr.Button("🚀 Generate", variant="primary")
+
         with gr.Column(scale=1):
-            output_images = gr.Gallery(
-                label="Generated Image",
-                type="pil",
+            output_images = gr.Gallery(label="Generated Images")
+            used_seed = gr.Number(label="Seed Used", interactive=False)
+            debug_log = gr.Textbox(
+                label="Debug Log Output",
+                lines=25,
+                interactive=False
             )
-            used_seed = gr.Number(
-                label="Seed Used",
-                interactive=False,
-            )
-    
-    gr.Markdown("### 💡 Example Prompts")
-    gr.Examples(
-        examples=examples,
-        inputs=[prompt],
-        cache_examples=False,
-    )
 
-     
-    # Connect the generate button
+    gr.Examples(examples=examples, inputs=[prompt], cache_examples=False)
+
     generate_btn.click(
         fn=generate_image,
         inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, num_images],
-         outputs=[output_images, used_seed],
-         )
-    prompt.submit(
-         fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, num_images],
-          outputs=[output_images, used_seed],
-       )
-
-
+        outputs=[output_images, used_seed, debug_log],
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()