Update app.py

app.py

@@ -17,6 +17,7 @@ from diffusers import FluxPipeline, FlowMatchEulerDiscreteScheduler
 from PIL import Image
 import traceback
 import numpy as np
+import gc  # 添加垃圾回收
 
 # 移除 Compel（FLUX 不兼容，简化处理）
 COMPEL_AVAILABLE = False
@@ -59,23 +60,36 @@ 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 initialize_model():
-    """优化的模型初始化函数（基于官方示例）"""
+    """优化的模型初始化函数（针对ZeroGPU优化）"""
     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"📦 Loading fixed model: {FIXED_MODEL}")
         
-        # 基础模型加载（官方示例）
+        # ZeroGPU优化的模型加载
         pipeline = FluxPipeline.from_pretrained(
             FIXED_MODEL,
-            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+            # 添加内存优化参数
+            variant=None,
+            use_safetensors=True
         )
         
         # 优化调度器（默认 FlowMatchEulerDiscreteScheduler）
@@ -84,30 +98,33 @@ def initialize_model():
         )
         pipeline = pipeline.to(device)
         
-        # 统一数据类型和优化（官方推荐）
+        # ZeroGPU专用优化
         if torch.cuda.is_available():
-            pipeline.vae.to(torch.bfloat16)
-            pipeline.enable_sequential_cpu_offload()  # 官方 VRAM 优化
-            pipeline.enable_vae_slicing()  # 保留 VAE 分片
+            # 使用更保守的内存设置
+            pipeline.enable_model_cpu_offload()  # 改用model_cpu_offload
+            pipeline.enable_vae_slicing()
+            pipeline.enable_vae_tiling()  # 添加VAE tiling
         
         # 移除 torch.compile（Spaces 不稳定）
-        print("✅ Model initialization complete (no torch.compile for stability)")
+        print("✅ Model initialization complete (ZeroGPU optimized)")
         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
         return False
 
 def enhance_prompt(prompt: str, style: str) -> str:
     """增强提示词"""
-    quality_terms = ", ".join(QUALITY_ENHANCERS)
+    # 限制质量词数量，避免过长
+    quality_terms = ", ".join(QUALITY_ENHANCERS[:5])  # 只取前5个
     
     style_terms = ""
     if style in STYLE_ENHANCERS:
-        style_terms = ", " + ", ".join(STYLE_ENHANCERS[style])
+        style_terms = ", " + ", ".join(STYLE_ENHANCERS[style][:3])  # 只取前3个
     
     style_suffix = STYLE_PRESETS.get(style, "")
     
@@ -122,12 +139,18 @@ def enhance_prompt(prompt: str, style: str) -> str:
     enhanced_parts.append(quality_terms)
     
     enhanced_prompt = ", ".join(filter(None, enhanced_parts))
+    
+    # 限制总长度，避免超出模型限制
+    if len(enhanced_prompt) > 500:
+        enhanced_prompt = enhanced_prompt[:500] + "..."
+    
     return enhanced_prompt
 
 def apply_spaces_decorator(func):
-    """应用 spaces 装饰器"""
+    """应用 spaces 装饰器，增加更长的超时时间"""
     if SPACES_AVAILABLE:
-        return spaces.GPU(duration=60)(func)
+        # 增加超时时间到120秒，并设置更大的内存限制
+        return spaces.GPU(duration=120)(func)
     return func
 
 def create_metadata_content(prompt, enhanced_prompt, seed, steps, cfg_scale, width, height, style):
