Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler | |
| def generate_image(image): | |
| if image is not None: | |
| sketch = image['layers'][0] | |
| with torch.no_grad(): # Disable gradient calculation for inference | |
| model_output = pipe(prompt, num_inference_steps=20, generator=torch.manual_seed(0), image=sketch) | |
| generated_image = model_output.images[0] | |
| return generated_image | |
| controlnet_model_name_or_path = "./controlnet" | |
| pretrained_model_name_or_path = "runwayml/stable-diffusion-v1-5" | |
| controlnet = ControlNetModel.from_pretrained(controlnet_model_name_or_path, torch_dtype=torch.float16, conditioning_channels=3) | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained(pretrained_model_name_or_path, controlnet=controlnet, torch_dtype=torch.float16, safety_checker=None) | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe.enable_model_cpu_offload() | |
| iface = gr.Interface( | |
| fn=generate_image, | |
| inputs=[ | |
| gr.ImageEditor(sources=(), brush=gr.Brush(colors=["#ffb266", #building | |
| "#4059ff", #parking | |
| "#66ff66", #grass | |
| #"#009900", #forest | |
| "#cce5ff", #water | |
| #"#c0c0c0", #path | |
| "#606060" #road | |
| ], color_mode="fixed")) | |
| #gr.Sketchpad() | |
| #gr.Image(shape=(512, 512), source="canvas", tool="sketch", image_mode="RGB", invert_colors=False, brush_color="black"), | |
| #gr.Dropdown(choices=allowed_colors, label="Brush Color") | |
| ], | |
| outputs="image" | |
| ) | |
| iface.launch() | |