Spaces:
Runtime error
Runtime error
import gradio as gr | |
import argparse | |
from einops import rearrange | |
#from glide_text2im import dist_util, logger | |
from torchvision.utils import make_grid | |
from glide_text2im.script_util import ( | |
model_and_diffusion_defaults, | |
create_model_and_diffusion, | |
args_to_dict, | |
add_dict_to_argparser, | |
) | |
from glide_text2im.image_datasets_sketch import get_tensor | |
from glide_text2im.train_util import TrainLoop | |
from glide_text2im.glide_util import sample | |
import torch | |
import os | |
import torch as th | |
import torchvision.utils as tvu | |
import torch.distributed as dist | |
from PIL import Image | |
import cv2 | |
import numpy as np | |
from huggingface_hub import hf_hub_download | |
def run(image, mode, sample_c=1.3, num_samples=3, sample_step=100): | |
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") | |
parser, parser_up = create_argparser() | |
args = parser.parse_args() | |
args_up = parser_up.parse_args() | |
#dist_util.setup_dist() | |
if mode == 'sketch': | |
args.mode = 'coco-edge' | |
args_up.mode = 'coco-edge' | |
args.model_path = hf_hub_download(repo_id="tfwang/PITI", filename="base.pt") | |
args.sr_model_path = hf_hub_download(repo_id="tfwang/PITI", filename="upsample.pt") | |
elif mode == 'mask': | |
args.mode = 'coco' | |
args_up.mode = 'coco' | |
args.model_path = hf_hub_download(repo_id="tfwang/PITI", filename="base_mask.pt") | |
args.sr_model_path = hf_hub_download(repo_id="tfwang/PITI", filename="upsample_mask.pt") | |
args.val_data_dir = image | |
args.sample_c = sample_c | |
args.num_samples = num_samples | |
options=args_to_dict(args, model_and_diffusion_defaults(0.).keys()) | |
model, diffusion = create_model_and_diffusion(**options) | |
options_up=args_to_dict(args_up, model_and_diffusion_defaults(True).keys()) | |
model_up, diffusion_up = create_model_and_diffusion(**options_up) | |
if args.model_path: | |
print('loading model') | |
model_ckpt = torch.load(args.model_path, map_location="cpu") | |
model.load_state_dict( | |
model_ckpt , strict=True ) | |
if args.sr_model_path: | |
print('loading sr model') | |
model_ckpt2 = torch.load(args.sr_model_path, map_location="cpu") | |
model_up.load_state_dict( | |
model_ckpt2 , strict=True ) | |
model.to(device) | |
model_up.to(device) | |
model.eval() | |
model_up.eval() | |
########### dataset | |
# logger.log("creating data loader...") | |
if args.mode == 'coco': | |
pil_image = image | |
label_pil = pil_image.convert("RGB").resize((256, 256), Image.NEAREST) | |
label_tensor = get_tensor()(label_pil) | |
data_dict = {"ref":label_tensor.unsqueeze(0).repeat(args.num_samples, 1, 1, 1)} | |
elif args.mode == 'coco-edge': | |
# pil_image = Image.open(image) | |
pil_image = image | |
label_pil = pil_image.convert("L").resize((256, 256), Image.NEAREST) | |
im_dist = cv2.distanceTransform(255-np.array(label_pil), cv2.DIST_L1, 3) | |
im_dist = np.clip((im_dist) , 0, 255).astype(np.uint8) | |
im_dist = Image.fromarray(im_dist).convert("RGB") | |
label_tensor = get_tensor()(im_dist)[:1] | |
data_dict = {"ref":label_tensor.unsqueeze(0).repeat(args.num_samples, 1, 1, 1)} | |
print("sampling...") | |
sampled_imgs = [] | |
grid_imgs = [] | |
img_id = 0 | |
while (True): | |
if img_id >= args.num_samples: | |
break | |
model_kwargs = data_dict | |
with th.no_grad(): | |
samples_lr =sample( | |
glide_model= model, | |
glide_options= options, | |
side_x= 64, | |
side_y= 64, | |
prompt=model_kwargs, | |
batch_size= args.num_samples, | |
guidance_scale=args.sample_c, | |
device=device, | |
prediction_respacing= str(sample_step), | |
upsample_enabled= False, | |
upsample_temp=0.997, | |
mode = args.mode, | |
) | |
samples_lr = samples_lr.clamp(-1, 1) | |
tmp = (127.5*(samples_lr + 1.0)).int() | |
model_kwargs['low_res'] = tmp/127.5 - 1. | |
samples_hr =sample( | |
glide_model= model_up, | |
glide_options= options_up, | |
side_x=256, | |
side_y=256, | |
prompt=model_kwargs, | |
batch_size=args.num_samples, | |
guidance_scale=1, | |
device=device, | |
prediction_respacing= "fast27", | |
upsample_enabled=True, | |
upsample_temp=0.997, | |
mode = args.mode, | |
) | |
samples_hr = samples_hr | |
for hr in samples_hr: | |
hr = 255. * rearrange((hr.cpu().numpy()+1.0)*0.5, 'c h w -> h w c') | |
sample_img = Image.fromarray(hr.astype(np.uint8)) | |
sampled_imgs.append(sample_img) | |
img_id += 1 | |
grid_imgs.append(samples_hr) | |
grid = torch.stack(grid_imgs, 0) | |
grid = rearrange(grid, 'n b c h w -> (n b) c h w') | |
grid = make_grid(grid, nrow=2) | |
# to image | |
grid = 255. * rearrange((grid+1.0)*0.5, 'c h w -> h w c').cpu().numpy() | |
return Image.fromarray(grid.astype(np.uint8)) | |
def create_argparser(): | |
defaults = dict( | |
data_dir="", | |
val_data_dir="", | |
model_path="./base_edge.pt", | |
sr_model_path="./upsample_edge.pt", | |
encoder_path="", | |
schedule_sampler="uniform", | |
lr=1e-4, | |
weight_decay=0.0, | |
lr_anneal_steps=0, | |
batch_size=2, | |
microbatch=-1, # -1 disables microbatches | |
ema_rate="0.9999", # comma-separated list of EMA values | |
log_interval=100, | |
save_interval=20000, | |
resume_checkpoint="", | |
use_fp16=False, | |
fp16_scale_growth=1e-3, | |
sample_c=1., | |
sample_respacing="100", | |
uncond_p=0.2, | |
num_samples=3, | |
finetune_decoder = False, | |
mode = '', | |
) | |
defaults_up = defaults | |
defaults.update(model_and_diffusion_defaults()) | |
parser = argparse.ArgumentParser() | |
add_dict_to_argparser(parser, defaults) | |
defaults_up.update(model_and_diffusion_defaults(True)) | |
parser_up = argparse.ArgumentParser() | |
add_dict_to_argparser(parser_up, defaults_up) | |
return parser, parser_up | |
image = gr.outputs.Image(type="pil", label="Sampled results") | |
css = ".output-image{height: 528px !important} .output-carousel .output-image{height:272px !important} a{text-decoration: underline}" | |
iface = gr.Interface(fn=run, inputs=[ | |
gr.inputs.Image(type="pil", label="Input Sketch or Mask" ) , | |
# gr.Image(image_mode="L", source="canvas", type="pil", shape=(256,256), invert_colors=False, tool="editor"), | |
gr.inputs.Radio(label="Input Mode - The type of your input", choices=["mask", "sketch"],default="sketch"), | |
gr.inputs.Slider(label="sample_c - The strength of classifier-free guidance",default=1.4, minimum=1.0, maximum=2.0), | |
gr.inputs.Slider(label="Number of samples - How many samples you wish to generate", default=2, step=1, minimum=1, maximum=8), | |
gr.inputs.Slider(label="Number of Steps - How many steps you want to use", default=100, step=10, minimum=50, maximum=1000), | |
], | |
outputs=[image], | |
css=css, | |
title="Generate images from sketches with PITI", | |
description="<div>By uploading a sketch map or a semantic map and pressing submit, you can generate images based on your input. As the computing device is CPU, the running may be slow.</div>", | |
examples=[["1.png", "sketch", 1.3, 4, 100], ["2.png", "sketch", 1.3, 4, 100],["3.png", "sketch", 1.3, 4, 100],["4.png", "mask", 1.3, 4, 100],["5.png", "mask", 1.3, 4, 100],["6.png", "mask", 1.3, 4, 100]]) | |
iface.launch(enable_queue=True) | |