Update app.py

app.py

@@ -11,7 +11,7 @@ import warnings
 # 忽略警告
 warnings.filterwarnings("ignore")
 
-# ==================== 1. 分辨率配置 (严格保持原样) ====================
+# ==================== 1. 分辨率配置 ====================
 RES_CHOICES = {
     "1024": [
         "720x1280 (9:16)",
@@ -65,46 +65,82 @@ def get_resolution(resolution_str):
         return width - width % 8, height - height % 8
     return 1024, 1024
 
-# ==================== 2. 模型加载与优化 ====================
+# ==================== 2. 模型加载与核心优化 ====================
 print("🚀 Loading Z-Image-Turbo pipeline...")
 
-# 加载模型
+# 必须设置为 True，才能加载 Z-Image 自定义的 Pipeline 和 Transformer 类
+# 否则无法调用 set_attention_backend
 pipe = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=True,
     use_safetensors=True,
+    trust_remote_code=True, 
 )
 
-# 【核心加速】使用 FlowMatchEulerDiscreteScheduler 并设置 shift=3.0
-# 这是 Turbo 模型速度快的关键，能让模型在 4-10 步内生成高质量图片
-scheduler_config = dict(pipe.scheduler.config)
-scheduler_config.pop("algorithm_type", None) 
-pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
-    scheduler_config, 
-    shift=3.0
-)
+# 使用 FlowMatchEulerDiscreteScheduler 并设置 shift=3.0
+try:
+    scheduler_config = dict(pipe.scheduler.config)
+    scheduler_config.pop("algorithm_type", None) 
+    pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
+        scheduler_config, 
+        shift=3.0
+    )
+    print("✅ Scheduler optimized with shift=3.0")
+except Exception as e:
+    print(f"⚠️ Scheduler config warning: {e}")
 
 # 移动到 GPU
 pipe.to("cuda")
 
-# 【显存优化】开启 xFormers (如果环境支持)
-try:
-    pipe.enable_xformers_memory_efficient_attention()
-    print("✅ XFormers enabled")
-except Exception as e:
-    print(f"⚠️ XFormers not available: {e}")
+# 尝试按顺序启用最快的后端
+def enable_best_attention_backend(pipeline):
+    # 检查 pipeline.transformer 是否有 set_attention_backend 方法
+    # 这是 Z-Image 自定义类特有的
+    if hasattr(pipeline.transformer, "set_attention_backend"):
+        try:
+            # 优先尝试 Flash Attention 2 (A100/A10G)
+            print("⚡ Attempting to set backend to 'flash_attention_2'...")
+            pipeline.transformer.set_attention_backend("flash_attention_2")
+            print("✅ Attention backend set to: flash_attention_2")
+            return
+        except Exception as e:
+            print(f"ℹ️ Flash Attention 2 not available: {e}")
+        
+        try:
+            # 其次尝试 xFormers (T4/V100 通用)
+            print("⚡ Attempting to set backend to 'xformers'...")
+            pipeline.transformer.set_attention_backend("xformers")
+            print("✅ Attention backend set to: xformers")
+            return
+        except Exception as e:
+            print(f"ℹ️ xFormers not available: {e}")
+            
+        try:
+            # 最后使用 PyTorch 2.0 Native SDPA
+            print("⚡ Setting backend to 'native' (SDPA)...")
+            pipeline.transformer.set_attention_backend("native")
+            print("✅ Attention backend set to: native")
+        except Exception as e:
+            print(f"⚠️ Could not set custom attention backend: {e}")
+    else:
+        print("⚠️ Warning: Transformer model does not support 'set_attention_backend'. Custom code might not be loaded.")
+        # 如果加载失败，尝试标准的 xformers
+        try:
+            pipeline.enable_xformers_memory_efficient_attention()
+            print("✅ Standard xFormers enabled as fallback")
+        except:
+            pass
+
+# 执行后端设置
+enable_best_attention_backend(pipe)
 
-# 【显存优化】开启 VAE 切片
+# VAE 内存优化
 try:
     pipe.vae.enable_slicing()
 except:
     pass
 
