Update app.py

app.py

@@ -1,27 +1,36 @@
 import gradio as gr
-import os
-import requests
+from fastapi import FastAPI, Request
+import uvicorn
+from transformers import AutoTokenizer, AutoModel
+import torch
 
-def predict(message, history):
-    # 这里保持你之前的调用逻辑
-    api_url = "https://router.huggingface.co/jiang1002/chatglm-6b-adgen"
-    headers = {"Authorization": f"Bearer {os.environ.get('HF_TOKEN')}"}
-    payload = {"inputs": message, "options": {"wait_for_model": True}}
-    
-    try:
-        response = requests.post(api_url, headers=headers, json=payload, timeout=60)
-        if response.status_code == 200:
-            result = response.json()
-            # 兼容处理返回结果
-            return result[0].get("generated_text") if isinstance(result, list) else str(result)
-        else:
-            return f"Error: {response.status_code}"
-    except Exception as e:
-        return f"发生错误: {str(e)}"
+# 初始化 FastAPI
+app = FastAPI()
 
-# 创建界面
-demo = gr.ChatInterface(fn=predict, title="ChatGLM 广告生成助手")
+# 加载模型 (这是为了让你的 Space 真正拥有处理能力)
+model_path = "jiang1002/chatglm-6b-adgen"
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_path, trust_remote_code=True).float()
+model.eval()
+
+# --- 核心：给别人用的 API 接口 ---
+@app.post("/generate")
+async def generate(request: Request):
+    data = await request.json()
+    prompt = data.get("text", "")
+    # 调用本地加载好的模型生成
+    response, _ = model.chat(tokenizer, prompt, history=[])
+    return {"result": response}
+
+# --- 同时也保留网页界面方便你自己测 ---
+def chat_func(msg, hist):
+    res, _ = model.chat(tokenizer, msg, history=hist)
+    return res
+
+demo = gr.ChatInterface(fn=chat_func, title="ChatGLM API Server")
+
+# 把 Gradio 挂载到 FastAPI 上
+app = gr.mount_gradio_app(app, demo, path="/")
 
 if __name__ == "__main__":
-    # ✅ 必须设置 server_name 和 server_port，否则 Docker 内部访问不到
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)