Update app.py

app.py

@@ -1,60 +1,86 @@
 # app.py
-
 import gradio as gr
-from transformers import AutoModel
 import torch
 import numpy as np
 from PIL import Image
+from transformers import AutoModel
+import warnings
+warnings.filterwarnings('ignore')
 
-def load_model():
-    # 加载模型
-    model = AutoModel.from_pretrained("jadechoghari/vfusion3d", trust_remote_code=True)
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
+# 全局变量存储模型实例
+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(input_image)
-        
-        # 加载模型
-        model = load_model()
+            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()
+        # 转换为tensor并归一化
+        image_tensor = torch.from_numpy(np.array(input_image)).float() / 255.0
         image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0)
         
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        image_tensor = image_tensor.to(device)
-        
         # 模型推理
         with torch.no_grad():
-            output = model(image_tensor)
-            
-        return output
-        
+            try:
+                output = model(image_tensor)
+                return output, "处理成功"
+            except Exception as e:
+                return None, f"模型推理错误: {str(e)}"
+                
     except Exception as e:
-        return f"错误: {str(e)}"
+        return None, f"处理错误: {str(e)}"
 
 # 创建Gradio界面
 demo = gr.Interface(
     fn=process_image,
     inputs=[
-        gr.Image(type="pil", label="上传图片")
+        gr.Image(
+            type="pil",
+            label="上传图片",
+            tool="select"
+        )
     ],
     outputs=[
-        gr.Model3D(label="生成的3D模型"),
-        gr.Text(label="处理状态")
+        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()
+    theme=gr.themes.Soft(),
+    allow_flagging="never"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+# 启动应用
+demo.launch()