app.py

import os
import sys
import warnings
warnings.filterwarnings("ignore")

try:
    import gradio as gr
    from PIL import Image
    import torch
    print(f"✅ 基础库加载成功，PyTorch: {torch.__version__}")
except ImportError as e:
    print(f"❌ 基础库加载失败: {e}")
    sys.exit(1)

try:
    from transformers import AutoTokenizer, AutoModel, AutoProcessor, Blip2ForConditionalGeneration
    print("✅ Transformers 加载成功")
except ImportError as e:
    print(f"❌ Transformers 加载失败: {e}")
    # 使用简化版本
    USE_AI_MODELS = False
else:
    USE_AI_MODELS = True

# 环境配置
IS_SPACES = os.environ.get("SPACE_ID") is not None
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"环境: {'HF Spaces' if IS_SPACES else 'Local'}, 设备: {device}, AI模型: {USE_AI_MODELS}")

# 全局变量
processor = None
blip_model = None
tokenizer = None
chat_model = None

def load_image_model():
    """加载图像理解模型"""
    global processor, blip_model
    
    if not USE_AI_MODELS:
        return False
        
    try:
        print("📷 加载图像理解模型...")
        
        # 使用更小更稳定的BLIP模型
        model_name = "Salesforce/blip2-opt-2.7b"
        
        processor = AutoProcessor.from_pretrained(model_name)
        blip_model = Blip2ForConditionalGeneration.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            device_map="auto" if device == "cuda" else None,
            load_in_8bit=device == "cuda"
        )
        
        if device == "cpu":
            blip_model = blip_model.to("cpu")
            
        print("✅ 图像模型加载成功")
        return True
        
    except Exception as e:
        print(f"❌ 图像模型加载失败: {str(e)}")
        return False

def load_chat_model():
    """加载对话模型"""
    global tokenizer, chat_model
    
    if not USE_AI_MODELS:
        return False
        
    try:
        print("💬 加载对话模型...")
        
        model_name = "THUDM/chatglm2-6b-int4"
        tokenizer = AutoTokenizer.from_pretrained(
            model_name,
            trust_remote_code=True,
            use_fast=False
        )
        
        chat_model = AutoModel.from_pretrained(
            model_name,
            trust_remote_code=True,
            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
            low_cpu_mem_usage=True
        )
        
        if device == "cuda":
            chat_model = chat_model.half().cuda()
        chat_model.eval()
        
        print("✅ 对话模型加载成功")
        return True
        
    except Exception as e:
        print(f"❌ 对话模型加载失败: {str(e)}")
        return False

def describe_image_ai(image):
    """使用AI模型描述图像"""
    if not processor or not blip_model:
        return None
        
    try:
        # 预处理图像
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
            
        # 调整图像大小，避免内存问题
        max_size = 384
        if max(image.size) > max_size:
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        # 确保图像是RGB模式
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        inputs = processor(image, return_tensors="pt")
        
        if device == "cuda":
            inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            generated_ids = blip_model.generate(
                **inputs,
                max_new_tokens=20,  # 减少token数量
                min_length=5,       # 最小长度
                num_beams=2,        # 减少beam数量
                do_sample=False,    # 关闭采样
                early_stopping=True,
                pad_token_id=processor.tokenizer.pad_token_id
            )
        
        caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
        
        # 验证描述质量
        if len(caption) < 3 or caption.lower() in ['a', 'an', 'the', '']:
            print(f"⚠️ AI描述质量差: '{caption}'")
            return None
            
        print(f"✅ AI图像描述: {caption}")
        return caption
        
    except Exception as e:
        print(f"❌ AI图像描述失败: {str(e)}")
        return None

