Update app.py

app.py

@@ -1,5 +1,125 @@
+# ===== 必须首先导入spaces =====
+try:
+    import spaces
+    SPACES_AVAILABLE = True
+    print("✅ Spaces available - ZeroGPU mode")
+except ImportError:
+    SPACES_AVAILABLE = False
+    print("⚠️ Spaces not available - running in regular mode")
+
+# ===== 其他导入 =====
+import os
+from datetime import datetime
+import random
+import torch
+import gradio as gr
+from diffusers import AutoPipelineForText2Image, FlowMatchEulerDiscreteScheduler
+from PIL import Image
+import traceback
+import numpy as np
+import gc
+
+# ===== 配置常量 =====
+COMPEL_AVAILABLE = False
+print("⚠️ Compel disabled for FLUX compatibility")
+
+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"
+}
+
+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内存"""
+    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])
+    
+    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))
+    
+    if len(enhanced_prompt) > 800:
+        enhanced_prompt = enhanced_prompt[:800]
+    
+    return enhanced_prompt
+
+
+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
+
+
+# ===== 模型相关函数 =====
 def initialize_model():
-    """优化后的模型初始化函数"""
+    """优化的模型初始化函数"""
     global pipeline, device, model_loaded
     
     if model_loaded and pipeline is not None:
@@ -13,12 +133,11 @@ def initialize_model():
         
         print(f"📦 Loading fixed model: {FIXED_MODEL}")
         
-        # 优化的模型加载
         pipeline = AutoPipelineForText2Image.from_pretrained(
             FIXED_MODEL,
-            torch_dtype=torch.bfloat16,
-            use_safetensors=True,
-            variant=None
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            variant=None,
+            use_safetensors=True
         )
         
         pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
@@ -26,14 +145,11 @@ def initialize_model():
         )
         pipeline = pipeline.to(device)
         
-        # 关键优化:使用sequential而不是model cpu offload
         if torch.cuda.is_available():
-            pipeline.to(device) # 改这里!
+            # 关键优化:使用sequential代替model cpu offload
+            pipeline.enable_sequential_cpu_offload()
             pipeline.enable_vae_slicing()
             pipeline.enable_vae_tiling()
-            
-            # 可选:如果内存充足,可以完全不用offload
-            # pipeline.to(device)  # 全部加载到GPU,最快但吃内存
         
         print("✅ Model initialization complete (Optimized)")
         model_loaded = True
@@ -49,16 +165,16 @@ def initialize_model():
 
 @apply_spaces_decorator
 def generate_image(prompt: str, style: str, negative_prompt: str = "", 
-                  steps: int = 20, cfg_scale: float = 3.5, 
+                  steps: int = 15, cfg_scale: float = 3.5, 
                   seed: int = -1, width: int = 1024, height: int = 1024, 
                   progress=gr.Progress()):
-    """优化后的图像生成函数"""
+    """图像生成函数(优化版本)"""
     try:
         if not prompt or prompt.strip() == "":
             return None, "", "❌ Please enter a prompt"
         
         # 优化的参数限制
-        steps = max(10, min(steps, 25))  # 降低最大步数
+        steps = max(10, min(steps, 25))
         width = min(width, 1024)
         height = min(height, 1024)
         
@@ -72,11 +188,10 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "",
         if seed == -1:
             seed = random.randint(0, np.iinfo(np.int32).max)
         
-        # 不要过度增强提示词
         enhanced_prompt = enhance_prompt(prompt.strip(), style)
         
         if not negative_prompt.strip():
-            negative_prompt = "(low quality, worst quality:1.4), blurry, deformed"
+            negative_prompt = "(low quality, worst quality:1.4), (bad anatomy, bad hands:1.2), blurry, deformed"
         
         generator = torch.Generator("cpu").manual_seed(seed)
         
@@ -94,7 +209,7 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "",
                 guidance_scale=cfg_scale,
                 width=width,
                 height=height,
-                max_sequence_length=512,  # 从256改到512!
+                max_sequence_length=512,  # 从256改到512
                 generator=generator,
                 output_type="pil"
             )
@@ -123,7 +238,7 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "",
         
     except torch.cuda.OutOfMemoryError as e:
         cleanup_memory()
-        error_msg = "❌ GPU memory insufficient. Try: 768x768 or fewer steps"
+        error_msg = "❌ GPU memory insufficient. Try 768x768 or fewer steps."
         print(f"CUDA OOM: {error_msg}")
         return None, "", error_msg
         
@@ -135,48 +250,214 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "",
         return None, "", f"❌ Generation failed: {error_msg}"
 
 
-def enhance_prompt(prompt: str, style: str) -> str:
-    """优化的提示词增强"""
-    # 减少质量词,避免提示词过长
-    quality_terms = ", ".join(QUALITY_ENHANCERS[:3])  # 从5减到3
-    
-    style_terms = ""
-    if style in STYLE_ENHANCERS:
-        style_terms = ", " + ", ".join(STYLE_ENHANCERS[style][:2])  # 从3减到2
-    
-    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))
-    
-    # 提高长度限制,但不要太长
-    if len(enhanced_prompt) > 800:  # 从500提高到800
-        enhanced_prompt = enhanced_prompt[:800]
+# ===== 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 =====
+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.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
+            )
+            
+            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)
+                
+                return (
+                    image, info, metadata, actual_seed, image,
+                    gr.update(visible=True), gr.update(visible=True)
+                )
+            else:
+                return (
+                    None, info, "", 0, None,
+                    gr.update(visible=False), gr.update(visible=False)
+                )
+        
+        generate_button.click(
+            fn=on_generate,
+            inputs=[
+                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
+            ],
+            show_progress=True
+        )
+        
+        prompt_input.submit(
+            fn=on_generate,
+            inputs=[
+                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
+            ],
+            show_progress=True
+        )
     
-    return enhanced_prompt
+    return interface
 
 
-# UI中的步数滑块也要调整
-steps_input = gr.Slider(
-    label="Steps (15-20 recommended for speed)",
-    minimum=10,
-    maximum=25,  # 从30降到25
-    value=15,    # 从20降到15
-    step=1
-)
-
-# 添加尺寸预设方便快速选择
-size_preset = gr.Radio(
-    label="Size Preset (smaller = faster)",
-    choices=["768x768 (Fast)", "1024x1024 (Quality)"],
-    value="768x768 (Fast)"
-)
+# ===== 启动应用 =====
+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'}")
+    
+    app = create_interface()
+    app.queue(max_size=5, default_concurrency_limit=1)
+    
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        share=False
+    )