app.py

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=4, step=1, minimum=1, maximum=16),
    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)