Spaces:
Running
Running
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) |