EleutherAI/VQGAN_CLIP

1import os2os.system('pip freeze')3 4import torch5torch.hub.download_url_to_file('https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1', 'vqgan_imagenet_f16_16384.yaml')6torch.hub.download_url_to_file('https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fckpts%2Flast.ckpt&dl=1', 'vqgan_imagenet_f16_16384.ckpt')7import argparse8import math9from pathlib import Path10import sys11sys.path.insert(1, './taming-transformers')12from base64 import b64encode13from omegaconf import OmegaConf14from PIL import Image15from taming.models import cond_transformer, vqgan16import taming.modules 17from torch import nn, optim18from torch.nn import functional as F19from torchvision import transforms20from torchvision.transforms import functional as TF21from tqdm.notebook import tqdm22from CLIP import clip23import kornia.augmentation as K24import numpy as np25import imageio26from PIL import ImageFile, Image27ImageFile.LOAD_TRUNCATED_IMAGES = True28import gradio as gr29torch.hub.download_url_to_file('https://images.pexels.com/photos/158028/bellingrath-gardens-alabama-landscape-scenic-158028.jpeg', 'garden.jpeg')30torch.hub.download_url_to_file('https://images.pexels.com/photos/68767/divers-underwater-ocean-swim-68767.jpeg', 'coralreef.jpeg')31torch.hub.download_url_to_file('https://images.pexels.com/photos/803975/pexels-photo-803975.jpeg', 'cabin.jpeg')32def sinc(x):33    return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))34def lanczos(x, a):35    cond = torch.logical_and(-a < x, x < a)36    out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))37    return out / out.sum()38def ramp(ratio, width):39    n = math.ceil(width / ratio + 1)40    out = torch.empty([n])41    cur = 042    for i in range(out.shape[0]):43        out[i] = cur44        cur += ratio45    return torch.cat([-out[1:].flip([0]), out])[1:-1]46def resample(input, size, align_corners=True):47    n, c, h, w = input.shape48    dh, dw = size49    input = input.view([n * c, 1, h, w])50    if dh < h:51        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)52        pad_h = (kernel_h.shape[0] - 1) // 253        input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')54        input = F.conv2d(input, kernel_h[None, None, :, None])55    if dw < w:56        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)57        pad_w = (kernel_w.shape[0] - 1) // 258        input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')59        input = F.conv2d(input, kernel_w[None, None, None, :])60    input = input.view([n, c, h, w])61    return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)62class ReplaceGrad(torch.autograd.Function):63    @staticmethod64    def forward(ctx, x_forward, x_backward):65        ctx.shape = x_backward.shape66        return x_forward67    @staticmethod68    def backward(ctx, grad_in):69        return None, grad_in.sum_to_size(ctx.shape)70replace_grad = ReplaceGrad.apply71class ClampWithGrad(torch.autograd.Function):72    @staticmethod73    def forward(ctx, input, min, max):74        ctx.min = min75        ctx.max = max76        ctx.save_for_backward(input)77        return input.clamp(min, max)78    @staticmethod79    def backward(ctx, grad_in):80        input, = ctx.saved_tensors81        return grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0), None, None82clamp_with_grad = ClampWithGrad.apply83def vector_quantize(x, codebook):84    d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T85    indices = d.argmin(-1)86    x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook87    return replace_grad(x_q, x)88class Prompt(nn.Module):89    def __init__(self, embed, weight=1., stop=float('-inf')):90        super().__init__()91        self.register_buffer('embed', embed)92        self.register_buffer('weight', torch.as_tensor(weight))93        self.register_buffer('stop', torch.as_tensor(stop))94    def forward(self, input):95        input_normed = F.normalize(input.unsqueeze(1), dim=2)96        embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)97        dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)98        dists = dists * self.weight.sign()99        return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()100def parse_prompt(prompt):101    vals = prompt.rsplit(':', 2)102    vals = vals + ['', '1', '-inf'][len(vals):]103    return vals[0], float(vals[1]), float(vals[2])104class MakeCutouts(nn.Module):105    def __init__(self, cut_size, cutn, cut_pow=1.):106        super().__init__()107        self.cut_size = cut_size108        self.cutn = cutn109        self.cut_pow = cut_pow110        self.augs = nn.Sequential(111            # K.RandomHorizontalFlip(p=0.5),112            # K.RandomVerticalFlip(p=0.5),113            # K.RandomSolarize(0.01, 0.01, p=0.7),114            # K.RandomSharpness(0.3,p=0.4),115            # K.RandomResizedCrop(size=(self.cut_size,self.cut_size), scale=(0.1,1),  ratio=(0.75,1.333), cropping_mode='resample', p=0.5),116            # K.RandomCrop(size=(self.cut_size,self.cut_size), p=0.5),117            K.RandomAffine(degrees=15, translate=0.1, p=0.7, padding_mode='border'),118            K.RandomPerspective(0.7,p=0.7),119            K.ColorJitter(hue=0.1, saturation=0.1, p=0.7),120            K.RandomErasing((.1, .4), (.3, 1/.3), same_on_batch=True, p=0.7),121            122)123        self.noise_fac = 0.1124        self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))125        self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))126    def forward(self, input):127        sideY, sideX = input.shape[2:4]128        max_size = min(sideX, sideY)129        min_size = min(sideX, sideY, self.cut_size)130        cutouts = []131        132        for _ in range(self.cutn):133            # size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)134            # offsetx = torch.randint(0, sideX - size + 1, ())135            # offsety = torch.randint(0, sideY - size + 1, ())136            # cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]137            # cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))138            # cutout = transforms.Resize(size=(self.cut_size, self.cut_size))(input)139            140            cutout = (self.av_pool(input) + self.max_pool(input))/2141            cutouts.append(cutout)142        batch = self.augs(torch.cat(cutouts, dim=0))143        if self.noise_fac:144            facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)145            batch = batch + facs * torch.randn_like(batch)146        return batch147def load_vqgan_model(config_path, checkpoint_path):148    config = OmegaConf.load(config_path)149    if config.model.target == 'taming.models.vqgan.VQModel':150        model = vqgan.VQModel(**config.model.params)151        model.eval().requires_grad_(False)152        model.init_from_ckpt(checkpoint_path)153    elif config.model.target == 'taming.models.vqgan.GumbelVQ':154        model = vqgan.GumbelVQ(**config.model.params)155        model.eval().requires_grad_(False)156        model.init_from_ckpt(checkpoint_path)157    elif config.model.target == 'taming.models.cond_transformer.Net2NetTransformer':158        parent_model = cond_transformer.Net2NetTransformer(**config.model.params)159        parent_model.eval().requires_grad_(False)160        parent_model.init_from_ckpt(checkpoint_path)161        model = parent_model.first_stage_model162    else:163        raise ValueError(f'unknown model type: {config.model.target}')164    del model.loss165    return model166def resize_image(image, out_size):167    ratio = image.size[0] / image.size[1]168    area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])169    size = round((area * ratio)**0.5), round((area / ratio)**0.5)170    return image.resize(size, Image.LANCZOS)171model_name = "vqgan_imagenet_f16_16384" 172images_interval =  50173width =  280174height = 280175init_image = ""176seed = 42177args = argparse.Namespace(178    noise_prompt_seeds=[],179    noise_prompt_weights=[],180    size=[width, height],181    init_image=init_image,182    init_weight=0.,183    clip_model='ViT-B/32',184    vqgan_config=f'{model_name}.yaml',185    vqgan_checkpoint=f'{model_name}.ckpt',186    step_size=0.15,187    cutn=4,188    cut_pow=1.,189    display_freq=images_interval,190    seed=seed,191)192device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')193print('Using device:', device)194model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)195perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)196def inference(text, seed, step_size, max_iterations, width, height, init_image, init_weight, target_images, cutn, cut_pow):197    torch.cuda.empty_cache()198    torch.cuda.memory_summary(device=None, abbreviated=False)199    all_frames = []200    size=[width, height]201    texts = text202    init_weight=init_weight203    if init_image:204        init_image = init_image.name205    else:206        init_image = ""207    if target_images:208        target_images = target_images.name209    else:210        target_images = ""211    max_iterations = max_iterations212    model_names={"vqgan_imagenet_f16_16384": 'ImageNet 16384',"vqgan_imagenet_f16_1024":"ImageNet 1024", 'vqgan_openimages_f16_8192':'OpenImages 8912',213                    "wikiart_1024":"WikiArt 1024", "wikiart_16384":"WikiArt 16384", "coco":"COCO-Stuff", "faceshq":"FacesHQ", "sflckr":"S-FLCKR"}214    name_model = model_names[model_name]215    if target_images == "None" or not target_images:216        target_images = []217    else:218        target_images = target_images.split("|")219        target_images = [image.strip() for image in target_images]220    texts = [phrase.strip() for phrase in texts.split("|")]221    if texts == ['']:222        texts = []223    from urllib.request import urlopen224    if texts:225        print('Using texts:', texts)226    if target_images:227        print('Using image prompts:', target_images)228    if seed is None or seed == -1:229        seed = torch.seed()230    else:231        seed = seed232    torch.manual_seed(seed)233    print('Using seed:', seed)234    # clock=deepcopy(perceptor.visual.positional_embedding.data)235    # perceptor.visual.positional_embedding.data = clock/clock.max()236    # perceptor.visual.positional_embedding.data=clamp_with_grad(clock,0,1)237    cut_size = perceptor.visual.input_resolution238    f = 2**(model.decoder.num_resolutions - 1)239    make_cutouts = MakeCutouts(cut_size, cutn, cut_pow=cut_pow)240    toksX, toksY = size[0] // f, size[1] // f241    sideX, sideY = toksX * f, toksY * f242    if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':243        e_dim = 256244        n_toks = model.quantize.n_embed245        z_min = model.quantize.embed.weight.min(dim=0).values[None, :, None, None]246        z_max = model.quantize.embed.weight.max(dim=0).values[None, :, None, None]247    else:248        e_dim = model.quantize.e_dim249        n_toks = model.quantize.n_e250        z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]251        z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]252    # z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]253    # z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]254    # normalize_imagenet = transforms.Normalize(mean=[0.485, 0.456, 0.406],255    #                                            std=[0.229, 0.224, 0.225])256    if init_image:257        if 'http' in init_image:258            img = Image.open(urlopen(init_image))259        else:260            img = Image.open(init_image)261        pil_image = img.convert('RGB')262        pil_image = pil_image.resize((sideX, sideY), Image.LANCZOS)263        pil_tensor = TF.to_tensor(pil_image)264        z, *_ = model.encode(pil_tensor.to(device).unsqueeze(0) * 2 - 1)265    else:266        one_hot = F.one_hot(torch.randint(n_toks, [toksY * toksX], device=device), n_toks).float()267        # z = one_hot @ model.quantize.embedding.weight268        if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':269            z = one_hot @ model.quantize.embed.weight270        else:271            z = one_hot @ model.quantize.embedding.weight272        z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2) 273        z = torch.rand_like(z)*2274    z_orig = z.clone()275    z.requires_grad_(True)276    opt = optim.Adam([z], lr=step_size)277    normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],278                                    std=[0.26862954, 0.26130258, 0.27577711])279    pMs = []280    for prompt in texts:281        txt, weight, stop = parse_prompt(prompt)282        embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()283        pMs.append(Prompt(embed, weight, stop).to(device))284    for prompt in target_images:285        path, weight, stop = parse_prompt(prompt)286        img = Image.open(path)287        pil_image = img.convert('RGB')288        img = resize_image(pil_image, (sideX, sideY))289        batch = make_cutouts(TF.to_tensor(img).unsqueeze(0).to(device))290        embed = perceptor.encode_image(normalize(batch)).float()291        pMs.append(Prompt(embed, weight, stop).to(device))292    for seed, weight in zip(args.noise_prompt_seeds, args.noise_prompt_weights):293        gen = torch.Generator().manual_seed(seed)294        embed = torch.empty([1, perceptor.visual.output_dim]).normal_(generator=gen)295        pMs.append(Prompt(embed, weight).to(device))296    def synth(z):297        if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':298            z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)299        else:300            z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)301        return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)302    @torch.no_grad()303    def checkin(i, losses):304        losses_str = ', '.join(f'{loss.item():g}' for loss in losses)305        tqdm.write(f'i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}')306        out = synth(z)307        # TF.to_pil_image(out[0].cpu()).save('progress.png')308        # display.display(display.Image('progress.png'))309        res = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)310        print(f'gpu: {res.gpu}%, gpu-mem: {res.memory}%')311    def ascend_txt():312        # global i313        out = synth(z)314        iii = perceptor.encode_image(normalize(make_cutouts(out))).float()315        316        result = []317        if init_weight:318            result.append(F.mse_loss(z, z_orig) * init_weight / 2)319            #result.append(F.mse_loss(z, torch.zeros_like(z_orig)) * ((1/torch.tensor(i*2 + 1))*init_weight) / 2)320        for prompt in pMs:321            result.append(prompt(iii))322        img = np.array(out.mul(255).clamp(0, 255)[0].cpu().detach().numpy().astype(np.uint8))[:,:,:]323        img = np.transpose(img, (1, 2, 0))324        # imageio.imwrite('./steps/' + str(i) + '.png', np.array(img))325        img = Image.fromarray(img).convert('RGB')326        all_frames.append(img)327        return result, np.array(img)328    def train(i):329        opt.zero_grad()330        lossAll, image = ascend_txt()331        if i % args.display_freq == 0:332            checkin(i, lossAll)333        334        loss = sum(lossAll)335        loss.backward()336        opt.step()337        with torch.no_grad():338            z.copy_(z.maximum(z_min).minimum(z_max))339        return image340    i = 0341    try:342        with tqdm() as pbar:343            while True:344                image = train(i)345                if i == max_iterations:346                    break347                i += 1348                pbar.update()349    except KeyboardInterrupt:350        pass351    writer = imageio.get_writer('test.mp4', fps=20)352 353    for im in all_frames:354        writer.append_data(np.array(im))355    writer.close()356   # all_frames[0].save('out.gif',357              # save_all=True, append_images=all_frames[1:], optimize=False, duration=80, loop=0)358    return image, 'test.mp4'359    360def load_image( infilename ) :361    img = Image.open( infilename )362    img.load()363    data = np.asarray( img, dtype="int32" )364    return data365title = "VQGAN + CLIP"366description = "Gradio demo for VQGAN + CLIP. To use it, simply add your text, or click one of the examples to load them. Read more at the links below."367article = "<p style='text-align: center'>Originally made by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings). The original BigGAN+CLIP method was by https://twitter.com/advadnoun. Added some explanations and modifications by Eleiber#8347, pooling trick by Crimeacs#8222 (https://twitter.com/EarthML1) and the GUI was made with the help of Abulafia#3734. | <a href='https://colab.research.google.com/drive/1ZAus_gn2RhTZWzOWUpPERNC0Q8OhZRTZ'>Colab</a> | <a href='https://github.com/CompVis/taming-transformers'>Taming Transformers Github Repo</a> | <a href='https://github.com/openai/CLIP'>CLIP Github Repo</a> | Special thanks to BoneAmputee (https://twitter.com/BoneAmputee) for suggestions and advice</p>"368gr.Interface(369    inference, 370    [gr.inputs.Textbox(label="Text Input"),371     gr.inputs.Number(default=42, label="seed"),372     gr.inputs.Slider(minimum=0.1, maximum=0.9, default=0.6, label='step size'),373    gr.inputs.Slider(minimum=1, maximum=500, default=100, label='max iterations', step=1),374    gr.inputs.Slider(minimum=200, maximum=600, default=256, label='width', step=1),375    gr.inputs.Slider(minimum=200, maximum=600, default=256, label='height', step=1),376    gr.inputs.Image(type="file", label="Initial Image (Optional)", optional=True),377    gr.inputs.Slider(minimum=0.0, maximum=15.0, default=0.0, label='Initial Weight', step=1.0),378    gr.inputs.Image(type="file", label="Target Image (Optional)", optional=True),379    gr.inputs.Slider(minimum=1, maximum=40, default=1, label='cutn', step=1),380    gr.inputs.Slider(minimum=1.0, maximum=40.0, default=1.0, label='cut_pow', step=1.0)381     ], 382    [gr.outputs.Image(type="numpy", label="Output Image"),gr.outputs.Video(label="Output Video")],383    title=title,384    description=description,385    article=article,386    examples=[387              ['a garden by james gurney',42,0.6, 100, 256, 256, 'garden.jpeg', 0.0, 'garden.jpeg',1,1.0],388              ['coral reef city artstationHQ',1000,0.6, 110, 200, 200, 'coralreef.jpeg', 0.0, 'coralreef.jpeg',1,1.0],389              ['a cabin in the mountains unreal engine',98,0.6, 120, 280, 280, 'cabin.jpeg', 0.0, 'cabin.jpeg',1,1.0]390    ]391    ).launch(enable_queue=True)392