Update app.py

app.py

@@ -1,26 +1,32 @@
 import os
 import torch
-import gradio as gr
 import spaces
+import gradio as gr
 from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# 1. 加载模型 Qwen2.5
+# 1. 明确模型名称
 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
 
 print(f"Loading model {MODEL_NAME}...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+# 【关键点】ZeroGPU 启动时全局只能加载到 CPU / float16，不能直接分配 GPU
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
+    torch_dtype=torch.float16,
+    device_map="cpu"
 )
-print("Model loaded successfully!")
+print("Model loaded successfully on CPU!")
 
-# 2. 绑定 ZeroGPU 核心推理函数
+# 2. ZeroGPU 核心推理函数（必须处于全局最外层）
 @spaces.GPU
 def generate_response(system_prompt: str, user_prompt: str, temperature: float = 0.7, max_tokens: int = 2048):
     try:
+        # 运行时把模型移入 ZeroGPU 动态分配的 GPU 中
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        
         messages = [
             {"role": "system", "content": system_prompt},
             {"role": "user", "content": user_prompt}
@@ -32,7 +38,7 @@ def generate_response(system_prompt: str, user_prompt: str, temperature: float =
             add_generation_prompt=True
         )
 
-        inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device)
+        inputs = tokenizer([prompt_text], return_tensors="pt").to(device)
         
         with torch.no_grad():
             outputs = model.generate(
@@ -50,21 +56,21 @@ def generate_response(system_prompt: str, user_prompt: str, temperature: float =
     except Exception as e:
         return f"Error: {str(e)}"
 
-# 3. 创建 FastAPI 实例
+# 3. 创建 FastAPI 实例（供 n8n 直连）
 app = FastAPI()
 
-# 4. 专供 n8n 调用的直连 POST 接口（不走 Gradio 队列）
 @app.post("/api/generate")
-async def api_generate(payload: dict):
-    system_prompt = payload.get("system_prompt", "你是一个专业的小红书美食博主，只输出 JSON 数据。")
+def api_generate(payload: dict):
+    system_prompt = payload.get("system_prompt", "你是专业的小红书美食博主。")
     user_prompt = payload.get("user_prompt", "")
     temperature = float(payload.get("temperature", 0.7))
     max_tokens = int(payload.get("max_tokens", 2048))
     
+    # 直接调用全局 @spaces.GPU 装饰的函数
     result = generate_response(system_prompt, user_prompt, temperature, max_tokens)
     return {"result": result}
 
-# 5. 构建 Gradio 可视化界面
+# 4. 构建 Gradio 可视化 UI
 with gr.Blocks() as demo:
     gr.Markdown("# 通用 LLM API Space (ZeroGPU)")
     sys_in = gr.Textbox(label="System Prompt", value="你是一个专业的小红书美食博主，只输出 JSON 数据。")
@@ -75,5 +81,5 @@ with gr.Blocks() as demo:
     btn = gr.Button("Submit")
     btn.click(fn=generate_response, inputs=[sys_in, user_in, temp_in, tokens_in], outputs=out)
 
-# 6. 将 Gradio 挂载到 FastAPI 根路径
+# 挂载 Gradio 页面
 app = gr.mount_gradio_app(app, demo, path="/")