import openai
import torch
from diffusers import StableDiffusionPipeline
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 获取 API 密钥
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
openai.api_key = OPENAI_API_KEY

# 加载 Stable Diffusion 模型
pipe = StableDiffusionPipeline.from_pretrained("SfinOe/stable-diffusion-v1.5", torch_dtype=torch.float16).to("cpu")

def generate_prompt(description):
    """ 使用 OpenAI 生成图像生成的提示 """
    prompt = f"Generate a detailed prompt for stable diffusion image generation based on the description: {description}"
    response = openai.Completion.create(
        engine="text-davinci-003",  # 或其他 GPT-3 模型
        prompt=prompt,
        max_tokens=100
    )
    return response.choices[0].text.strip()

def generate_image_from_prompt(prompt):
    """ 使用 Stable Diffusion 生成图像 """
    image = pipe(prompt).images[0]
    return image