@@ -147,20 +170,26 @@ Model: FLUX.1-dev
 """
 
 @apply_spaces_decorator
-def generate_image(prompt: str, style: str, negative_prompt: str = "", steps: int = 28, cfg_scale: float = 3.5, 
+def generate_image(prompt: str, style: str, negative_prompt: str = "", steps: int = 20, cfg_scale: float = 3.5, 
                   seed: int = -1, width: int = 1024, height: int = 1024, progress=gr.Progress()):
-    """图像生成函数（基于官方示例）"""
-    if not prompt or prompt.strip() == "":
-        return None, "", ""
-    
-    # 初始化模型
-    progress(0.1, desc="Loading model...")
-    if not initialize_model():
-        return None, "", "❌ Failed to load model"
-    
-    progress(0.3, desc="Processing prompt...")
-    
+    """图像生成函数（ZeroGPU优化版本）"""
     try:
+        if not prompt or prompt.strip() == "":
+            return None, "", "❌ Please enter a prompt"
+        
+        # 参数验证和限制
+        steps = max(10, min(steps, 30))  # 限制步数范围，避免超时
+        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...")
+        
         # 处理 seed
         if seed == -1:
             seed = random.randint(0, np.iinfo(np.int32).max)
@@ -175,29 +204,39 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "", steps: in
         # 生成参数（官方示例：generator 用 cpu）
         generator = torch.Generator("cpu").manual_seed(seed)
         
-        progress(0.5, desc="Generating image...")
-        print(f"🔄 Starting inference: prompt='{enhanced_prompt[:50]}...', steps={steps}, guidance={cfg_scale}")
+        progress(0.4, desc="Starting generation...")
+        print(f"🔥 Starting inference: steps={steps}, guidance={cfg_scale}, size={width}x{height}")
+        
+        # 清理内存
+        cleanup_memory()
         
         # 直接使用标准 pipeline（无 Compel，添加 max_sequence_length）
-        result = pipeline(
-            prompt=enhanced_prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=steps,
-            guidance_scale=cfg_scale,
-            width=width,
-            height=height,
-            max_sequence_length=512,  # 官方必须参数！
-            generator=generator
-        )
+        with torch.no_grad():  # 确保不计算梯度
+            result = pipeline(
+                prompt=enhanced_prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                max_sequence_length=256,  # 减少序列长度，节省内存
+                generator=generator,
+                output_type="pil"
+            )
+        
         image = result.images[0]
         print("✅ Inference complete")
         
-        progress(0.9, desc="Saving image...")
+        progress(0.9, desc="Finalizing...")
+        
+        # 立即清理内存
+        del result
+        cleanup_memory()
         
         # 保存图像
         filename = f"IMG_{seed}.png"
         filepath = os.path.join(SAVE_DIR, filename)
-        image.save(filepath, quality=95, optimize=True)
+        image.save(filepath, format="PNG", optimize=True)
         
         # 创建元数据内容
         metadata_content = create_metadata_content(
@@ -211,13 +250,20 @@ def generate_image(prompt: str, style: str, negative_prompt: str = "", steps: in
         
         return image, generation_info, metadata_content
         
+    except torch.cuda.OutOfMemoryError as e:
+        cleanup_memory()
+        error_msg = "❌ GPU memory insufficient. Try reducing image size or steps."
+        print(f"CUDA OOM: {error_msg}")
+        return None, "", error_msg
+        
     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}"
-        
-# ===== CSS 样式（保持不变） =====
+
+# ===== CSS 样式（保持不变）=====
 css = """
 /* 全局容器 */
 .gradio-container {
@@ -466,33 +512,33 @@ def create_interface():
                             precision=0
                         )
                     
-                    # 宽度选择
+                    # 宽度选择（降低最大值）
                     with gr.Group(elem_classes=["controls-section"]):
                         width_input = gr.Slider(
                             label="Width",
                             minimum=512,
-                            maximum=2048,
+                            maximum=1024,  # 降低最大值
                             value=1024,
                             step=64
                         )
                     
-                    # 高度选择
+                    # 高度选择（降低最大值）
                     with gr.Group(elem_classes=["controls-section"]):
                         height_input = gr.Slider(
                             label="Height",
                             minimum=512,
-                            maximum=2048,
+                            maximum=1024,  # 降低最大值
                             value=1024,
                             step=64
                         )
                     
-                    # 高级参数
+                    # 高级参数（调整默认值）
                     with gr.Group(elem_classes=["controls-section"]):
                         steps_input = gr.Slider(
                             label="Steps",
                             minimum=10,
-                            maximum=50,
-                            value=28,
+                            maximum=30,  # 降低最大值
+                            value=20,    # 降低默认值
                             step=1
                         )
                         
@@ -555,7 +601,10 @@ def create_interface():
             
             if image is not None:
                 # 提取实际使用的 seed
-                actual_seed = seed if seed != -1 else int(info.split("Seed:")[1].split("|")[0].strip())
+                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,  # 图片输出
@@ -582,7 +631,7 @@ def create_interface():
             if image_data is not None:
                 filename = f"IMG_{seed_val}.png"
                 filepath = os.path.join(SAVE_DIR, filename)
-                image_data.save(filepath, quality=95, optimize=True)
+                image_data.save(filepath, format="PNG", optimize=True)
                 return filepath
             return None
             
@@ -671,7 +720,7 @@ if __name__ == "__main__":
     print(f"🔧 CUDA: {'✅ Available' if torch.cuda.is_available() else '❌ Not Available'}")
     
     app = create_interface()
-    app.queue(max_size=10, default_concurrency_limit=2)
+    app.queue(max_size=5, default_concurrency_limit=1)  # 降低并发限制
     
     app.launch(
         server_name="0.0.0.0",