def describe_image_basic(image):
    """基础图像分析（不依赖AI）"""
    try:
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)
        
        width, height = image.size
        mode = image.mode
        aspect_ratio = width / height
        
        # 基本尺寸分析
        if aspect_ratio > 1.5:
            orientation = "横向构图"
        elif aspect_ratio < 0.7:
            orientation = "纵向构图" 
        else:
            orientation = "方形构图"
        
        # 颜色分析
        try:
            colors = image.getcolors(maxcolors=256*256*256)
            if colors:
                # 分析主要颜色
                sorted_colors = sorted(colors, key=lambda x: x[0], reverse=True)
                dominant_color = sorted_colors[0][1]
                
                # 简单色彩描述
                r, g, b = dominant_color[:3] if len(dominant_color) >= 3 else (128, 128, 128)
                if r > g and r > b:
                    color_tone = "偏红色调"
                elif g > r and g > b:
                    color_tone = "偏绿色调"
                elif b > r and b > g:
                    color_tone = "偏蓝色调"
                else:
                    color_tone = "中性色调"
            else:
                color_tone = "色彩丰富"
        except:
            color_tone = "色彩复杂"
        
        description = f"一幅{orientation}的艺术作品，{color_tone}，尺寸为{width}x{height}像素"
        return description
        
    except Exception as e:
        print(f"基础图像分析失败: {str(e)}")
        return "一幅艺术作品"

def analyze_artwork(image):
    """综合艺术品分析"""
    if image is None:
        return "请上传艺术品图像"
    
    print("🎨 开始分析艺术品...")
    
    # 尝试AI描述，失败则使用基础描述
    ai_description = describe_image_ai(image) if USE_AI_MODELS else None
    
    if ai_description:
        description = ai_description
        analysis_type = "AI智能分析"
    else:
        description = describe_image_basic(image)
        analysis_type = "基础视觉分析"
    
    # 生成分析报告
    analysis = f"""
🎨 艺术品分析报告 ({analysis_type})

📷 图像描述: {description}

🎯 艺术解读:
从视觉构成来看，这幅作品体现了艺术家对空间、色彩和形式的独特理解。作品通过画面元素的组织和安排，创造出特定的视觉效果和情感氛围。

🖌️ 技法特点:
• 构图安排体现了艺术家的空间意识
• 色彩运用展现了对色彩关系的掌握
• 整体风格反映了特定的艺术追求

💭 欣赏建议:
建议从构图平衡、色彩和谐、主题表达等角度深入欣赏这件作品。每一次观看都可能发现新的细节和感受。

---
💡 提示: 您可以在下方继续提问，探讨更多艺术话题
"""
    
    return analysis

def chat_with_ai(message, history):
    """智能对话功能"""
    if not message or not message.strip():
        return history, ""
    
    history = history or []
    
    # 尝试使用ChatGLM
    if chat_model and tokenizer:
        try:
            response, _ = chat_model.chat(tokenizer, message, history=[])
            history.append([message, response])
            return history, ""
        except Exception as e:
            print(f"AI对话失败: {str(e)}")
            # 继续使用预设回复
    
    # 预设的艺术话题回复
    responses = {
        "颜色": "色彩是艺术作品中最直观的视觉元素。不同颜色能够唤起不同的情感反应，艺术家通过色彩的冷暖对比、明暗变化来营造画面氛围。",
        "构图": "构图是艺术作品的骨架结构。经典的构图法则包括三分法、黄金分割、对称与均衡等，这些都有助于创造视觉焦点和引导观者视线。",
        "风格": "艺术风格是艺术家个性和时代特征的体现。从古典主义的严谨到印象派的光影变化，每种风格都有其独特的表现语言。",
        "技法": "绘画技法包括笔触运用、色彩调和、明暗处理等。掌握技法是艺术表达的基础，但真正的艺术超越技法，直达情感和思想。",
        "历史": "艺术史是人类文明发展的重要组成部分。了解作品的历史背景有助于更好地理解艺术家的创作意图和作品的文化价值。"
    }
    
    # 简单的关键词匹配
    response = "这是个很有趣的艺术问题。"
    for keyword, reply in responses.items():
        if keyword in message:
            response = reply
            break
    else:
        # 默认回复
        if "?" in message or "？" in message:
            response = "关于这个问题，建议从作品的视觉元素、创作背景和个人感受三个维度来思考。艺术欣赏很大程度上是主观的，重要的是培养自己的审美感受力。"
        else:
            response = "您提到的观点很有见地。艺术作品往往承载着丰富的内涵，值得我们从多个角度去解读和欣赏。"
    
    history.append([message, response])
    return history, ""

