Update model.py

model.py

@@ -1,25 +1,31 @@
-from diffusers import StableDiffusionPipeline
-from transformers import pipeline
+import openai
 import torch
+from diffusers import StableDiffusionPipeline
+import os
+from dotenv import load_dotenv
 
-# 加载Stable Diffusion模型
-def load_stable_diffusion_model():
-    model = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
-    model.to("cuda" if torch.cuda.is_available() else "cpu")  # 如果有GPU，使用GPU加速
-    return model
+# 加载环境变量
+load_dotenv()
 
-# 加载GPT模型
-def load_gpt_model():
-    gpt_model = pipeline("text-generation", model="gpt2")
-    return gpt_model
+# 获取 API 密钥
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+openai.api_key = OPENAI_API_KEY
 
-# 生成图像的函数
-def generate_image_from_prompt(model, prompt):
-    image = model(prompt).images[0]
-    return image
+# 加载 Stable Diffusion 模型
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v-1-4-original", torch_dtype=torch.float16).to("cuda")
 
-# 使用GPT生成图像的详细提示
-def generate_prompt(gpt_model, description):
-    prompt = gpt_model(f"Create a detailed image prompt based on this: {description}")
-    return prompt[0]['generated_text']
+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