Update app.py

app.py

@@ -17,6 +17,7 @@ from diffusers import ZImagePipeline
 from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
 
 # ==================== Environment Variables ==================================
+# 使用 Z-Image 基础模型
 MODEL_PATH = os.environ.get("MODEL_PATH", "Tongyi-MAI/Z-Image")
 ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "true").lower() == "true"
 ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "true").lower() == "true"
@@ -105,7 +106,6 @@ def load_models(model_path, enable_compile=False, attention_backend="native"):
 
     pipe.transformer = transformer
     
-    # 尝试设置指定的 Backend，如果失败则回退
     try:
         pipe.transformer.set_attention_backend(attention_backend)
     except Exception as e:
@@ -121,7 +121,7 @@ def load_models(model_path, enable_compile=False, attention_backend="native"):
 
 def generate_image(pipe, prompt, negative_prompt="", width=1024, height=1024, seed=42, guidance_scale=4.0, num_inference_steps=30, shift=3.0, max_sequence_length=512, progress=gr.Progress(track_tqdm=True)):
     generator = torch.Generator("cuda").manual_seed(seed)
-    # Z-Image 使用 FlowMatch
+    # Z-Image 使用 FlowMatch 调度器
     scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)
     pipe.scheduler = scheduler
 
@@ -134,7 +134,7 @@ def generate_image(pipe, prompt, negative_prompt="", width=1024, height=1024, se
         num_inference_steps=num_inference_steps,
         generator=generator,
         max_sequence_length=max_sequence_length,
-        cfg_normalization=False, # 官方推荐
+        cfg_normalization=False,
     ).images[0]
 
     return image
@@ -145,13 +145,13 @@ def warmup_model(pipe, resolutions):
     for res_str in resolutions:
         try:
             w, h = get_resolution(res_str)
-            for i in range(1): # 减少基础模型的预热次数以节省资源
+            for i in range(1):
                 generate_image(pipe, prompt=dummy_prompt, width=w, height=h, num_inference_steps=28, guidance_scale=4.0, seed=42 + i)
         except Exception as e:
             print(f"Warmup failed for {res_str}: {e}")
     print("Warmup completed.")
 
-# Global Pipe Variable
+# 全局变量
 pipe = None
 
 def init_app():
@@ -161,7 +161,7 @@ def init_app():
         print(f"Model loaded: {MODEL_PATH}. Compile: {ENABLE_COMPILE}, Backend: {ATTENTION_BACKEND}")
 
         if ENABLE_WARMUP:
-            # 仅预热常用分辨率以加快启动
+            # 仅预热基础分辨率
             warmup_res = ["1024x1024 ( 1:1 )"]
             warmup_model(pipe, warmup_res)
 
@@ -197,22 +197,18 @@ def generate(prompt, negative_prompt, width=1024, height=1024, seed=42, steps=30
 
     return gallery_images, str(new_seed), int(new_seed)
 
-# Initialize
+# 初始化
 init_app()
 
-# ==================== AoTI Optimization ====================
+# ==================== AoTI 优化 ====================
 if pipe is not None:
     try:
-        # 针对 ZeroGPU 的优化配置
         pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
         spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
     except Exception as e:
         print(f"Warning: Failed to load AoTI blocks: {e}")
 
-# ==================== UI Construction ====================
-custom_css = ".fillable{max-width: 1230px !important}"
-
-with gr.Blocks(title="Z-Image Demo", css=custom_css) as demo:
+with gr.Blocks(title="Z-Image Demo") as demo:
     gr.Markdown(
         """<div align="center">
 # ⚡️ Z-Image Generation Demo
@@ -223,8 +219,8 @@ with gr.Blocks(title="Z-Image Demo", css=custom_css) as demo:
 
     with gr.Row():
         with gr.Column(scale=1):
-            prompt_input = gr.Textbox(label="Prompt (提示词)", lines=3, placeholder="输入你想要生成的图像描述...")
-            negative_prompt_input = gr.Textbox(label="Negative Prompt (负面提示词)", lines=2, placeholder="输入你不想要出现在图像中的内容...")
+            prompt_input = gr.Textbox(label="Prompt (提示词)", lines=3, placeholder="输入图像描述...")
+            negative_prompt_input = gr.Textbox(label="Negative Prompt (负面提示词)", lines=2, placeholder="输入不想要出现的内容...")
             
             with gr.Row():
                 width = gr.Slider(label="Width (宽)", minimum=512, maximum=2048, value=1024, step=64)
@@ -248,7 +244,7 @@ with gr.Blocks(title="Z-Image Demo", css=custom_css) as demo:
             output_gallery = gr.Gallery(
                 label="Generated Images", columns=1, rows=1, height=600, object_fit="contain", format="png", interactive=False
             )
-            used_seed = gr.Textbox(label="Seed Used (本次使用的种子)", interactive=False)
+            used_seed = gr.Textbox(label="Seed Used (使用的种子)", interactive=False)
 
     generate_btn.click(
         generate,