Update app1.py

app1.py

@@ -2,10 +2,9 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
+from transformers import AutoConfig, AutoModel
 import diffusers
 import io
-import logging
-from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # ------------------------
 # GLOBAL LOG BUFFER
@@ -21,45 +20,39 @@ diffusers.utils.logging.set_verbosity_info()
 
 log("Loading Z-Image-Turbo pipeline...")
 
-# ------------------------
-# Load FP8 text encoder + tokenizer
-# ------------------------
-log("Loading FP8 Qwen3-4B tokenizer + text encoder...")
-fp8_tokenizer = AutoTokenizer.from_pretrained(
-    "jiangchengchengNLP/qwen3-4b-fp8-scaled"
-)
-fp8_text_encoder = AutoModelForCausalLM.from_pretrained(
-    "jiangchengchengNLP/qwen3-4b-fp8-scaled",
-    device_map="auto",
-    torch_dtype=torch.bfloat16,  # can replace with torch.float8_e4m3fn if PyTorch nightly supports
-)
-
-# ------------------------
-# Load Z-Image-Turbo
-# ------------------------
+# Load main pipeline
 pipe = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
     attn_implementation="kernels-community/vllm-flash-attn3",
 )
-
-# Inject FP8 tokenizer + text encoder
-pipe.tokenizer = fp8_tokenizer
-pipe.text_encoder = fp8_text_encoder
 pipe.to("cuda")
 
 # ------------------------
-# Pipeline debug info
+# Load FP8 Text Encoder
+# ------------------------
+log("Loading FP8 text encoder...")
+config = AutoConfig.from_pretrained("jiangchengchengNLP/qwen3-4b-fp8-scaled")
+fp8_encoder = AutoModel.from_pretrained(
+    "jiangchengchengNLP/qwen3-4b-fp8-scaled",
+    config=config,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+pipe.text_encoder = fp8_encoder
+log("FP8 text encoder loaded successfully.")
+log(f"Tokenizer remains: {pipe.tokenizer.__class__.__name__}")
+
+# ------------------------
+# PIPELINE DEBUG INFO
 # ------------------------
 def pipeline_debug_info(pipe):
     info = []
     info.append("=== PIPELINE DEBUG INFO ===")
-
     try:
         tr = pipe.transformer.config
         info.append(f"Transformer Class: {pipe.transformer.__class__.__name__}")
-        # Z-Image-Turbo keys
         info.append(f"Hidden dim: {tr.get('hidden_dim')}")
         info.append(f"Attention heads: {tr.get('num_heads')}")
         info.append(f"Depth (layers): {tr.get('depth')}")
@@ -69,7 +62,6 @@ def pipeline_debug_info(pipe):
     except Exception as e:
         info.append(f"Transformer diagnostics failed: {e}")
 
-    # VAE info
     try:
         vae = pipe.vae.config
         info.append(f"VAE latent channels: {vae.latent_channels}")
@@ -79,7 +71,6 @@ def pipeline_debug_info(pipe):
 
     return "\n".join(info)
 
-
 def latent_shape_info(h, w, pipe):
     try:
         c = pipe.vae.config.latent_channels
@@ -90,9 +81,8 @@ def latent_shape_info(h, w, pipe):
     except Exception as e:
         return f"Latent shape calc failed: {e}"
 
-
 # ------------------------
-# IMAGE GENERATOR
+# IMAGE GENERATION
 # ------------------------
 @spaces.GPU
 def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, num_images):
@@ -111,6 +101,7 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
     else:
         log(f"Seed: {seed}")
 
+    # Clamp images to 1–3
     num_images = min(max(1, int(num_images)), 3)
 
     # Debug pipeline info
@@ -131,7 +122,7 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
         output_type="pil",
     )
 
-    # Tensor diagnostics (shapes only)
+    # Latent diagnostics
     try:
         log(f"VAE latent channels: {pipe.vae.config.latent_channels}")
         log(f"VAE scaling factor: {pipe.vae.config.scaling_factor}")
@@ -144,7 +135,6 @@ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_s
 
     return result.images, seed, log_buffer.getvalue()
 
-
 # ------------------------
 # GRADIO UI
 # ------------------------
@@ -157,7 +147,7 @@ examples = [
 ]
 
 with gr.Blocks(title="Z-Image-Turbo Multi Image Demo") as demo:
-    gr.Markdown("# 🎨 Z-Image-Turbo — Multi Image ")
+    gr.Markdown("# 🎨 Z-Image-Turbo — Multi Image (FP8 Text Encoder)")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -181,7 +171,7 @@ with gr.Blocks(title="Z-Image-Turbo Multi Image Demo") as demo:
             generate_btn = gr.Button("🚀 Generate", variant="primary")
 
         with gr.Column(scale=1):
-            output_images = gr.Gallery(label="Generated Images")
+            output_images = gr.Gallery(label="Generated Images", type="pil")
             used_seed = gr.Number(label="Seed Used", interactive=False)
             debug_log = gr.Textbox(label="Debug Log Output", lines=25, interactive=False)