Update app.py

app.py

@@ -1,13 +1,13 @@
-# ===== 必须首先导入spaces =====
+# ===== ZeroGPU 超时优化终极版 =====
+
 try:
     import spaces
     SPACES_AVAILABLE = True
-    print("✅ Spaces available - ZeroGPU mode")
+    print("✅ ZeroGPU mode enabled")
 except ImportError:
     SPACES_AVAILABLE = False
-    print("⚠️ Spaces not available - running in regular mode")
+    print("⚠️ Running in regular mode")
 
-# ===== 其他导入 =====
 import os
 from datetime import datetime
 import random
@@ -18,190 +18,126 @@ from PIL import Image
 import traceback
 import numpy as np
 import gc
+import warnings
+warnings.filterwarnings('ignore')
 
-# ===== 配置常量 =====
-COMPEL_AVAILABLE = False
-print("⚠️ Compel disabled for FLUX compatibility")
+# ===== 配置 =====
+FIXED_MODEL = "aoxo/flux.1dev-abliterated"
+SAVE_DIR = "generated_images"
+os.makedirs(SAVE_DIR, exist_ok=True)
 
 STYLE_PRESETS = {
     "None": "",
-    "Realistic": "photorealistic, 8k, ultra-detailed, cinematic lighting, masterpiece",
-    "Anime": "anime style, detailed, high quality, masterpiece, best quality",
-    "Comic": "comic book style, bold outlines, vibrant colors, cel shading",
-    "Watercolor": "watercolor illustration, soft gradients, pastel palette"
+    "Realistic": "photorealistic, detailed",
+    "Anime": "anime style, high quality",
+    "Comic": "comic book style",
+    "Watercolor": "watercolor painting"
 }
 
-FIXED_MODEL = "aoxo/flux.1dev-abliterated"
-
-QUALITY_ENHANCERS = [
-    "detailed anatomy", "perfect anatomy", "soft skin",
-    "high resolution", "masterpiece", "best quality",
-    "professional photography", "natural lighting"
-]
-
-STYLE_ENHANCERS = {
-    "Realistic": ["photorealistic", "ultra realistic", "natural lighting"],
-    "Anime": ["anime style", "high quality anime", "detailed eyes"],
-    "Comic": ["comic book style", "bold outlines", "vibrant colors"],
-    "Watercolor": ["watercolor style", "artistic", "soft gradients"]
-}
-
-SAVE_DIR = "generated_images"
-os.makedirs(SAVE_DIR, exist_ok=True)
-
 # ===== 全局变量 =====
 pipeline = None
 device = None
 model_loaded = False
 
 
-# ===== 工具函数(必须在装饰器之前定义) =====
 def cleanup_memory():
-    """清理GPU内存"""
+    """激进的内存清理"""
+    gc.collect()
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.synchronize()
-    gc.collect()
 
 
-def enhance_prompt(prompt: str, style: str) -> str:
-    """增强提示词"""
-    quality_terms = ", ".join(QUALITY_ENHANCERS[:3])
-    
-    style_terms = ""
-    if style in STYLE_ENHANCERS:
-        style_terms = ", " + ", ".join(STYLE_ENHANCERS[style][:2])
-    
+def apply_spaces_decorator(func):
+    """ZeroGPU 装饰器 - 60秒限制"""
+    if SPACES_AVAILABLE:
+        # ZeroGPU 实际只给 60 秒!
+        return spaces.GPU(duration=60)(func)
+    return func
+
+
+def enhance_prompt_minimal(prompt: str, style: str) -> str:
+    """最小化提示词增强 - 严格控制长度"""
     style_suffix = STYLE_PRESETS.get(style, "")
     
-    enhanced_parts = [prompt.strip()]
-    
     if style_suffix:
-        enhanced_parts.append(style_suffix)
-    
-    if style_terms:
-        enhanced_parts.append(style_terms.lstrip(", "))
-    
-    enhanced_parts.append(quality_terms)
-    
-    enhanced_prompt = ", ".join(filter(None, enhanced_parts))
+        enhanced = f"{prompt}, {style_suffix}, masterpiece"
+    else:
+        enhanced = f"{prompt}, masterpiece"
     
-    if len(enhanced_prompt) > 800:
-        enhanced_prompt = enhanced_prompt[:800]
+    # CLIP 硬限制: 77 tokens ≈ 200-250 字符
+    if len(enhanced) > 200:
+        enhanced = prompt[:180] + ", masterpiece"
+        print(f"⚠️ Prompt truncated to fit CLIP limit")
     
