Update app.py

app.py

@@ -1,31 +1,77 @@
-# app.py (Gradio界面)
+# app.py (Gradio界面) - 优化版
 import gradio as gr
-import requests
+import os
+import sys
+import logging
 from PIL import Image
 import numpy as np
 from sklearn.cluster import KMeans
 import time
 import random
-import os
-import torch
+import warnings
+
+# 抑制警告
+warnings.filterwarnings("ignore")
+
+# 设置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# 检查模型管理器是否存在
+try:
+    from models.model_manager import ModelManager
+    MODELS_AVAILABLE = True
+except ImportError as e:
+    logger.warning(f"模型管理器导入失败: {e}")
+    logger.info("将使用简化版本运行")
+    MODELS_AVAILABLE = False
+    ModelManager = None
 
-# 模型管理器
-from models.model_manager import ModelManager
+class SimpleModelManager:
+    """简化的模型管理器，用于演示模式"""
+    def __init__(self):
+        self.device = "cpu"
+        logger.info("使用简化模型管理器")
+    
+    def generate_caption(self, image):
+        return "这是一件时尚的服装设计，具有现代感和优雅的风格。"
+    
+    def generate_image(self, prompt, **kwargs):
+        # 创建占位图像
+        width, height = kwargs.get('width', 512), kwargs.get('height', 512)
+        color = (random.randint(100, 255), random.randint(100, 255), random.randint(100, 255))
+        img = Image.new('RGB', (width, height), color=color)
+        return img
+    
+    def cleanup(self):
+        pass
 
-# 初始化模型管理器
-model_manager = ModelManager()
+# 根据是否有模型选择管理器
+if MODELS_AVAILABLE:
+    try:
+        model_manager = ModelManager()
+        logger.info("使用完整模型管理器")
+    except Exception as e:
+        logger.error(f"初始化完整模型管理器失败: {e}")
+        model_manager = SimpleModelManager()
+else:
+    model_manager = SimpleModelManager()
 
 def upload_and_analyze(image_path):
     """分析上传的图片"""
     try:
         if image_path is None:
-            return {}, {}, []
+            return {}, {}, gr.Radio(choices=[])
         
         # 打开图片
         image = Image.open(image_path).convert('RGB')
         
         # 生成图像描述
-        caption = model_manager.generate_caption(image)
+        try:
+            caption = model_manager.generate_caption(image)
+        except Exception as e:
+            logger.error(f"生成描述失败: {e}")
+            caption = "时尚服装设计"
         
         # 基于图像描述进行智能分析
         analysis_result = analyze_image_content(image, caption)
@@ -40,87 +86,94 @@ def upload_and_analyze(image_path):
         
     except Exception as e:
         error_result = {"错误": f"分析失败: {str(e)}"}
-        return error_result, {}, []
+        logger.error(f"上传分析失败: {e}")
+        return error_result, {}, gr.Radio(choices=[])
 
 def analyze_image_content(image, caption):
     """基于图像和描述进行深度分析"""
-    # 分析图像颜色
-    colors = extract_dominant_colors(image)
-    
-    # 根据描述推断风格类型
-    style_type = infer_style_from_caption(caption)
-    
-    # 根据描述推断服装类别
-    clothing_category = infer_clothing_category(caption)
-    
-    # 根据风格推荐适合场景
-    suitable_scenes = get_suitable_scenes(style_type)
-    
-    return {
-        "图像描述": caption,
-        "检测到的颜色": colors,
-        "风格类型": style_type,
-        "服装类别": clothing_category,
-        "适合场景": suitable_scenes,
-        "图像尺寸": f"{image.width} x {image.height}",
-        "分析时间": time.strftime("%Y-%m-%d %H:%M:%S")
-    }
+    try:
+        # 分析图像颜色
+        colors = extract_dominant_colors(image)
+        
+        # 根据描述推断风格类型
+        style_type = infer_style_from_caption(caption)
+        
+        # 根据描述推断服装类别
+        clothing_category = infer_clothing_category(caption)
+        
+        # 根据风格推荐适合场景
+        suitable_scenes = get_suitable_scenes(style_type)
+        
+        return {
+            "图像描述": caption,
+            "检测到的颜色": colors,
+            "风格类型": style_type,
+            "服装类别": clothing_category,
+            "适合场景": suitable_scenes,
+            "图像尺寸": f"{image.width} x {image.height}",
+            "分析时间": time.strftime("%Y-%m-%d %H:%M:%S")
+        }
+    except Exception as e:
+        logger.error(f"图像内容分析失败: {e}")
+        return {
+            "错误": f"分析失败: {str(e)}",
+            "图像描述": caption,
+            "风格类型": "休闲风"
+        }
 
 def extract_dominant_colors(image):
     """提取图像主要颜色"""
