Rename app.py to model.py

app.py → model.py

@@ -1,31 +1,25 @@
-import gradio as gr
-import torch
-from transformers import pipeline
 from diffusers import StableDiffusionPipeline
+from transformers import pipeline
+import torch
 
 # 加载Stable Diffusion模型
-stable_diff_pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
+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
 
 # 加载GPT模型
-gpt_model = pipeline("text-generation", model="gpt2")
-
-# 生成提示
-def generate_prompt(description):
-    prompt = gpt_model(f"Create a detailed image prompt based on this: {description}")
-    return prompt[0]['generated_text']
+def load_gpt_model():
+    gpt_model = pipeline("text-generation", model="gpt2")
+    return gpt_model
 
-# 生成图像
-def generate_image(prompt):
-    image = stable_diff_pipe(prompt).images[0]
+# 生成图像的函数
+def generate_image_from_prompt(model, prompt):
+    image = model(prompt).images[0]
     return image
 
-# 创建Gradio界面
-def interface(description):
-    prompt = generate_prompt(description)
-    image = generate_image(prompt)
-    return prompt, image
-
-# Gradio UI组件
-iface = gr.Interface(fn=interface, inputs="text", outputs=["text", "image"])
+# 使用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']
 
-iface.launch()