-    return enhanced_prompt
+    return enhanced
 
 
-def create_metadata_content(prompt, enhanced_prompt, seed, steps, cfg_scale, width, height, style):
-    """创建元数据内容"""
-    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    return f"""Generated Image Metadata
-======================
-Timestamp: {timestamp}
-Original Prompt: {prompt}
-Enhanced Prompt: {enhanced_prompt}
-Seed: {seed}
-Steps: {steps}
-CFG Scale: {cfg_scale}
-Dimensions: {width}x{height}
-Style: {style}
-Model: FLUX.1-dev
-"""
-
-
-# ===== 装饰器定义(必须在使用之前) =====
-def apply_spaces_decorator(func):
-    """应用 spaces 装饰器,增加更长的超时时间"""
-    if SPACES_AVAILABLE:
-        return spaces.GPU(duration=120)(func)
-    return func
-
-
-# ===== 模型相关函数 =====
+# ===== 分离模型初始化（不使用 GPU 装饰器）=====
 def initialize_model():
-    """优化的模型初始化函数"""
+    """模型初始化 - 不占用 GPU 时间"""
     global pipeline, device, model_loaded
     
     if model_loaded and pipeline is not None:
         return True
     
     try:
-        cleanup_memory()
-        
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"🖥️ Using device: {device}")
+        print(f"🖥️ Device: {device}")
         
-        print(f"📦 Loading fixed model: {FIXED_MODEL}")
+        print(f"📦 Loading: {FIXED_MODEL}")
         
         pipeline = AutoPipelineForText2Image.from_pretrained(
             FIXED_MODEL,
-            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-            variant=None,
-            use_safetensors=True
+            dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            use_safetensors=True,
         )
         
         pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
             pipeline.scheduler.config
         )
+        
+        # 关键优化：不用 offload，直接全部加载
         pipeline = pipeline.to(device)
         
+        # 只保留最必要的优化
         if torch.cuda.is_available():
-            # 关键优化:使用sequential代替model cpu offload
-            pipeline.enable_sequential_cpu_offload()
             pipeline.enable_vae_slicing()
             pipeline.enable_vae_tiling()
         
-        print("✅ Model initialization complete (Optimized)")
+        print("✅ Model ready")
         model_loaded = True
         return True
         
     except Exception as e:
-        print(f"❌ Critical model loading error: {e}")
-        print(traceback.format_exc())
-        cleanup_memory()
-        model_loaded = False
+        print(f"❌ Init failed: {e}")
         return False
 
 
 @apply_spaces_decorator
-def generate_image(prompt: str, style: str, negative_prompt: str = "", 
-                  steps: int = 15, cfg_scale: float = 3.5, 
-                  seed: int = -1, width: int = 1024, height: int = 1024, 
-                  progress=gr.Progress()):
-    """图像生成函数(优化版本)"""
+def generate_image_fast(prompt: str, style: str, negative_prompt: str, 
+                       steps: int, cfg_scale: float, seed: int, 
+                       width: int, height: int):
+    """超快速生成 - 必须在 60 秒内完成"""
     try:
-        if not prompt or prompt.strip() == "":
-            return None, "", "❌ Please enter a prompt"
-        
-        # 优化的参数限制
-        steps = max(10, min(steps, 25))
-        width = min(width, 1024)
-        height = min(height, 1024)
-        
-        progress(0.1, desc="Initializing model...")
-        if not initialize_model():
-            cleanup_memory()
-            return None, "", "❌ Failed to initialize model"
-        
-        progress(0.2, desc="Processing prompt...")
+        print(f"⏱️ GPU timer started (60s limit)")
         
         if seed == -1:
-            seed = random.randint(0, np.iinfo(np.int32).max)
+            seed = random.randint(0, 999999)
         
-        enhanced_prompt = enhance_prompt(prompt.strip(), style)
+        enhanced_prompt = enhance_prompt_minimal(prompt, style)
         
-        if not negative_prompt.strip():
-            negative_prompt = "(low quality, worst quality:1.4), (bad anatomy, bad hands:1.2), blurry, deformed"
+        if not negative_prompt:
+            negative_prompt = "low quality, blurry"
         
         generator = torch.Generator("cpu").manual_seed(seed)
         
-        progress(0.4, desc="Starting generation...")
-        print(f"🔥 Inference: steps={steps}, guidance={cfg_scale}, size={width}x{height}")
+        print(f"🚀 Generating: {steps} steps, {width}x{height}")
         
         cleanup_memory()
         