-    # 调整图像大小以提高处理速度
-    image = image.resize((150, 150))
-    
-    # 转换为numpy数组
-    img_array = np.array(image)
-    
-    # 重塑为颜色列表
-    pixels = img_array.reshape(-1, 3)
-    
-    # 使用KMeans聚类找到主要颜色
-    kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
-    kmeans.fit(pixels)
-    
-    # 将RGB值转换为颜色名称
-    color_names = []
-    for color in kmeans.cluster_centers_:
-        color_name = rgb_to_color_name(color)
-        color_names.append(color_name)
-    
-    return color_names
+    try:
+        # 调整图像大小以提高处理速度
+        image = image.resize((150, 150))
+        
+        # 转换为numpy数组
+        img_array = np.array(image)
+        
+        # 重塑为颜色列表
+        pixels = img_array.reshape(-1, 3)
+        
+        # 使用KMeans聚类找到主要颜色
+        kmeans = KMeans(n_clusters=3, random_state=42, n_init=10)
+        kmeans.fit(pixels)
+        
+        # 将RGB值转换为颜色名称
+        color_names = []
+        for color in kmeans.cluster_centers_:
+            color_name = rgb_to_color_name(color)
+            color_names.append(color_name)
+        
+        return color_names
+    except Exception as e:
+        logger.error(f"颜色提取失败: {e}")
+        return ["主色调", "辅助色", "点缀色"]
 
 def rgb_to_color_name(rgb):
     """将RGB值转换为颜色名称"""
-    r, g, b = rgb.astype(int)
-    
-    # 简单的颜色映射
-    if r > 200 and g > 200 and b > 200:
-        return "白色"
-    elif r < 50 and g < 50 and b < 50:
-        return "黑色"
-    elif r > g and r > b:
-        if r > 150:
-            return "红色"
-        else:
-            return "深红色"
-    elif g > r and g > b:
-        if g > 150:
-            return "绿色"
-        else:
-            return "深绿色"
-    elif b > r and b > g:
-        if b > 150:
-            return "蓝色"
+    try:
+        r, g, b = rgb.astype(int)
+        
+        # 简单的颜色映射
+        if r > 200 and g > 200 and b > 200:
+            return "白色"
+        elif r < 50 and g < 50 and b < 50:
+            return "黑色"
+        elif r > g and r > b:
+            return "红色系" if r > 150 else "深红色系"
+        elif g > r and g > b:
+            return "绿色系" if g > 150 else "深绿色系"
+        elif b > r and b > g:
+            return "蓝色系" if b > 150 else "深蓝色系"
+        elif r > 150 and g > 150:
+            return "黄色系"
+        elif r > 100 and b > 100:
+            return "紫色系"
+        elif g > 100 and b > 100:
+            return "青色系"
         else:
-            return "深蓝色"
-    elif r > 150 and g > 150:
-        return "黄色"
-    elif r > 100 and b > 100:
-        return "紫色"
-    elif g > 100 and b > 100:
-        return "青色"
-    else:
-        return "灰色"
+            return "灰色系"
+    except:
+        return "未知色调"
 
 def infer_style_from_caption(caption):
     """根据图像描述推断风格类型"""
@@ -177,43 +230,51 @@ def get_suitable_scenes(style_type):
 
 def generate_personalized_suggestions(analysis_result, caption):
     """基于分析结果生成个性化建议"""
