Update app.py

app.py

@@ -2,12 +2,68 @@ import torch
 import spaces
 import gradio as gr
 import time
+import re
 from diffusers import DiffusionPipeline
 import warnings
 
 # 忽略一些警告
 warnings.filterwarnings("ignore")
 
+# ==================== 分辨率配置 ====================
+RES_CHOICES = {
+    "1024": [
+        "1024x1024 (1:1)",
+        "1152x896 (9:7)",
+        "896x1152 (7:9)",
+        "1152x864 (4:3)",
+        "864x1152 (3:4)",
+        "1248x832 (3:2)",
+        "832x1248 (2:3)",
+        "1280x720 (16:9)",
+        "720x1280 (9:16)",
+        "1344x576 (21:9)",
+        "576x1344 (9:21)",
+    ],
+    "1280": [
+        "1280x1280 (1:1)",
+        "1440x1120 (9:7)",
+        "1120x1440 (7:9)",
+        "1472x1104 (4:3)",
+        "1104x1472 (3:4)",
+        "1536x1024 (3:2)",
+        "1024x1536 (2:3)",
+        "1536x864 (16:9)",
+        "864x1536 (9:16)",
+        "1680x720 (21:9)",
+        "720x1680 (9:21)",
+    ],
+    "1536": [
+        "1536x1536 (1:1)",
+        "1728x1344 (9:7)",
+        "1344x1728 (7:9)",
+        "1728x1296 (4:3)",
+        "1296x1728 (3:4)",
+        "1872x1248 (3:2)",
+        "1248x1872 (2:3)",
+        "2048x1152 (16:9)",
+        "1152x2048 (9:16)",
+        "2016x864 (21:9)",
+        "864x2016 (9:21)",
+    ],
+}
+
+def get_resolution(resolution_str):
+    """从分辨率字符串提取宽高"""
+    match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution_str)
+    if match:
+        width = int(match.group(1))
+        height = int(match.group(2))
+        # 确保是8的倍数
+        width = width - width % 8
+        height = height - height % 8
+        return width, height
+    return 1024, 1024
+
 # ==================== 模型加载优化 ====================
 print("🚀 Loading Z-Image-Turbo pipeline...")
 start_time = time.time()
