Spaces:
Runtime error
Runtime error
File size: 3,134 Bytes
29fad46 |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
from diffusers.utils import load_image, make_image_grid
import gradio as gr
import cv2
import numpy as np
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, UniPCMultistepScheduler
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
controlnet = ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16, use_safetensors=True)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, use_safetensors=True
).to(device)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()
def create_mask(img):
img = cv2.imread(img, cv2.IMREAD_GRAYSCALE)
# Create a binary mask
mask = cv2.threshold(img, 254, 255, cv2.THRESH_BINARY)[1]
# Invert the mask
inverted_mask = cv2.bitwise_not(mask)
return mask
def load_and_resize_images(image_path, mask_path, target_size=(512, 512)):
init_image = load_image(image_path)
init_image = init_image.resize(target_size)
mask_image = load_image(mask_path)
mask_image = mask_image.resize(target_size)
return init_image, mask_image
def make_inpaint_condition(image, image_mask):
image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
assert image.shape[0:1] == image_mask.shape[0:1]
image[image_mask > 0.5] = -1.0 # set as masked pixel
image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
image = torch.from_numpy(image)
return image
def generate_inpainting(init_image, mask_image, prompt, negative_prompt=None):
control_image = make_inpaint_condition(init_image, mask_image)
output = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=50,
eta=1.0,
image=init_image,
mask_image=mask_image,
control_image=control_image,
).images[0]
return output
def process_images_and_generate_inpainting(image_path, mask_path, prompt, negative_prompt=None):
init_image, mask_image = load_and_resize_images(image_path, mask_path)
output_image = generate_inpainting(init_image, mask_image, prompt, negative_prompt)
return output_image
maskGen = gr.Interface(
fn=create_mask,
inputs=gr.Image(type='filepath', label="Original Image"),
outputs=gr.Image(label="Masked Image Image"),
allow_flagging="never")
inpaintImg = gr.Interface(
fn=process_images_and_generate_inpainting,
inputs=[
gr.Image(type='filepath', label="Original Image"),
gr.Image(type='filepath', label="Masked Image"),
gr.Textbox(label="Prompt"),
gr.Textbox(label="Negative Prompt")
],
outputs=[
gr.Image(label="Result Image")
],
allow_flagging="never")
demo = gr.TabbedInterface(
[maskGen, inpaintImg],
["Generate Mask", "Make new Image"]
)
demo.launch(share=True)
|