Update app.py

app.py

@@ -1,54 +1,26 @@
 import gradio as gr
-from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
-import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration
 from diffusers import StableDiffusionPipeline
-import cv2
-import numpy as np
+import torch
 
-# 使用T5模型生成文本描述
-model_name = "t5-large"  # 你可以根据需求选择不同版本的T5
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-t5_model = T5ForConditionalGeneration.from_pretrained(model_name)
+# 使用BART模型生成文本描述
+prompt_generator = pipeline("text2text-generation", model="facebook/bart-large-cnn")
 
 def generate_prompt(description: str) -> str:
-    # 使用T5模型生成详细的图像生成提示
-    input_text = f"将这个描述扩展为一个详细的图像生成提示：{description}"
-    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
-    outputs = t5_model.generate(inputs["input_ids"], max_length=150, num_beams=5, early_stopping=True)
-    prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # 根据简短描述生成详细的图像生成提示
+    prompt = prompt_generator(f"将这个描述扩展为一个详细的图像生成提示：{description}", max_length=150)[0]['generated_text']
     return prompt
 
-# 加载 ControlNet Canny 边缘检测模型
-controlnet = ControlNetModel.from_pretrained("lllyasviel/controlnet-canny-sdxl-1.0")
-stable_diffusion = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-1.0")
+# 加载 ByteDance/SDXL-Lightning 模型
+sdxl_pipeline = StableDiffusionPipeline.from_pretrained("ByteDance/SDXL-Lightning")
+sdxl_pipeline.to("cpu")  # 使用 CPU
 
-# 使用 ControlNet 管道
-pipe = StableDiffusionControlNetPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-1.0", 
-    controlnet=controlnet
-)
-pipe.to("cpu")  # 使用CPU
-
-def generate_image_with_controlnet(prompt: str):
-    # 生成 Canny 边缘图像并传入 ControlNet
-    # 使用模型生成图像并提取边缘
-    image = pipe(prompt).images[0]
-    
-    # 转换为灰度图像
-    image = np.array(image.convert('L'))  # 转为灰度图
-    
-    # 使用 Canny 边缘检测
-    canny_edge_image = cv2.Canny(image, 100, 200)  # 进行 Canny 边缘检测
-    
-    # 将 Canny 边缘图像转换为适用于 ControlNet 的格式
-    canny_edge_image = torch.from_numpy(canny_edge_image).unsqueeze(0).unsqueeze(0).float() / 255.0  # 规范化
-    generated_image = pipe(prompt=prompt, control_image=canny_edge_image).images[0]
-    
-    return generated_image
+def generate_image(prompt: str):
+    # 根据提示生成图像
+    image = sdxl_pipeline(prompt).images[0]
+    return image
 
 # 使用Whisper模型进行语音转文本
-from transformers import WhisperProcessor, WhisperForConditionalGeneration
 processor = WhisperProcessor.from_pretrained("openai/whisper-large")
 model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-large")
 
@@ -64,21 +36,20 @@ def process_input(description: str, creativity: float, include_background: bool)
     prompt = generate_prompt(description)
     if include_background:
         prompt += " 添加详细的生动背景。"
-    image = generate_image_with_controlnet(prompt)
+    image = generate_image(prompt)
     return prompt, image
 
 # 处理音频输入和生成图像
 def process_audio_input(audio):
     description = transcribe_audio(audio)
     prompt = generate_prompt(description)
-    image = generate_image_with_controlnet(prompt)
+    image = generate_image(prompt)
     return prompt, image
 
 # Gradio界面部分
 text_input = gr.Textbox(label="请输入描述", placeholder="例如：天空中的魔法树屋")
 creativity_slider = gr.Slider(minimum=0, maximum=1, step=0.1, label="创意程度 (0 到 1)", value=0.7)
 background_checkbox = gr.Checkbox(label="是否添加背景", value=True)
-
 audio_input = gr.Audio(type="numpy", label="用语音描述图像")
 
 # 创建文本输入的界面