Update app1.py

app1.py

@@ -2,7 +2,7 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
-from transformers import AutoConfig, AutoModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import diffusers
 import io
 
@@ -20,43 +20,33 @@ diffusers.utils.logging.set_verbosity_info()
 
 log("Loading Z-Image-Turbo pipeline...")
 
-# Load main 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
+# ------------------------
 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",
+    text_encoder=fp8_encoder,        # swap in FP8 encoder
+    tokenizer=fp8_tokenizer,         # swap tokenizer
 )
 pipe.to("cuda")
 
-# ------------------------
-# Load FP8 Text Encoder
-# ------------------------
-log("Loading FP8 text encoder...")
-from safetensors.torch import load_file
-import torch
-from transformers import AutoConfig
-from qwen.modeling_qwen import QwenForCausalLM  # ensure Qwen3 model class is imported
-
-# Load config
-config = AutoConfig.from_pretrained("jiangchengchengNLP/qwen3-4b-fp8-scaled")
-
-# Initialize model
-fp8_encoder = QwenForCausalLM(config)
-
-# Load weights manually
-weights = load_file("qwen3_4b_fp8_scaled.safetensors")
-fp8_encoder.load_state_dict(weights, strict=False)
-
-fp8_encoder.to("cuda")
-
-
 # ------------------------
 # PIPELINE DEBUG INFO
 # ------------------------
 def pipeline_debug_info(pipe):
-    info = []
-    info.append("=== PIPELINE DEBUG INFO ===")
+    info = ["=== PIPELINE DEBUG INFO ==="]
     try:
         tr = pipe.transformer.config
         info.append(f"Transformer Class: {pipe.transformer.__class__.__name__}")
@@ -108,10 +98,9 @@ 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
+    # Pipeline debug info
     log(pipeline_debug_info(pipe))
 
     generator = torch.Generator("cuda").manual_seed(int(seed))
@@ -153,7 +142,7 @@ examples = [
     ["Portrait of a wise old wizard..."],
 ]
 
-with gr.Blocks(title="Z-Image-Turbo Multi Image Demo") as demo:
+with gr.Blocks(title="Z-Image-Turbo Multi Image Demo (FP8 Text Encoder)") as demo:
     gr.Markdown("# 🎨 Z-Image-Turbo — Multi Image (FP8 Text Encoder)")
 
     with gr.Row():