Spaces:
Sleeping
Sleeping
File size: 2,409 Bytes
a6170f9 17a1c6f a6170f9 17a1c6f a6170f9 2def794 a6170f9 2def794 a6170f9 aee1326 17a1c6f 92a1c24 6291a59 a6170f9 17a1c6f a6170f9 17a1c6f 6291a59 92a1c24 a6170f9 6438e34 a6170f9 7c014c9 a6170f9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import gradio as gr
import cv2
import requests
import numpy as np
from PIL import Image
from io import BytesIO
from rembg import remove
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import torch
from diffusers import UniPCMultistepScheduler
def remove_background(input_image: Image.Image, to_grayscale: bool) -> Image.Image:
output_image = remove(input_image)
if to_grayscale:
output_image = convert_to_grayscale(output_image)
return output_image
def convert_to_grayscale(image: Image.Image) -> Image.Image:
return image.convert("L")
def canny_image(image: Image.Image) -> Image.Image:
np_image = np.array(image)
low_threshold = 100
high_threshold = 200
np_image = cv2.Canny(np_image, low_threshold, high_threshold)
np_image = np_image[:, :, None]
np_image = np.concatenate([np_image, np_image, np_image], axis=2)
return Image.fromarray(np_image)
def process_image(input_image: Image.Image, to_grayscale: bool, prompt: str) -> Image.Image:
output_image = remove_background(input_image, to_grayscale)
canny_output = canny_image(output_image)
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float32)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32
)
pipe.enable_model_cpu_offload()
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
generator = torch.manual_seed(2)
output = pipe(
prompt,
canny_output,
negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
generator=generator,
num_inference_steps=20,
)
return output.images[0]
image_input = gr.components.Image(label="Input Image")
grayscale_checkbox = gr.components.Checkbox(label="Convert output to grayscale", default=False)
prompt_input = gr.components.Textbox(lines=1, label="Prompt")
image_output = gr.components.Image(label="Output Image", type="pil")
gr.Interface(
fn=process_image,
inputs=[image_input, grayscale_checkbox, prompt_input],
outputs=image_output,
title="ControlNet",
description="Upload an image and a prompt to generate an image with the prompt in the style of the input image.",
).launch() |