-# 注意：移除了 pipe.transformer = torch.compile(...) 
-# 原因是它导致了 'Cannot construct ConstantVariable for value of type torch.device' 错误。
-# 目前仅靠 Scheduler 优化和 xFormers 已经足够快且稳定。
-
 # ==================== 3. 生成逻辑 ====================
 @spaces.GPU
 def generate_image(
@@ -117,37 +153,33 @@ def generate_image(
     seed, 
     randomize_seed,
     negative_prompt,
-    gallery_history, # 接收历史记录
+    gallery_history,
     progress=gr.Progress(track_tqdm=True)
 ):
     if gallery_history is None:
         gallery_history = []
         
     try:
-        # 1. 输入校验
         if not prompt or len(prompt.strip()) < 2:
             raise gr.Error("请输提示词 (Prompt)")
         
         prompt = prompt.strip()
         neg_prompt = negative_prompt.strip() if negative_prompt else None
         
-        # 2. 分辨率计算
         if use_custom_res:
             width = int(custom_width) - int(custom_width) % 8
             height = int(custom_height) - int(custom_height) % 8
         else:
             width, height = get_resolution(resolution_choice)
         
-        # 3. 种子处理
         if randomize_seed:
             seed = random.randint(0, 2**32 - 1)
         seed = int(seed)
         
-        # 4. 生成
         start_time = time.time()
         generator = torch.Generator("cuda").manual_seed(seed)
         
-        # 清理显存
+        # 清理显存确保最大空间
         torch.cuda.empty_cache()
         
         with torch.cuda.amp.autocast(dtype=torch.bfloat16):
@@ -156,7 +188,7 @@ def generate_image(
                 height=height,
                 width=width,
                 num_inference_steps=int(num_inference_steps),
-                guidance_scale=0.0, # Turbo 模型不需要 guidance_scale
+                guidance_scale=0.0,
                 generator=generator,
                 negative_prompt=neg_prompt,
                 max_sequence_length=512,
@@ -164,190 +196,135 @@ def generate_image(
             
         gen_time = time.time() - start_time
         
-        # 5. 构建历史记录
-        # Gallery 格式: [(image_path_or_obj, label), ...]
-        info_label = f"{width}x{height} | Seed: {seed} | {gen_time:.1f}s"
+        # 格式化历史记录
+        info_label = f"{width}x{height} | Steps: {num_inference_steps} | Seed: {seed} | {gen_time:.2f}s"
         gallery_history.insert(0, (image, info_label))
         
         return gallery_history, seed
         
     except Exception as e:
-        raise gr.Error(f"生成失败: {str(e)}")
+        raise gr.Error(f"生成错误: {str(e)}")
 
-# ==================== 4. UI 样式 (CSS) ====================
+# ==================== 4. UI 样式 ====================
 css = """
-/* 全局字体 */
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;800&display=swap');
 body, .gradio-container { font-family: 'Inter', sans-serif !important; }
 
-/* 标题样式 */
 .header-container { text-align: center; margin-bottom: 20px; }
 .header-title { 
     font-size: 2.5rem; font-weight: 800; margin: 0;
-    background: linear-gradient(135deg, #6366f1, #a855f7, #ec4899);
+    background: linear-gradient(135deg, #f59e0b, #ea580c);
     -webkit-background-clip: text; -webkit-text-fill-color: transparent;
 }
 .header-subtitle { font-size: 1rem; color: #6b7280; font-weight: 500; }
 
-/* 按钮样式 */
 .primary-btn { 
-    background: linear-gradient(90deg, #4f46e5 0%, #7c3aed 100%) !important; 
+    background: linear-gradient(90deg, #f59e0b 0%, #d97706 100%) !important; 
     border: none !important; 
     color: white !important; 
     font-weight: 600 !important;
     font-size: 1.1rem !important;
-    box-shadow: 0 4px 6px -1px rgba(79, 70, 229, 0.2) !important;
-    transition: all 0.2s !important;
+    box-shadow: 0 4px 6px -1px rgba(245, 158, 11, 0.2) !important;
 }
-.primary-btn:hover { transform: translateY(-2px); box-shadow: 0 10px 15px -3px rgba(79, 70, 229, 0.3) !important; }
+.primary-btn:hover { transform: translateY(-2px); box-shadow: 0 10px 15px -3px rgba(245, 158, 11, 0.3) !important; }
 
-/* 输入框和面板 */
 .panel-container { 
-    background: #ffffff; 
-    border: 1px solid #e5e7eb; 
-    border-radius: 16px; 
-    padding: 20px; 
-    box-shadow: 0 1px 3px rgba(0,0,0,0.05);
+    background: #ffffff; border: 1px solid #e5e7eb; border-radius: 12px; padding: 15px; 
 }
-/* 暗黑模式适配 */
 .dark .panel-container { background: #1f2937; border-color: #374151; }
-
-/* 画廊样式 */
-#output-gallery { min-height: 600px; }
 """
 
-# ==================== 5. Gradio 界面构建 ====================
-with gr.Blocks(theme=gr.themes.Soft(), css=css, title="Z-Image-Turbo") as demo:
+# ==================== 5. Gradio 界面 ====================
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange"), css=css, title="Z-Image-Turbo") as demo:
     
-    # 头部
     gr.HTML("""
         <div class="header-container">
             <h1 class="header-title">⚡ Z-Image-Turbo</h1>
-            <p class="header-subtitle">Ultra-Fast Generation • 8 Steps • Gallery History</p>
+            <p class="header-subtitle">Optimized Backend • 8 Steps • Gallery History</p>
         </div>
     """)
     
     with gr.Row():
-        # --- 左侧：控制面板 ---
+        # --- 控制面板 ---
         with gr.Column(scale=4, min_width=320):
             with gr.Group(elem_classes="panel-container"):
                 prompt = gr.Textbox(
-                    label="提示词 (Prompt)",
-                    placeholder="Describe your imagination...",
-                    lines=4,
-                    show_label=True
+                    label="Prompt",
+                    placeholder="Enter your prompt here...",
+                    lines=3
                 )
                 negative_prompt = gr.Textbox(
-                    label="反向提示词 (Negative Prompt)",
-                    placeholder="Blurry, low quality, ugly...",
-                    lines=2
+                    label="Negative Prompt",
+                    placeholder="Low quality, blurry...",
+                    lines=1
                 )
-                
-                generate_btn = gr.Button("🚀 Generate Image", elem_classes="primary-btn")
+                generate_btn = gr.Button("🚀 Generate", elem_classes="primary-btn")
             
             with gr.Group(elem_classes="panel-container"):
-                gr.Markdown("### 📐 尺寸设置")
-                
-                # 分辨率分类
-                with gr.Row():
-                    res_category = gr.Radio(
-                        choices=["1024", "1280", "1536"],
-                        value="1024",
-                        label="基准分辨率",
-                        container=False
-                    )
-                
-                # 具体分辨率下拉框
+                gr.Markdown("### 📐 Resolution")
+                res_category = gr.Radio(
+                    choices=["1024", "1280", "1536"],
+                    value="1024",
+                    label="Resolution Base",
+                    container=False
+                )
                 resolution_dropdown = gr.Dropdown(
                     choices=RES_CHOICES["1024"],
                     value=RES_CHOICES["1024"][0],
-                    label="选择比例",
-                    show_label=False,
-                    interactive=True
+                    label="Select Ratio",
+                    show_label=False
                 )
                 
-                # 自定义分辨率开关
-                with gr.Accordion("自定义尺寸 (高级)", open=False):
-                    use_custom_res = gr.Checkbox(label="启用自定义尺寸", value=False)
+                with gr.Accordion("Custom Size", open=False):
+                    use_custom_res = gr.Checkbox(label="Enable Custom", value=False)
                     with gr.Row(visible=False) as custom_res_row:
-                        width_slider = gr.Slider(512, 1536, value=1024, step=64, label="宽")
-                        height_slider = gr.Slider(512, 1536, value=1024, step=64, label="高")
+                        width_slider = gr.Slider(512, 1536, value=1024, step=64, label="W")
+                        height_slider = gr.Slider(512, 1536, value=1024, step=64, label="H")
 
-            with gr.Accordion("⚙️ 高级设置", open=False):
+            with gr.Accordion("⚙️ Settings", open=False):
                 with gr.Group(elem_classes="panel-container"):
-                    steps_slider = gr.Slider(
-                        minimum=4, maximum=20, value=8, step=1, 
-                        label="步数 (Steps) - 推荐 4-8"
-                    )
+                    steps_slider = gr.Slider(4, 20, value=8, step=1, label="Steps")
                     with gr.Row():
-                        random_seed = gr.Checkbox(label="随机种子", value=True)
-                        seed_input = gr.Number(label="种子", value=42, visible=False, precision=0)
+                        random_seed = gr.Checkbox(label="Random Seed", value=True)
+                        seed_input = gr.Number(label="Seed", value=42, visible=False, precision=0)
 
-        # --- 右侧：画廊展示 ---
+        # --- 画廊 ---
         with gr.Column(scale=6, min_width=500):
             output_gallery = gr.Gallery(
-                label="生成历史 (History)",
+                label="History",
                 value=[],
                 columns=[2],
                 rows=[2],
                 object_fit="contain",
                 height="auto",
-                elem_id="output-gallery",
-                show_label=True,
                 show_share_button=True,
                 show_download_button=True,
                 interactive=False
             )
-            
             with gr.Row():
-                last_seed_display = gr.Textbox(label="当前图种子", interactive=False, scale=3)
-                clear_history_btn = gr.Button("🗑️ 清空历史", scale=1, variant="secondary")
+                last_seed_display = gr.Textbox(label="Last Seed", interactive=False, scale=3)
+                clear_btn = gr.Button("🗑️ Clear", scale=1, variant="secondary")
 
-    # ==================== 6. 交互逻辑 ====================
-    
-    # 1. 切换分辨率分类时，更新下拉框
+    # 交互逻辑
     def update_resolution_list(category):
-        new_choices = RES_CHOICES[category]
-        return gr.Dropdown(choices=new_choices, value=new_choices[0])
+        return gr.Dropdown(choices=RES_CHOICES[category], value=RES_CHOICES[category][0])
 
-    res_category.change(
-        fn=update_resolution_list,
-        inputs=[res_category],
-        outputs=[resolution_dropdown]
-    )
+    res_category.change(update_resolution_list, inputs=res_category, outputs=resolution_dropdown)
     
-    # 2. 切换自定义分辨率显示
-    def toggle_custom(is_custom):
-        return gr.Row(visible=is_custom), gr.Dropdown(interactive=not is_custom)
-
     use_custom_res.change(
-        fn=toggle_custom,
-        inputs=[use_custom_res],
-        outputs=[custom_res_row, resolution_dropdown]
+        lambda x: (gr.Row(visible=x), gr.Dropdown(interactive=not x)),
+        inputs=use_custom_res, outputs=[custom_res_row, resolution_dropdown]
     )
     
-    # 3. 切换种子输入框
-    random_seed.change(
-        fn=lambda x: gr.Number(visible=not x),
-        inputs=[random_seed],
-        outputs=[seed_input]
-    )
+    random_seed.change(lambda x: gr.Number(visible=not x), inputs=random_seed, outputs=seed_input)
     
-    # 4. 生成按钮点击
     generate_btn.click(
         fn=generate_image,
-        inputs=[
-            prompt, resolution_dropdown, use_custom_res, width_slider, height_slider,
-            steps_slider, seed_input, random_seed, negative_prompt, output_gallery
-        ],
+        inputs=[prompt, resolution_dropdown, use_custom_res, width_slider, height_slider, steps_slider, seed_input, random_seed, negative_prompt, output_gallery],
         outputs=[output_gallery, last_seed_display]
     )
     
-    # 5. 清空历史
-    clear_history_btn.click(
-        fn=lambda: ([], ""),
-        outputs=[output_gallery, last_seed_display]
-    )
+    clear_btn.click(lambda: ([], ""), outputs=[output_gallery, last_seed_display])
 
 if __name__ == "__main__":
     demo.launch()