Update generate_image.py

generate_image.py

@@ -1,18 +1,24 @@
+# generate_image.py
+
 import os
-import requests
+import torch
+from diffusers import StableDiffusionPipeline
 from PIL import Image
-from io import BytesIO
 
-# 替换为你选定的模型 repo（如 stablediffusionapi/anything-v5）
-API_URL = "https://api-inference.huggingface.co/models/nitrosocke/epiCRealism"
-headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
+# 加载 Hugging Face Token（你需要在 Spaces 的 Secrets 设置中添加 HF_TOKEN）
+hf_token = os.getenv("HF_TOKEN")
 
-def generate_ghibli_avatar(prompt, reference_image=None):
-    # 模拟 ControlNet 的参考图传入能力（实际不支持，参考图可用于引导微调）
-    payload = {"inputs": prompt}
-    response = requests.post(API_URL, headers=headers, json=payload)
+# 使用 nitrosocke/Ghibli-Diffusion 模型
+pipe = StableDiffusionPipeline.from_pretrained(
+    "nitrosocke/Ghibli-Diffusion",
+    use_auth_token=hf_token
+).to("cuda" if torch.cuda.is_available() else "cpu")
 
-    if response.status_code != 200:
-        raise Exception(f"生成失败：{response.status_code}, {response.text}")
 
-    return Image.open(BytesIO(response.content))
+def generate_ghibli_avatar(prompt, reference_image=None):
+    """
+    生成吉卜力风格头像，仅使用文本 prompt，不依赖 reference image。
+    """
+    final_prompt = f"{prompt}, ghibli style"
+    image = pipe(final_prompt, num_inference_steps=30).images[0]
+    return image