| import gradio as gr | |
| import torch | |
| import numpy as np | |
| from PIL import Image | |
| from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler | |
| from diffusers.utils import make_image_grid | |
| import cv2 | |
| controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny") | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| "stable-diffusion-v1-5/stable-diffusion-v1-5", controlnet=controlnet | |
| ) | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe.to("cpu") | |
| def generate_image(input_image, text_prompt): | |
| original_image = np.array(input_image) | |
| low_threshold = 100 | |
| high_threshold = 200 | |
| edges = cv2.Canny(original_image, low_threshold, high_threshold) | |
| edges = edges[:, :, None] | |
| canny_image = np.concatenate([edges, edges, edges], axis=2) | |
| canny_image_pil = Image.fromarray(canny_image) | |
| output_image = pipe(text_prompt, image=canny_image_pil).images[0] | |
| result_grid = make_image_grid([input_image, canny_image_pil, output_image], rows=1, cols=3) | |
| return result_grid | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Image Transformation with ControlNet and Stable Diffusion (CPU Only)") | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image(type="pil", label="Upload Image") | |
| text_prompt = gr.Textbox(label="Enter a prompt for the transformation") | |
| generate_button = gr.Button("Generate Image") | |
| result = gr.Image(label="Result") | |
| generate_button.click(fn=generate_image, inputs=[input_image, text_prompt], outputs=result) | |
| demo.launch() |