# app.py
import gradio as gr
import torch
import numpy as np
from PIL import Image
from transformers import AutoModel
import warnings
warnings.filterwarnings('ignore')

# 全局变量存储模型实例
model = None

def initialize_model():
    global model
    try:
        if model is None:
            model = AutoModel.from_pretrained(
                "jadechoghari/vfusion3d", 
                trust_remote_code=True,
                device_map="auto"  # 自动处理设备分配
            )
    except Exception as e:
        print(f"模型加载错误: {str(e)}")
        return None
    return model

def process_image(input_image):
    if input_image is None:
        return None, "请上传图片"
    
    try:
        # 初始化模型
        model = initialize_model()
        if model is None:
            return None, "模型加载失败"

        # 确保输入图像是PIL Image格式
        if not isinstance(input_image, Image.Image):
            input_image = Image.fromarray(np.uint8(input_image))
        
        # 图像预处理
        input_image = input_image.resize((256, 256))
        
        # 转换为tensor并归一化
        image_tensor = torch.from_numpy(np.array(input_image)).float() / 255.0
        image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0)
        
        # 模型推理
        with torch.no_grad():
            try:
                output = model(image_tensor)
                return output, "处理成功"
            except Exception as e:
                return None, f"模型推理错误: {str(e)}"
                
    except Exception as e:
        return None, f"处理错误: {str(e)}"

# 创建Gradio界面
demo = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(
            type="pil",
            label="上传图片",
            tool="select"
        )
    ],
    outputs=[
        gr.Model3D(
            label="生成的3D模型",
            clear_color=[0.0, 0.0, 0.0, 0.0]
        ),
        gr.Textbox(
            label="处理状态",
            placeholder="等待处理..."
        )
    ],
    title="麒迹云台 - 2D转3D模型生成器",
    description="上传一张图片，AI将自动生成对应的3D模型。支持格式：jpg, png, jpeg",
    theme=gr.themes.Soft(),
    allow_flagging="never"
)

# 启动应用
demo.launch()