@@ -46,8 +102,10 @@ print(f"✅ Pipeline loaded in {load_time:.2f} seconds!")
 @spaces.GPU
 def generate_image(
     prompt, 
-    height=768, 
-    width=768, 
+    resolution_choice,
+    custom_height=None,
+    custom_width=None,
+    use_custom_res=False,
     num_inference_steps=8, 
     seed=42, 
     randomize_seed=True,
@@ -63,16 +121,25 @@ def generate_image(
         prompt = prompt.strip()
         negative_prompt = negative_prompt_text.strip() if negative_prompt_text else None
         
-        # 自动调整尺寸为8的倍数
-        height = max(512, int(height) - int(height) % 8)
-        width = max(512, int(width) - int(width) % 8)
+        # 确定分辨率
+        if use_custom_res and custom_height and custom_width:
+            # 使用自定义分辨率
+            height = max(512, int(custom_height) - int(custom_height) % 8)
+            width = max(512, int(custom_width) - int(custom_width) % 8)
+        else:
+            # 使用预设分辨率
+            width, height = get_resolution(resolution_choice)
         
         # 限制最大尺寸防止OOM（T4 GPU限制）
-        MAX_SIZE = 1152
+        MAX_SIZE = 1536
         if height > MAX_SIZE or width > MAX_SIZE:
             height = min(height, MAX_SIZE)
             width = min(width, MAX_SIZE)
         
+        # 检查宽高积是否过大
+        if height * width > 2048 * 2048:
+            return None, 0, "❌ Resolution too large! Maximum total pixels: 2048x2048"
+        
         # 生成随机种子
         if randomize_seed:
             seed = torch.randint(0, 2**32 - 1, (1,)).item()
@@ -113,34 +180,11 @@ def generate_image(
         return image, seed, info
         
     except torch.cuda.OutOfMemoryError:
-        return None, seed, "💥 Out of memory! Try smaller image size (e.g., 768x768)"
+        return None, seed, "💥 Out of memory! Try smaller image size"
     except Exception as e:
         error_msg = str(e)[:100]
         return None, seed, f"❌ Error: {error_msg}"
 
-# ==================== 快速生成函数（用于示例） ====================
-def quick_generate(prompt):
-    """为示例使用的快速生成函数"""
-    return generate_image(
-        prompt=prompt,
-        height=768,
-        width=768,
-        num_inference_steps=8,
-        seed=42,
-        randomize_seed=False
-    )
-
-# ==================== 示例提示词 ====================
-examples = [
-    ["A beautiful Chinese woman in traditional red Hanfu, intricate embroidery, cinematic lighting, photorealistic"],
-    ["Cyberpunk city at night, neon lights, rainy streets, futuristic architecture, Blade Runner style"],
-    ["Majestic dragon flying over ancient Chinese palace, sunset, epic fantasy art"],
-    ["Cute anime girl with pink hair, cyberpunk street background, vibrant colors"],
-    ["Fantasy landscape with floating islands, waterfalls, magical creatures, digital painting"],
-    ["Portrait of a wise old samurai, detailed armor, cherry blossoms, studio lighting"],
-    ["Steampunk airship flying over Victorian London, gears and cogs, detailed"],
-]
-
 # ==================== 主题配置 ====================
 custom_theme = gr.themes.Soft(
     primary_hue="yellow",
@@ -161,8 +205,7 @@ custom_theme = gr.themes.Soft(
     slider_color="*primary_500",
 )
 
-# ==================== Gradio界面 ====================
-
+# ==================== 全局样式 ====================
 global_styles = """
 /* 全局样式 */
 body {
@@ -231,27 +274,39 @@ button.secondary:hover {
 
 /* 信息显示 */
 #info-display {
-    background: #f8fafc !important;
+    background: rgba(71,85,105) !important;
     border: 1px solid #e2e8f0 !important;
     border-radius: 8px !important;
     font-size: 0.9rem !important;
     padding: 0.75rem !important;
 }
 
-/* 提示框 */
-#tips-box {
-    background: #fefce8 !important;
-    border-left: 4px solid #fbbf24 !important;
-    padding: 12px !important;
-    border-radius: 8px !important;
-    margin-top: 1rem !important;
-}
-
 /* 进度条 */
 .progress-bar {
     background: linear-gradient(90deg, #fbbf24 0%, #f59e0b 100%) !important;
 }
 
+/* 分辨率选择器 */
+.resolution-option {
+    border: 1px solid #e2e8f0 !important;
+    border-radius: 8px !important;
+    padding: 0.5rem !important;
+    margin: 0.25rem !important;
+    cursor: pointer !important;
+    transition: all 0.2s ease !important;
+}
+
+.resolution-option:hover {
+    background: #fefce8 !important;
+    border-color: #fbbf24 !important;
+}
+
+.resolution-option.selected {
+    background: linear-gradient(135deg, #fef3c7 0%, #fde68a 100%) !important;
+    border-color: #f59e0b !important;
+    font-weight: 600 !important;
+}
+
 /* 移动端适配 */
 @media (max-width: 768px) {
     .gradio-container {
@@ -322,6 +377,18 @@ button.secondary:hover {
 }
 """
 
+# ==================== 示例提示词 ====================
+examples = [
+    ["A beautiful Chinese woman in traditional red Hanfu, intricate embroidery, cinematic lighting, photorealistic"],
+    ["Cyberpunk city at night, neon lights, rainy streets, futuristic architecture, Blade Runner style"],
+    ["Majestic dragon flying over ancient Chinese palace, sunset, epic fantasy art"],
+    ["Cute anime girl with pink hair, cyberpunk street background, vibrant colors"],
+    ["Fantasy landscape with floating islands, waterfalls, magical creatures, digital painting"],
+    ["Portrait of a wise old samurai, detailed armor, cherry blossoms, studio lighting"],
+    ["Steampunk airship flying over Victorian London, gears and cogs, detailed"],
+]
+
+# ==================== Gradio界面 ====================
 with gr.Blocks(
     theme=custom_theme,
     title="Z-Image-Turbo • Ultra-Fast AI Image Generator",
@@ -336,7 +403,7 @@ with gr.Blocks(
         <h1 style="font-size: 2.5rem; font-weight: 800; margin-bottom: 0.5rem; background: linear-gradient(135deg, #fbbf24 0%, #f59e0b 50%, #d97706 100%); -webkit-background-clip: text; -webkit-text-fill-color: transparent; background-clip: text;">
         🎨 Z-Image-Turbo</h1>
         <p style="font-size: 1.1rem; color: #64748b; margin-bottom: 1.5rem;">
-        Generate stunning images in <strong>8 steps</strong> • Optimized for speed</p>
+        Generate stunning images in <strong>8 steps</strong> • Optimized for speed • Multiple resolutions</p>
         </div>
         """,
         elem_id="header"
@@ -376,26 +443,52 @@ with gr.Blocks(
                         scale=1
                     )
             
-            # 高级设置
-            with gr.Accordion("⚙️ Advanced Settings", open=False):
-                with gr.Row():
-                    height = gr.Slider(
-                        label="Height",
+            # 分辨率设置
+            with gr.Accordion("🖼️ Resolution Settings", open=True):
+                # 分辨率预设选择
+                resolution_preset = gr.Radio(
+                    label="Preset Resolutions",
+                    choices=["1024", "1280", "1536"],
+                    value="1024",
+                    type="value",
+                    elem_classes="resolution-preset"
+                )
+                
+                # 分辨率选项动态更新
+                resolution_choices = gr.Dropdown(
+                    label="Select Resolution",
+                    choices=RES_CHOICES["1024"],
+                    value="1024x1024 (1:1)",
+                    elem_classes="resolution-choices"
+                )
+                
+                # 自定义分辨率选项
+                use_custom_res = gr.Checkbox(
+                    label="Use Custom Resolution",
+                    value=False,
+                    info="Manually set width and height"
+                )
+                
+                with gr.Row(visible=False) as custom_res_row:
+                    custom_width = gr.Slider(
+                        label="Width",
                         minimum=512,
-                        maximum=1152,
-                        value=768,
+                        maximum=1536,
+                        value=1024,
                         step=64,
-                        info="512-1152 pixels"
+                        info="512-1536 pixels"
                     )
-                    width = gr.Slider(
-                        label="Width",
+                    custom_height = gr.Slider(
+                        label="Height",
                         minimum=512,
-                        maximum=1152,
-                        value=768,
+                        maximum=1536,
+                        value=1024,
                         step=64,
-                        info="512-1152 pixels"
+                        info="512-1536 pixels"
                     )
-                
+            
+            # 高级设置
+            with gr.Accordion("⚙️ Advanced Settings", open=False):
                 num_inference_steps = gr.Slider(
                     label="Inference Steps",
                     minimum=4,
@@ -418,16 +511,6 @@ with gr.Blocks(
                         visible=False,
                         scale=2
                     )
-                
-                # 显示/隐藏种子输入
-                def toggle_seed_visibility(randomize):
-                    return gr.Number(visible=not randomize)
-                
-                randomize_seed.change(
-                    toggle_seed_visibility,
-                    inputs=[randomize_seed],
-                    outputs=[seed]
-                )
             
             # 示例提示词
             gr.Examples(
@@ -474,12 +557,13 @@ with gr.Blocks(
             # 建议提示
             gr.Markdown(
                 """
-                <div style="background: #fefce8; border-left: 4px solid #fbbf24; padding: 0.75rem; border-radius: 0.5rem; margin-top: 1rem;">
+                <div style="background: rgba(71,85,105); border-left: 4px solid #fbbf24; padding: 0.75rem; border-radius: 0.5rem; margin-top: 1rem;">
                 <strong>💡 Tips for best results:</strong>
                 <ul style="margin: 0.5rem 0 0 1rem; padding-left: 0.5rem;">
                 <li>Use descriptive prompts</li>
                 <li>Start with 8 steps for speed</li>
-                <li>768x768 is optimal for T4 GPU</li>
+                <li>1024x1024 is optimal for most cases</li>
+                <li>Higher resolutions require more VRAM</li>
                 <li>Use negative prompts to exclude unwanted elements</li>
                 </ul>
                 </div>
@@ -502,6 +586,36 @@ with gr.Blocks(
     
     # ==================== 事件处理 ====================
     
+    # 更新分辨率选项
+    def update_resolution_choices(preset):
+        return gr.Dropdown(choices=RES_CHOICES[preset], value=RES_CHOICES[preset][0])
+    
+    resolution_preset.change(
+        update_resolution_choices,
+        inputs=[resolution_preset],
+        outputs=[resolution_choices]
+    )
+    
+    # 显示/隐藏自定义分辨率
+    def toggle_custom_resolution(use_custom):
+        return gr.Row(visible=use_custom), gr.Dropdown(visible=not use_custom)
+    
+    use_custom_res.change(
+        toggle_custom_resolution,
+        inputs=[use_custom_res],
+        outputs=[custom_res_row, resolution_choices]
+    )
+    
+    # 显示/隐藏种子输入
+    def toggle_seed_visibility(randomize):
+        return gr.Number(visible=not randomize)
+    
+    randomize_seed.change(
+        toggle_seed_visibility,
+        inputs=[randomize_seed],
+        outputs=[seed]
+    )
+    
     # 复制种子按钮
     def copy_seed_to_clipboard(s):
         if s != 0:
@@ -517,7 +631,8 @@ with gr.Blocks(
     # 生成按钮点击
     generate_btn.click(
         fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, negative_prompt],
+        inputs=[prompt, resolution_choices, custom_height, custom_width, use_custom_res, 
+                num_inference_steps, seed, randomize_seed, negative_prompt],
         outputs=[output_image, used_seed, info_display],
         api_name="generate"
     )
@@ -525,7 +640,8 @@ with gr.Blocks(
     # Enter键提交
     prompt.submit(
         fn=generate_image,
-        inputs=[prompt, height, width, num_inference_steps, seed, randomize_seed, negative_prompt],
+        inputs=[prompt, resolution_choices, custom_height, custom_width, use_custom_res,
+                num_inference_steps, seed, randomize_seed, negative_prompt],
         outputs=[output_image, used_seed, info_display]
     )