Update app1.py

app1.py

@@ -2,7 +2,7 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 import diffusers
 import io
 
@@ -18,30 +18,21 @@ def log(msg):
 # Enable diffusers debug logs
 diffusers.utils.logging.set_verbosity_info()
 
-log("Loading Z-Image-Turbo pipeline...")
-
-# ------------------------
-# Load FP8 text encoder
-# ------------------------
-log("Loading FP8 text encoder: Qwen/Qwen3-VL-2B-Instruct-FP8...")
-fp8_model_name = "Qwen/Qwen3-VL-2B-Instruct-FP8"
-fp8_tokenizer = AutoTokenizer.from_pretrained(fp8_model_name)
-fp8_encoder = AutoModelForCausalLM.from_pretrained(fp8_model_name, torch_dtype=torch.bfloat16)
-fp8_encoder.to("cuda")
-
 # ------------------------
-# Load main diffusion pipeline
+# LOAD PIPELINES
 # ------------------------
+log("Loading Z-Image-Turbo pipeline...")
 pipe = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
     attn_implementation="kernels-community/vllm-flash-attn3",
-    text_encoder=fp8_encoder,        # swap in FP8 encoder
-    tokenizer=fp8_tokenizer,         # swap tokenizer
 )
 pipe.to("cuda")
 
+log("Loading FP8 text encoder: Qwen/Qwen3-4B...")
+fp8_encoder = pipeline("text-generation", model="Qwen/Qwen3-4B", device=0)  # device=0 → CUDA
+
 # ------------------------
 # PIPELINE DEBUG INFO
 # ------------------------
@@ -98,14 +89,20 @@ 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)
 
-    # Pipeline debug info
+    # Run FP8 text encoder first
+    log("Encoding prompt with FP8 text encoder...")
+    encoded_prompt = fp8_encoder([{"role": "user", "content": prompt}])
+    log(f"FP8 encoding output: {encoded_prompt}")
+
+    # Debug pipeline info
     log(pipeline_debug_info(pipe))
 
     generator = torch.Generator("cuda").manual_seed(int(seed))
 
-    log("Running pipeline forward()...")
+    log("Running Z-Image-Turbo pipeline forward()...")
     result = pipe(
         prompt=prompt,
         height=int(height),
@@ -142,7 +139,7 @@ examples = [
     ["Portrait of a wise old wizard..."],
 ]
 
-with gr.Blocks(title="Z-Image-Turbo Multi Image Demo (FP8 Text Encoder)") as demo:
+with gr.Blocks(title="Z-Image-Turbo Multi Image Demo") as demo:
     gr.Markdown("# 🎨 Z-Image-Turbo — Multi Image (FP8 Text Encoder)")
 
     with gr.Row():