-    style_type = analysis_result["风格类型"]
-    clothing_category = analysis_result["服装类别"]
-    colors = analysis_result["检测到的颜色"]
-    
-    suggestions = {}
-    
-    # 根据检测到的风格生成建议
-    if style_type == "商务正装":
-        suggestions = {
-            "经典商务": f"保持{style_type}特色，搭配{colors[0]}系配饰",
-            "现代商务": f"在{style_type}基础上加入现代元素",
-            "休闲商务": f"将{style_type}与休闲元素结合",
-            "时尚商务": f"{style_type}融入时尚潮流元素"
-        }
-    elif style_type == "休闲风":
-        suggestions = {
-            "舒适休闲": f"强化{style_type}的舒适感，主色调{colors[0]}",
-            "时尚休闲": f"{style_type}加入时尚元素",
-            "运动休闲": f"{style_type}融入运动风格",
-            "优雅休闲": f"{style_type}提升优雅感"
-        }
-    elif style_type == "运动风":
-        suggestions = {
-            "专业运动": f"增强{style_type}的功能性",
-            "休闲运动": f"{style_type}与日常穿着结合",
-            "时尚运动": f"{style_type}加入潮流设计元素",
-            "户外运动": f"强化{style_type}的户外适应性"
-        }
-    else:
-        suggestions = {
-            f"经典{style_type}": f"保持原有{style_type}特色",
-            f"现代{style_type}": f"{style_type}加入现代元素",
-            f"融合风格": f"{style_type}与其他风格混搭",
-            f"个性化{style_type}": f"基于{colors[0]}色调的个性化{style_type}"
+    try:
+        style_type = analysis_result.get("风格类型", "休闲风")
+        clothing_category = analysis_result.get("服装类别", "服装单品")
+        colors = analysis_result.get("检测到的颜色", ["主色调"])
+        
+        suggestions = {}
+        
+        # 根据检测到的风格生成建议
+        if style_type == "商务正装":
+            suggestions = {
+                "经典商务": f"保持{style_type}特色，搭配{colors[0]}系配饰",
+                "现代商务": f"在{style_type}基础上加入现代元素",
+                "休闲商务": f"将{style_type}与休闲元素结合",
+                "时尚商务": f"{style_type}融入时尚潮流元素"
+            }
+        elif style_type == "休闲风":
+            suggestions = {
+                "舒适休闲": f"强化{style_type}的舒适感，主色调{colors[0]}",
+                "时尚休闲": f"{style_type}加入时尚元素",
+                "运动休闲": f"{style_type}融入运动风格",
+                "优雅休闲": f"{style_type}提升优雅感"
+            }
+        elif style_type == "运动风":
+            suggestions = {
+                "专业运动": f"增强{style_type}的功能性",
+                "休闲运动": f"{style_type}与日常穿着结合",
+                "时尚运动": f"{style_type}加入潮流设计元素",
+                "户外运动": f"强化{style_type}的户外适应性"
+            }
+        else:
+            suggestions = {
+                f"经典{style_type}": f"保持原有{style_type}特色",
+                f"现代{style_type}": f"{style_type}加入现代元素",
+                f"融合风格": f"{style_type}与其他风格混搭",
+                f"个性化{style_type}": f"基于{colors[0]}色调的个性化{style_type}"
+            }
+        
+        return suggestions
+    except Exception as e:
+        logger.error(f"生成建议失败: {e}")
+        return {
+            "经典设计": "传统经典的设计风格",
+            "现代风格": "融入现代设计元素",
+            "个性化": "独特的个性化设计"
         }
-    
-    return suggestions
 
 def generate_designs(selected_suggestion, progress=gr.Progress()):
     """根据选择的建议生成设计"""
@@ -232,17 +293,7 @@ def generate_designs(selected_suggestion, progress=gr.Progress()):
             "舒适休闲": "casual comfort wear, soft fabrics, relaxed fit, everyday style",
             "时尚休闲": "stylish casual outfit, trendy elements, urban chic",
             "运动休闲": "athleisure wear, sporty elements, comfortable and functional",
-            "优雅休闲": "elegant casual attire, sophisticated details, refined look",
-            "专业运动": "performance sportswear, technical fabrics, functional design",
-            "休闲运动": "casual athletic wear, versatile for sports and daily use",
-            "时尚运动": "fashion sportswear, trendy athletic style, streetwear influence",
-            "户外运动": "outdoor adventure wear, durable materials, weather-resistant",
-            "经典复古风": "vintage retro style, classic silhouette, nostalgic elements",
-            "现代复古风": "contemporary take on retro fashion, updated classics",
-            "融合风格": "fusion fashion, mixed styles, innovative combination",
-            "个性化街头风": "personalized streetwear, unique designs, urban style",
-            "经典优雅风": "timeless elegant fashion, sophisticated details, refined look",
-            "现代优雅风": "modern elegant attire, contemporary sophistication"
+            "优雅休闲": "elegant casual attire, sophisticated details, refined look"
         }
         
         prompt = design_prompts.get(selected_suggestion, "fashion design, stylish clothing")