# 创建界面
def create_interface():
    with gr.Blocks(
        title="AI艺术品讲解智能体", 
        theme=gr.themes.Soft()
    ) as demo:
        
        gr.HTML("""
        <div style="text-align: center; margin-bottom: 20px;">
            <h1>🎨 AI 艺术品讲解智能体</h1>
            <p>上传艺术品图像，获得专业分析与互动讨论</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                image_input = gr.Image(
                    label="📤 上传艺术品图像",
                    type="pil",
                    height=350
                )
                
                analyze_btn = gr.Button("🔍 开始分析", variant="primary", size="lg")
                
                gr.HTML("""
                <div style="margin-top: 15px; padding: 10px; background: #f8f9fa; border-radius: 8px; font-size: 0.9em;">
                    <b>💡 使用提示:</b><br>
                    • 支持JPG、PNG等常见格式<br>
                    • 建议上传清晰的艺术作品图像<br>
                    • 分析完成后可继续提问交流
                </div>
                """)
                
            with gr.Column(scale=2):
                analysis_output = gr.Textbox(
                    label="📊 艺术品分析",
                    lines=15,
                    max_lines=20,
                    placeholder="上传图像后点击'开始分析'按钮..."
                )
        
        # 对话区域
        gr.HTML("<hr style='margin: 30px 0;'><h3>💬 继续探讨</h3>")
        
        chatbot = gr.Chatbot(
            label="与AI艺术顾问对话",
            height=300,
            placeholder="分析完成后，您可以在这里继续讨论艺术相关话题"
        )
        
        chat_input = gr.Textbox(
            label="输入您的问题或观点",
            placeholder="例如：这幅作品的色彩运用有什么特点？",
            lines=2
        )
        
        with gr.Row():
            clear_chat_btn = gr.Button("🗑️ 清空对话", variant="secondary")
            clear_all_btn = gr.Button("🔄 重新开始", variant="secondary")
        
        # 事件绑定
        analyze_btn.click(
            fn=analyze_artwork,
            inputs=image_input,
            outputs=analysis_output
        )
        
        chat_input.submit(
            fn=chat_with_ai,
            inputs=[chat_input, chatbot],
            outputs=[chatbot, chat_input]
        )
        
        clear_chat_btn.click(
            fn=lambda: [],
            outputs=chatbot
        )
        
        def reset_all():
            return None, "", []
        
        clear_all_btn.click(
            fn=reset_all,
            outputs=[image_input, analysis_output, chatbot]
        )
        
        # 功能说明
        gr.HTML(f"""
        <div style="margin-top: 30px; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
                    color: white; border-radius: 10px;">
            <h4 style="margin-top: 0;">🔧 系统状态</h4>
            <p><strong>AI功能:</strong> {'✅ 已启用' if USE_AI_MODELS else '❌ 简化模式'}</p>
            <p><strong>运行环境:</strong> {'Hugging Face Spaces' if IS_SPACES else '本地环境'}</p>
            <p><strong>计算设备:</strong> {device.upper()}</p>
            
            <h4>📚 支持的讨论话题</h4>
            <p>色彩分析 • 构图技巧 • 艺术风格 • 绘画技法 • 艺术史背景</p>
        </div>
        """)
    
    return demo

# 主程序
if __name__ == "__main__":
    print("🚀 启动艺术品分析应用...")
    
    # 预加载模型（可选）
    if USE_AI_MODELS:
        print("📦 预加载AI模型...")
        image_model_loaded = load_image_model()
        chat_model_loaded = load_chat_model()
        print(f"模型状态: 图像模型{'✅' if image_model_loaded else '❌'}, 对话模型{'✅' if chat_model_loaded else '❌'}")
    
    # 创建并启动界面
    demo = create_interface()
    
    try:
        demo.launch(
            share=False,
            show_error=True,
            server_name="0.0.0.0" if IS_SPACES else None,
            server_port=7860 if IS_SPACES else None
        )
    except Exception as e:
        print(f"❌ 启动失败: {str(e)}")
        # 备用简化启动
        with gr.Blocks() as backup:
            gr.HTML(f"<h2>启动遇到问题</h2><p>错误: {str(e)}</p><p>请刷新页面重试</p>")
        backup.launch()