File size: 3,619 Bytes
5d82ba5 9ac174a c6ad3ed 5d82ba5 c6ad3ed 5d82ba5 b0e960e 5d82ba5 9ac174a 5d82ba5 b0e960e 5d82ba5 200468f b0e960e fca2f7e b0e960e 5d82ba5 b0e960e 5d82ba5 cb9a173 5d82ba5 b0e960e 5d82ba5 cb9a173 0d76975 b0e960e 5d82ba5 4b49fe7 5d82ba5 b0e960e 5d82ba5 b0e960e 5d82ba5 4b49fe7 5d82ba5 b0e960e 89ae04a adce2e3 b0e960e fca2f7e b0e960e |
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 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import os
import spaces
import gradio as gr
import torch
import yaml
import numpy as np
from munch import munchify
import torchvision.transforms as transforms
from torchvision.transforms import functional as F
from diffusers import (
AutoPipelineForInpainting,
)
from generate_dataset import outpainting_generator_rectangle, merge_images_horizontally
from ddim_with_prob import DDIMSchedulerCustom
transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((512, 512), interpolation=F.InterpolationMode.LANCZOS),
])
@spaces.GPU(duration=120)
def pref_inpainting(image,
box_width_ratio,
mask_random_start,
steps,
):
with open("./configs/paintreward_train_configs.yaml") as file:
config_dict= yaml.safe_load(file)
config = munchify(config_dict)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Current Device is {device}")
pipe_ours = AutoPipelineForInpainting.from_pretrained(
'./model_ckpt', torch_dtype=torch.float16, variant='fp16')
pipe_ours.scheduler = DDIMSchedulerCustom.from_config(pipe_ours.scheduler.config)
pipe_runway = AutoPipelineForInpainting.from_pretrained("runwayml/stable-diffusion-inpainting", torch_dtype=torch.float16, variant='fp16')
pipe_ours = pipe_ours.to(device)
pipe_runway = pipe_runway.to(device)
print('Loading pipeline')
color, mask = outpainting_generator_rectangle(image, box_width_ratio/100, mask_random_start)
mask = mask.convert('L')
color, mask = np.array(color).transpose(2, 0, 1), np.array(mask)
mask = mask[None, ...]
mask_ = np.zeros_like(mask)
mask_[mask < 125] = 0
mask_[mask >= 125] = 1
color = torch.from_numpy(color).to(device)
mask = torch.from_numpy(mask).to(device)
color, mask = transform(color), transform(mask)
res_ours = pipe_ours(prompt='', image=color, mask_image=mask, eta=config.eta).images[0]
print('Running inference ours')
res_runway = pipe_runway(prompt="", image=color, mask_image=mask).images[0]
print('Running inference runway')
# res.save(os.path.join('./', 'test.png'))
res_ours = merge_images_horizontally(color, res_ours, logo_path='./logo/pref_logo.png')
res_runway = merge_images_horizontally(color, res_runway, logo_path='./logo/runway_logo.png')
return res_ours, res_runway
inputs = [
gr.Image(type="pil", image_mode="RGBA", label='Input Image'), # shape=[512, 512]
gr.Slider(25, 45, value=35, step=1, label="box_width_ratio"),
gr.Slider(0, 256, value=125, step=1, label="mask_random_start"),
gr.Slider(30, 100, value=50, step=5, label="steps"),
]
outputs = [
gr.Image(type="pil", image_mode="RGBA", label='PrefPaint', container=True, width="100%"),
gr.Image(type="pil", image_mode="RGBA", label='RunwayPaint', container=True, width="100%"),
]
files = os.listdir("./assets")
examples = [
[f"./assets/{file_name}", 25, 125, 50] for file_name in files
]
with gr.Blocks() as demo:
iface = gr.Interface(
fn=pref_inpainting,
inputs=inputs,
outputs=outputs,
title="Inpainting with Human Preference (Only one GPU is available, you may need to queue.)",
description="Upload an image and start your inpainting (Currently, only outpainting masks are supported; other mask types will be available soon.).",
theme="default",
examples=examples,
# allow_flagging="never"
)
# iface.launch()
demo.launch() |