import gradio as gr from PIL import Image import torch from diffusers import AutoPipelineForInpainting from diffusers.utils import load_image def draw_on_image(image, prompt): print(image, prompt) if not prompt: return init_image = load_image( image["image"] ) mask_image = load_image( image["mask"] ) res_image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0] return res_image inputs = [ gr.Image(tool="sketch", label="Image", type="pil"), gr.Text(max_lines=1) ] if torch.cuda.is_available(): torch_dtype = torch.float32 device = "cuda" else: torch_dtype = torch.float16 device = "cpu" pipeline = AutoPipelineForInpainting.from_pretrained( "kandinsky-community/kandinsky-2-2-decoder-inpaint", torch_dtype=torch_dtype ) pipeline.to(device) app = gr.Interface(draw_on_image, inputs=inputs, outputs="image") app.queue() app.launch(share=True)