-        # 关键改动:提高max_sequence_length
-        with torch.no_grad():
+        # 极简推理参数
+        with torch.inference_mode():  # 比 no_grad 更快
             result = pipeline(
                 prompt=enhanced_prompt,
                 negative_prompt=negative_prompt,
@@ -209,236 +145,187 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "",
                 guidance_scale=cfg_scale,
                 width=width,
                 height=height,
-                max_sequence_length=512,  # 从256改到512
                 generator=generator,
                 output_type="pil"
             )
         
         image = result.images[0]
-        print("✅ Inference complete")
+        del result
+        cleanup_memory()
         
-        progress(0.9, desc="Finalizing...")
+        print(f"✅ Done in <60s")
+        return image, seed
         
-        del result
+    except Exception as e:
         cleanup_memory()
+        print(f"❌ Error: {e}")
+        raise e
+
+
+def generate_wrapper(prompt, style, neg_prompt, steps, cfg, seed, size_preset, progress=gr.Progress()):
+    """包装函数 - 处理 UI 逻辑"""
+    try:
+        if not prompt.strip():
+            return None, "❌ Enter a prompt", "", None
         
-        filename = f"IMG_{seed}.png"
-        filepath = os.path.join(SAVE_DIR, filename)
-        image.save(filepath, format="PNG", optimize=True)
+        # 解析尺寸
+        if size_preset == "512x512 (Ultra Fast)":
+            width = height = 512
+        elif size_preset == "768x768 (Fast)":
+            width = height = 768
+        else:
+            width = height = 1024
         
