Update app.py

app.py

@@ -2,28 +2,42 @@ import gradio as gr
 from transformers import pipeline
 from datasets import load_dataset
 
-# 加载数据集
-dataset = load_dataset("AntZet/home_decoration_objects_images", streaming=True)
-dataset = dataset['train'].take(100)  # 只取前100个样本加快加载
+# 加载数据集（使用更小的样本集）
+try:
+    dataset = load_dataset("AntZet/home_decoration_objects_images", streaming=True)
+    dataset = dataset['train'].take(50)  # 只取50个样本加快加载
+except Exception as e:
+    print(f"加载数据集失败: {e}")
+    dataset = []
 
-# 初始化模型
-style_advisor = pipeline("text-generation", model="gpt2")  # 使用较小的gpt2模型
+# 初始化模型（使用更小的模型）
+try:
+    style_advisor = pipeline("text-generation", model="distilgpt2")  # 改用更小的distilgpt2
+except Exception as e:
+    print(f"加载模型失败: {e}")
+    style_advisor = None
 
 def get_advice(style):
-    # 从数据集中找出匹配风格的例子
-    examples = [ex for ex in dataset if ex['style'].lower() == style.lower()]
+    if not dataset or not style_advisor:
+        return "系统初始化失败，请检查后台日志", []
     
-    if not examples:
-        return "未找到该风格，请尝试：工业风、北欧风等", []
-    
-    # 生成建议
-    prompt = f"如何将Airbnb房间装修成{style}风格？请给出3条具体建议"
-    advice = style_advisor(prompt, max_length=150)[0]['generated_text']
-    
-    # 获取示例图片
-    example_images = [ex['image'] for ex in examples[:3]]
-    
-    return advice, example_images
+    try:
+        # 从数据集中找出匹配风格的例子
+        examples = [ex for ex in dataset if ex['style'].lower() == style.lower()]
+        
+        if not examples:
+            return "未找到该风格，请尝试：工业风、北欧风等", []
+        
+        # 生成建议
+        prompt = f"如何将Airbnb房间装修成{style}风格？请给出3条具体建议"
+        advice = style_advisor(prompt, max_length=150)[0]['generated_text']
+        
+        # 获取示例图片
+        example_images = [ex['image'] for ex in examples[:3]]
+        
+        return advice, example_images
+    except Exception as e:
+        return f"生成建议时出错: {str(e)}", []
 
 # 创建界面
 with gr.Blocks() as demo:
@@ -41,4 +55,5 @@ with gr.Blocks() as demo:
         outputs=[advice_output, gallery]
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()