Update app.py

app.py

@@ -1,58 +1,35 @@
 import gradio as gr
 from transformers import pipeline
-from datasets import load_dataset
 
-# 加载数据集（使用更小的样本集）
+# 初始化文本生成管道
 try:
-    dataset = load_dataset("AntZet/home_decoration_objects_images", streaming=True)
-    dataset = dataset['train'].take(50)  # 只取50个样本加快加载
+    generator = pipeline("text-generation", model="distilgpt2")
 except Exception as e:
-    print(f"加载数据集失败: {e}")
-    dataset = []
+    print(f"模型加载失败: {e}")
+    generator = None
 
-# 初始化模型（使用更小的模型）
-try:
-    style_advisor = pipeline("text-generation", model="distilgpt2")  # 改用更小的distilgpt2
-except Exception as e:
-    print(f"加载模型失败: {e}")
-    style_advisor = None
-
-def get_advice(style):
-    if not dataset or not style_advisor:
-        return "系统初始化失败，请检查后台日志", []
+def generate_advice(style):
+    if not generator:
+        return "系统初始化失败，请检查transformers库是否安装", ""
     
+    prompt = f"如何将Airbnb房间装修成{style}风格？请给出3条具体建议"
     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
+        result = generator(prompt, max_length=200)
+        return result[0]["generated_text"], ""
     except Exception as e:
-        return f"生成建议时出错: {str(e)}", []
+        return f"生成建议时出错: {str(e)}", ""
 
-# 创建界面
+# 创建简单界面
 with gr.Blocks() as demo:
     gr.Markdown("## 🏡 Airbnb装修助手")
     style_input = gr.Textbox(label="输入想要的风格")
     submit_btn = gr.Button("生成建议")
-    
-    with gr.Row():
-        advice_output = gr.Textbox(label="装修建议")
-        gallery = gr.Gallery(label="风格示例")
+    advice_output = gr.Textbox(label="装修建议")
     
     submit_btn.click(
-        fn=get_advice,
+        fn=generate_advice,
         inputs=style_input,
-        outputs=[advice_output, gallery]
+        outputs=advice_output
     )
 
 if __name__ == "__main__":