-        metadata_content = create_metadata_content(
-            prompt, enhanced_prompt, seed, steps, cfg_scale, width, height, style
+        # 限制步数
+        steps = max(8, min(steps, 15))
+        
+        progress(0.1, desc="Initializing...")
+        
+        # 预加载模型（不计入 GPU 时间）
+        if not initialize_model():
+            return None, "❌ Model init failed", "", None
+        
+        progress(0.2, desc="Generating (30-50s)...")
+        
+        # 调用 GPU 函数
+        image, actual_seed = generate_image_fast(
+            prompt, style, neg_prompt, steps, cfg, seed, width, height
         )
         
-        progress(1.0, desc="Complete!")
+        progress(0.9, desc="Saving...")
         
-        generation_info = f"Prompt: {prompt}\nSeed: {seed} | Size: {width}×{height} | Steps: {steps} | CFG: {cfg_scale}"
+        filename = f"IMG_{actual_seed}.png"
+        filepath = os.path.join(SAVE_DIR, filename)
+        image.save(filepath)
         
-        return image, generation_info, metadata_content
+        metadata = f"""Generated: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+Prompt: {prompt}
+Style: {style}
+Seed: {actual_seed}
+Steps: {steps} | CFG: {cfg}
+Size: {width}x{height}
+"""
         
-    except torch.cuda.OutOfMemoryError as e:
-        cleanup_memory()
-        error_msg = "❌ GPU memory insufficient. Try 768x768 or fewer steps."
-        print(f"CUDA OOM: {error_msg}")
-        return None, "", error_msg
+        info = f"Seed: {actual_seed} | {width}×{height} | {steps} steps"
+        
+        progress(1.0, desc="Complete!")
+        
+        return image, info, metadata, image
         
     except Exception as e:
         cleanup_memory()
-        error_msg = str(e)
-        print(f"❌ Generation error: {error_msg}")
-        print(traceback.format_exc())
-        return None, "", f"❌ Generation failed: {error_msg}"
+        error_msg = f"Generation failed: {str(e)[:100]}"
+        print(f"❌ {error_msg}")
+        return None, error_msg, "", None
 
 
-# ===== CSS样式 =====
-css = """
-/* 保持原有CSS不变 */
-.gradio-container {
-    max-width: 100% !important;
-    margin: 0 !important;
-    padding: 0 !important;
-    background: linear-gradient(135deg, #e6a4f2 0%, #1197e4 100%) !important;
-    min-height: 100vh !important;
-}
-
-.main-content {
-    background: rgba(255, 255, 255, 0.95) !important;
-    border-radius: 20px !important;
-    padding: 20px !important;
-    margin: 15px !important;
-    box-shadow: 0 10px 25px rgba(0,0,0,0.2) !important;
-}
-
-.title {
-    text-align: center !important;
-    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
-    -webkit-background-clip: text !important;
-    -webkit-text-fill-color: transparent !important;
-    font-size: 2rem !important;
-    margin-bottom: 15px !important;
-    font-weight: bold !important;
-}
-
-.generate-btn {
-    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
-    color: white !important;
-    border: none !important;
-    padding: 15px 25px !important;
-    border-radius: 25px !important;
-    font-size: 16px !important;
-    font-weight: bold !important;
-    width: 100% !important;
-}
-"""
-
-
-# ===== 创建UI =====
+# ===== UI =====
 def create_interface():
-    with gr.Blocks(css=css, title="NSFW FLUX Image Generator") as interface:
-        with gr.Column(elem_classes=["main-content"]):
-            gr.HTML('<div class="title">NSFW FLUX Image Generator</div>')
-            
-            gr.HTML('<div class="warning-box">⚠️ 18+ CONTENT WARNING ⚠️</div>')
-            
-            with gr.Row():
-                with gr.Column(scale=2):
-                    prompt_input = gr.Textbox(
-                        label="Main Prompt",
-                        placeholder="beautiful woman, detailed portrait...",
-                        lines=6,
-                        elem_classes=["prompt-box"]
-                    )
-                    
-                    gr.HTML('''
-                        <div style="background: rgba(255, 193, 7, 0.1); padding: 10px; border-radius: 8px; margin: 10px 0;">
-                            <small><strong>💡 Optimized Settings:</strong><br>
-                            • Max sequence length: 512 tokens (supports longer prompts)<br>
-                            • Recommended steps: 15-20 for best speed/quality<br>
-                            • Try 768x768 for faster generation</small>
-                        </div>
-                    ''')
-                    
-                    negative_prompt_input = gr.Textbox(
-                        label="Negative Prompt (Optional)",
-                        placeholder="low quality, blurry...",
-                        lines=3,
-                        elem_classes=["prompt-box"]
-                    )
+    with gr.Blocks(title="Fast FLUX Generator") as interface:
+        gr.HTML('<h1 style="text-align:center">⚡ Fast FLUX Generator</h1>')
+        
+        gr.HTML('''
+            <div style="background:#fff3cd;padding:10px;border-radius:8px;margin:10px 0;">
+                <strong>⚠️ ZeroGPU Limits:</strong><br>
+                • 60 second GPU timeout (hard limit)<br>
+                • Recommended: 512x512 or 768x768, 10-15 steps<br>
+                • Keep prompts under 200 characters
+            </div>
+        ''')
+        
+        with gr.Row():
+            with gr.Column(scale=2):
+                prompt_input = gr.Textbox(
+                    label="Prompt (keep it short!)",
+                    placeholder="woman, portrait, detailed",
+                    lines=4,
+                    max_lines=4
+                )
                 
-                with gr.Column(scale=1):
-                    with gr.Group():
-                        style_input = gr.Radio(
-                            label="Style Preset",
-                            choices=list(STYLE_PRESETS.keys()),
-                            value="Realistic"
-                        )
-                    
-                    with gr.Group():
-                        seed_input = gr.Number(
-                            label="Seed (-1 for random)",
-                            value=-1,
-                            precision=0
-                        )
-                    
-                    with gr.Group():
-                        size_preset = gr.Radio(
-                            label="Size (smaller = faster)",
-                            choices=["768x768", "1024x1024"],
-                            value="768x768"
-                        )
-                    
-                    with gr.Group():
-                        steps_input = gr.Slider(
-                            label="Steps (15-20 recommended)",
-                            minimum=10,
-                            maximum=25,
-                            value=15,
-                            step=1
-                        )
-                        
-                        cfg_input = gr.Slider(
-                            label="CFG Scale",
-                            minimum=1.0,
-                            maximum=15.0,
-                            value=3.5,
-                            step=0.1
-                        )
-                    
-                    generate_button = gr.Button(
-                        "GENERATE",
-                        elem_classes=["generate-btn"],
-                        variant="primary"
-                    )
-            
-            image_output = gr.Image(
-                label="Generated Image",
-                elem_classes=["image-output"],
-                show_label=False
-            )
-            
-            generation_info = gr.Textbox(
-                label="Generation Info",
-                interactive=False,
-                visible=False
-            )
-            
-            metadata_content = gr.Textbox(visible=False)
-            current_seed = gr.Number(visible=False)
-            current_image = gr.Image(visible=False)
-            
-            with gr.Row(visible=False) as download_row:
-                download_image_btn = gr.Button("Save Image", size="sm")
-                download_metadata_btn = gr.Button("Save Metadata", size="sm")
-        
-        def parse_size(size_str):
-            """解析尺寸字符串"""
-            size = int(size_str.split('x')[0])
-            return size, size
-        
-        def on_generate(prompt, style, neg_prompt, steps, cfg, seed, size_preset):
-            width, height = parse_size(size_preset)
-            image, info, metadata = generate_image(
-                prompt, style, neg_prompt, steps, cfg, seed, width, height
-            )
+                negative_prompt_input = gr.Textbox(
+                    label="Negative Prompt",
+                    placeholder="low quality, blurry",
+                    lines=2
+                )
             
-            if image is not None:
-                try:
-                    actual_seed = seed if seed != -1 else int(info.split("Seed:")[1].split("|")[0].strip())
-                except:
-                    actual_seed = seed if seed != -1 else random.randint(0, 999999)
+            with gr.Column(scale=1):
+                style_input = gr.Radio(
+                    label="Style",
+                    choices=list(STYLE_PRESETS.keys()),
+                    value="Realistic"
+                )
                 
-                return (
-                    image, info, metadata, actual_seed, image,
-                    gr.update(visible=True), gr.update(visible=True)
+                seed_input = gr.Number(
+                    label="Seed (-1 = random)",
+                    value=-1,
+                    precision=0
                 )
-            else:
-                return (
-                    None, info, "", 0, None,
-                    gr.update(visible=False), gr.update(visible=False)
+                
+                size_preset = gr.Radio(
+                    label="Size (smaller = faster)",
+                    choices=[
+                        "512x512 (Ultra Fast)",
+                        "768x768 (Fast)",
+                        "1024x1024 (Slow)"
+                    ],
+                    value="768x768 (Fast)"
+                )
+                
+                steps_input = gr.Slider(
+                    label="Steps (10-15 recommended)",
+                    minimum=8,
+                    maximum=15,
+                    value=12,
+                    step=1
+                )
+                
+                cfg_input = gr.Slider(
+                    label="CFG Scale",
+                    minimum=1.0,
+                    maximum=10.0,
+                    value=3.5,
+                    step=0.5
+                )
+                
+                generate_button = gr.Button(
+                    "🚀 GENERATE (30-50s)",
+                    variant="primary",
+                    size="lg"
                 )
         
+        image_output = gr.Image(label="Result", show_label=False)
+        
+        generation_info = gr.Textbox(
+            label="Info",
+            interactive=False,
+            visible=True
+        )
+        
+        metadata_content = gr.Textbox(visible=False)
+        current_image = gr.Image(visible=False)
+        
         generate_button.click(
-            fn=on_generate,
+            fn=generate_wrapper,
             inputs=[
-                prompt_input, style_input, negative_prompt_input, 
+                prompt_input, style_input, negative_prompt_input,
                 steps_input, cfg_input, seed_input, size_preset
             ],
             outputs=[
-                image_output, generation_info, metadata_content, 
-                current_seed, current_image, generation_info, download_row
+                image_output, generation_info, 
+                metadata_content, current_image
             ],
             show_progress=True
         )
         
         prompt_input.submit(
-            fn=on_generate,
+            fn=generate_wrapper,
             inputs=[
-                prompt_input, style_input, negative_prompt_input, 
+                prompt_input, style_input, negative_prompt_input,
                 steps_input, cfg_input, seed_input, size_preset
             ],
             outputs=[
-                image_output, generation_info, metadata_content, 
-                current_seed, current_image, generation_info, download_row
+                image_output, generation_info, 
+                metadata_content, current_image
             ],
             show_progress=True
         )
@@ -446,14 +333,17 @@ def create_interface():
     return interface
 
 
-# ===== 启动应用 =====
 if __name__ == "__main__":
-    print("🎨 Starting NSFW FLUX Image Generator (Optimized)...")
-    print(f"🔧 Fixed Model: {FIXED_MODEL}")
-    print(f"🔧 CUDA: {'✅ Available' if torch.cuda.is_available() else '❌ Not Available'}")
+    print("🚀 Starting Fast FLUX Generator")
+    print(f"🔧 Model: {FIXED_MODEL}")
+    print(f"🔧 CUDA: {torch.cuda.is_available()}")
+    
+    # 预加载模型
+    print("📦 Pre-loading model...")
+    initialize_model()
     
     app = create_interface()
-    app.queue(max_size=5, default_concurrency_limit=1)
+    app.queue(max_size=3, default_concurrency_limit=1)
     
     app.launch(
         server_name="0.0.0.0",