@@ -255,24 +306,27 @@ def generate_designs(selected_suggestion, progress=gr.Progress()):
             try:
                 progress(0.2 + i*0.25, desc=f"生成设计方案 {i+1}/3...")
                 
-                # 生成真实的设计图像
+                # 生成设计图像
                 image = model_manager.generate_image(
                     prompt=f"{prompt}, design {i+1}, high detail, fashion illustration",
                     negative_prompt="blurry, low quality, distorted, text, watermark",
-                    num_inference_steps=30
+                    num_inference_steps=20,  # 减少步数以提高速度
+                    width=512,
+                    height=512
                 )
                 
                 if image:
                     design_images.append(image)
                     design_choices.append(f"{selected_suggestion} 设计方案 {i+1}")
+                    
             except Exception as e:
-                print(f"生成设计 {i+1} 失败: {e}")
+                logger.error(f"生成设计 {i+1} 失败: {e}")
                 # 创建占位图像
                 width, height = 512, 512
                 img = Image.new('RGB', (width, height), 
-                               color=(random.randint(0, 255), 
-                                      random.randint(0, 255), 
-                                      random.randint(0, 255)))
+                               color=(random.randint(100, 200), 
+                                      random.randint(100, 200), 
+                                      random.randint(100, 200)))
                 design_images.append(img)
                 design_choices.append(f"{selected_suggestion} 设计方案 {i+1}")
         
@@ -280,7 +334,7 @@ def generate_designs(selected_suggestion, progress=gr.Progress()):
         return design_images, gr.Radio(choices=design_choices, value=design_choices[0] if design_choices else None)
         
     except Exception as e:
-        print(f"设计生成错误: {e}")
+        logger.error(f"设计生成错误: {e}")
         return [], gr.Radio(choices=[])
 
 def generate_3d_fitting(selected_design, progress=gr.Progress()):
@@ -294,38 +348,25 @@ def generate_3d_fitting(selected_design, progress=gr.Progress()):
         # 生成3D试穿效果的提示词
         fitting_prompt = f"3D fashion fitting, virtual try-on, {selected_design}, realistic human model, full body, studio lighting"
         
-        progress(0.3, desc="生成3D模型...")
-        
-        # 使用模型生成3D试穿图像
-        if model_manager.controlnet_pipeline:
-            # 如果有ControlNet，使用更高级的生成
-            try:
-                # 这里简化了ControlNet的使用，实际需要姿势图等输入
-                image = model_manager.controlnet_pipeline(
-                    prompt=fitting_prompt,
-                    negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs",
-                    num_inference_steps=35,
-                    guidance_scale=8.0
-                ).images[0]
-                progress(0.9, desc="渲染3D效果")
-                return image
-            except Exception as e:
-                print(f"使用ControlNet生成失败: {e}")
+        progress(0.5, desc="生成3D效果...")
         
-        # 回退到普通SD模型
-        progress(0.4, desc="使用标准模型生成...")
+        # 生成图像
         image = model_manager.generate_image(
             prompt=fitting_prompt,
             negative_prompt="blurry, distorted, low quality, unrealistic, extra limbs",
-            num_inference_steps=35
+            num_inference_steps=25,
+            width=512,
+            height=768  # 适合全身图像
         )
         
         progress(0.9, desc="完成3D渲染")
         return image
         
     except Exception as e:
