Update app.py

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
 from transformers import pipeline, T5ForConditionalGeneration, T5Tokenizer
 import torch
-import stable_diffusion_webnn  # 假设stable-diffusion-v1.5-webnn的库名为 stable_diffusion_webnn
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from diffusers import StableDiffusionPipeline
+import cv2
+import numpy as np
 
 # 使用T5模型生成文本描述
-model_name = "t5-large"  # 可以根据需求选择不同版本的T5
+model_name = "t5-large"  # 你可以根据需求选择不同版本的T5
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 t5_model = T5ForConditionalGeneration.from_pretrained(model_name)
 
@@ -16,15 +19,33 @@ def generate_prompt(description: str) -> str:
     prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return prompt
 
-# 使用 stable-diffusion-v1.5-webnn 库加载 Stable Diffusion 模型
-# 这里假设 stable_diffusion_webnn 可以直接加载模型并生成图像
-text2image_pipeline = stable_diffusion_webnn.StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
-text2image_pipeline.to("cpu")  # 使用CPU
+# 加载 ControlNet Canny 边缘检测模型
+controlnet = ControlNetModel.from_pretrained("lllyasviel/controlnet-canny-sdxl-1.0")
+stable_diffusion = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-1.0")
 
-def generate_image(prompt: str):
-    # 根据提示生成图像
-    image = text2image_pipeline(prompt).images[0]
-    return image
+# 使用 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
 
 # 使用Whisper模型进行语音转文本
 from transformers import WhisperProcessor, WhisperForConditionalGeneration
@@ -43,20 +64,21 @@ def process_input(description: str, creativity: float, include_background: bool)
     prompt = generate_prompt(description)
     if include_background:
         prompt += " 添加详细的生动背景。"
-    image = generate_image(prompt)
+    image = generate_image_with_controlnet(prompt)
     return prompt, image
 
 # 处理音频输入和生成图像
 def process_audio_input(audio):
     description = transcribe_audio(audio)
     prompt = generate_prompt(description)
-    image = generate_image(prompt)
+    image = generate_image_with_controlnet(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="用语音描述图像")
 
 # 创建文本输入的界面