-        print(f"3D试穿生成错误: {e}")
-        return None
+        logger.error(f"3D试穿生成错误: {e}")
+        # 创建占位图像
+        img = Image.new('RGB', (512, 768), color=(200, 200, 200))
+        return img
 
 def create_gradio_interface():
     """创建Gradio用户界面"""
@@ -355,21 +396,10 @@ def create_gradio_interface():
         with gr.Tab("3D试穿效果"):
             fitting_result = gr.Image(label="3D试穿效果", height=500)
         
-        # 添加示例图片
-        examples_dir = "examples"
-        if os.path.exists(examples_dir):
-            example_files = [os.path.join(examples_dir, f) for f in os.listdir(examples_dir) 
-                            if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
+        # 内存管理
+        with gr.Row():
+            cleanup_btn = gr.Button("清理内存", variant="secondary")
             
-            gr.Examples(
-                examples=example_files[:4],  # 最多显示4个示例
-                inputs=image_input,
-                outputs=[analysis_output, suggestions_output, suggestion_choice],
-                fn=upload_and_analyze,
-                cache_examples=True,
-                label="示例图片"
-            )
-        
         # 事件绑定
         analyze_btn.click(
             fn=upload_and_analyze,
@@ -389,69 +419,30 @@ def create_gradio_interface():
             outputs=[fitting_result]
         )
         
-        # 添加清理按钮
-        clear_btn = gr.Button("清理内存", variant="secondary")
-        clear_btn.click(
+        cleanup_btn.click(
             fn=model_manager.cleanup,
             inputs=[],
             outputs=[]
         )
-        gr.Markdown("> **提示**: 生成图像后，点击'清理内存'按钮可以释放GPU资源")
+        
+        gr.Markdown("> **提示**: 如果遇到内存问题，请点击'清理内存'按钮")
     
     return demo
 
 if __name__ == "__main__":
-    # 检查并创建示例目录
-    examples_dir = "examples"
-    if not os.path.exists(examples_dir):
-        os.makedirs(examples_dir)
-        print(f"创建了示例目录: {examples_dir}")
-        print("请在此目录中添加示例图片以便在界面中使用")
+    # 检查环境
+    logger.info(f"Python版本: {sys.version}")
+    logger.info(f"当前工作目录: {os.getcwd()}")
     
+    # 创建并启动界面
     demo = create_gradio_interface()
-    demo.queue(concurrency_count=1)  # 限制并发以避免内存问题
+    demo.queue(
+        concurrency_count=1,  # 限制并发
+        max_size=10  # 队列大小
+    )
     demo.launch(
         server_name="0.0.0.0", 
-        server_port=7860, 
-        share=True,
-        favicon_path="favicon.ico" if os.path.exists("favicon.ico") else None
-    )
-    # 在 create_gradio_interface() 函数中添加内存管理按钮
-def create_gradio_interface():
-    with gr.Blocks(title="AI时尚设计师", theme="soft") as demo:
-        # ... [现有代码] ...
-        
-        # 添加内存管理部分
-        with gr.Row():
-            gr.Markdown("### 内存管理")
-            unload_caption_btn = gr.Button("卸载描述模型")
-            unload_sd_btn = gr.Button("卸载设计模型")
-            unload_controlnet_btn = gr.Button("卸载3D模型")
-            full_cleanup_btn = gr.Button("清理所有模型", variant="stop")
-        
-        # 内存管理事件
-        unload_caption_btn.click(
-            fn=lambda: model_manager.unload_model("caption"),
-            inputs=[],
-            outputs=[]
-        )
-        
-        unload_sd_btn.click(
-            fn=lambda: model_manager.unload_model("sd"),
-            inputs=[],
-            outputs=[]
-        )
-        
-        unload_controlnet_btn.click(
-            fn=lambda: model_manager.unload_model("controlnet"),
-            inputs=[],
-            outputs=[]
-        )
-        
-        full_cleanup_btn.click(
-            fn=model_manager.cleanup,
-            inputs=[],
-            outputs=[]
-        )
-        
-    return demo
+        server_port=7860,
+        share=False,  # 在Spaces环境中设为False
+        show_error=True
+    )