app.py

import gradio as gr
from git.repo.base import Repo
from os.path import exists as path_exists

if not (path_exists(f"stylegan_xl")):
    Repo.clone_from("https://github.com/autonomousvision/stylegan_xl", "stylegan_xl")

import sys
sys.path.append('./CLIP')
sys.path.append('./stylegan_xl')

import io
import os, time, glob
import pickle
import shutil
import numpy as np
from PIL import Image
import torch
import torch.nn.functional as F
import requests
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
import clip
import unicodedata
import re
from tqdm import tqdm
from torchvision.transforms import Compose, Resize, ToTensor, Normalize
from IPython.display import display
from einops import rearrange
import dnnlib
import legacy
import subprocess

torch.cuda.empty_cache()
device = torch.device('cuda:0')
print('Using device:', device, file=sys.stderr)

def fetch(url_or_path):
    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):
        r = requests.get(url_or_path)
        r.raise_for_status()
        fd = io.BytesIO()
        fd.write(r.content)
        fd.seek(0)
        return fd
    return open(url_or_path, 'rb')

def fetch_model(url_or_path,network_name):
    print(network_name)
    torch.hub.download_url_to_file(f'{url_or_path}',f'{network_name}')
    print(os.listdir())

def slugify(value, allow_unicode=False):
    """
    Taken from https://github.com/django/django/blob/master/django/utils/text.py
    Convert to ASCII if 'allow_unicode' is False. Convert spaces or repeated
    dashes to single dashes. Remove characters that aren't alphanumerics,
    underscores, or hyphens. Convert to lowercase. Also strip leading and
    trailing whitespace, dashes, and underscores.
    """
    value = str(value)
    if allow_unicode:
        value = unicodedata.normalize('NFKC', value)
    else:
        value = unicodedata.normalize('NFKD', value).encode('ascii', 'ignore').decode('ascii')
    value = re.sub(r'[^\w\s-]', '', value.lower())
    return re.sub(r'[-\s]+', '-', value).strip('-_')

def norm1(prompt):
    "Normalize to the unit sphere."
    return prompt / prompt.square().sum(dim=-1,keepdim=True).sqrt()

def spherical_dist_loss(x, y):
    x = F.normalize(x, dim=-1)
    y = F.normalize(y, dim=-1)
    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)

def prompts_dist_loss(x, targets, loss):
    if len(targets) == 1: # Keeps consitent results vs previous method for single objective guidance 
      return loss(x, targets[0])
    distances = [loss(x, target) for target in targets]
    return torch.stack(distances, dim=-1).sum(dim=-1)  

class MakeCutouts(torch.nn.Module):
    def __init__(self, cut_size, cutn, cut_pow=1.):
        super().__init__()
        self.cut_size = cut_size
        self.cutn = cutn
        self.cut_pow = cut_pow

    def forward(self, input):
        sideY, sideX = input.shape[2:4]
        max_size = min(sideX, sideY)
        min_size = min(sideX, sideY, self.cut_size)
        cutouts = []
        for _ in range(self.cutn):
            size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)
            offsetx = torch.randint(0, sideX - size + 1, ())
            offsety = torch.randint(0, sideY - size + 1, ())
            cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]
            cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))
        return torch.cat(cutouts)

make_cutouts = MakeCutouts(224, 32, 0.5)

def embed_image(image):
  n = image.shape[0]
  cutouts = make_cutouts(image)
  embeds = clip_model.embed_cutout(cutouts)
  embeds = rearrange(embeds, '(cc n) c -> cc n c', n=n)
  return embeds

def embed_url(url):
  image = Image.open(fetch(url)).convert('RGB')
  return embed_image(TF.to_tensor(image).to(device).unsqueeze(0)).mean(0).squeeze(0)

class CLIP(object):
  def __init__(self):
    clip_model = "ViT-B/32"
    self.model, _ = clip.load(clip_model)
    self.model = self.model.requires_grad_(False)
    self.normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                                          std=[0.26862954, 0.26130258, 0.27577711])

  @torch.no_grad()
  def embed_text(self, prompt):
      "Normalized clip text embedding."
      return norm1(self.model.encode_text(clip.tokenize(prompt).to(device)).float())

  def embed_cutout(self, image):
      "Normalized clip image embedding."
      return norm1(self.model.encode_image(self.normalize(image)))
  
clip_model = CLIP()

#@markdown #**Model selection** 🎭

Models = ["imagenet256", "Pokemon", "FFHQ"]

#@markdown ---

network_url = {
    "imagenet256":"https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet256.pkl",
    #"Imagenet512": "https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet512.pkl",
    #"Imagenet1024": "https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/imagenet1024.pkl",
    "Pokemon": "https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/pokemon256.pkl",
    "FFHQ": "https://s3.eu-central-1.amazonaws.com/avg-projects/stylegan_xl/models/ffhq256.pkl"
}

for Model in Models:
    network_name = network_url[Model].split("/")[-1]
    if not (path_exists(network_name)):
        fetch_model(network_url[Model],network_name)

def load_current_model(current_model="imagenet256.pkl"):
    with dnnlib.util.open_url(current_model) as f:
        G = legacy.load_network_pkl(f)['G_ema'].to(device)

    zs = torch.randn([10000, G.mapping.z_dim], device=device)
    cs = torch.zeros([10000, G.mapping.c_dim], device=device)
    for i in range(cs.shape[0]):
        cs[i,i//10]=1
    w_stds = G.mapping(zs, cs)
    w_stds = w_stds.reshape(10, 1000, G.num_ws, -1)
    w_stds=w_stds.std(0).mean(0)[0]
    w_all_classes_avg = G.mapping.w_avg.mean(0)
    return(G,w_stds,w_all_classes_avg)

G, w_stds, w_all_classes_avg = load_current_model()
previousModel = 'imagenet256'
def run(prompt,steps,model):
    global G, w_stds, w_all_classes_avg, previousModel
    if(model == 'imagenet256' and previousModel != 'imagenet256'):
        G, w_stds, w_all_classes_avg = load_current_model('imagenet256.pkl')
    #if(model == 'imagenet512' and previousModel != 'imagenet512'):
    #    G, w_stds, w_all_classes_avg = load_current_model('imagenet512.pkl')
    #elif(model=='imagenet1024' and previousModel != 'imagenet1024'):
    #    G, w_stds, w_all_classes_avg = load_current_model('imagenet1024.pkl')
    elif(model=='pokemon256' and previousModel != 'pokemon256'):
        G, w_stds, w_all_classes_avg = load_current_model('pokemon256.pkl')
    elif(model=='ffhq256' and previousModel != 'ffhq256'):
        G, w_stds, w_all_classes_avg = load_current_model('ffhq256.pkl')
    previousModel = model

    texts = prompt
    steps = steps
    seed = -1  # @param {type:"number"}

    # @markdown ---

    if seed == -1:
        seed = np.random.randint(0, 9e9)
        print(f"Your random seed is: {seed}")

    texts = [frase.strip() for frase in texts.split("|") if frase]

    targets = [clip_model.embed_text(text) for text in texts]

    tf = Compose(
        [
            # Resize(224),
            lambda x: torch.clamp((x + 1) / 2, min=0, max=1),
        ]
    )

    initial_batch = 2  # actually that will be multiplied by initial_image_steps
    initial_image_steps = 8

    def get_image(timestring):
        os.makedirs(f"samples/{timestring}", exist_ok=True)
        torch.manual_seed(seed)
        with torch.no_grad():
            qs = []
            losses = []
            for _ in range(initial_image_steps):
                a = torch.randn([initial_batch, 512], device=device) * 0.4 + w_stds * 0.4
                q = (a - w_all_classes_avg) / w_stds
                images = G.synthesis(
                    (q * w_stds + w_all_classes_avg).unsqueeze(1).repeat([1, G.num_ws, 1])
                )
                embeds = embed_image(images.add(1).div(2))
                loss = prompts_dist_loss(embeds, targets, spherical_dist_loss).mean(0)
                i = torch.argmin(loss)
                qs.append(q[i])
                losses.append(loss[i])
            qs = torch.stack(qs)
            losses = torch.stack(losses)
            i = torch.argmin(losses)
            q = qs[i].unsqueeze(0).repeat([G.num_ws, 1]).requires_grad_()

        # Sampling loop
        q_ema = q
        print(q.shape)
        opt = torch.optim.AdamW([q], lr=0.05, betas=(0.0, 0.999), weight_decay=0.025)
        loop = tqdm(range(steps))
        for i in loop:
            opt.zero_grad()
            w = q * w_stds
            image = G.synthesis((q * w_stds + w_all_classes_avg)[None], noise_mode="const")
            embed = embed_image(image.add(1).div(2))
            loss = prompts_dist_loss(embed, targets, spherical_dist_loss).mean()
            loss.backward()
            opt.step()
            loop.set_postfix(loss=loss.item(), q_magnitude=q.std().item())

            q_ema = q_ema * 0.98 + q * 0.02
            image = G.synthesis(
                (q_ema * w_stds + w_all_classes_avg)[None], noise_mode="const"
            )
            
            pil_image = TF.to_pil_image(image[0].add(1).div(2).clamp(0, 1))
            pil_image.save(f"samples/{timestring}/{i:04}.jpg")
            
            if (i+1) % steps == 0:
                #/usr/bin/
                subprocess.call(['ffmpeg', '-r', '60', '-i', f'samples/{timestring}/%04d.jpg', '-vcodec', 'libx264', '-crf','18','-pix_fmt','yuv420p', f'{timestring}.mp4'])
                shutil.rmtree(f"samples/{timestring}")
                pil_image = TF.to_pil_image(image[0].add(1).div(2).clamp(0, 1))
                return(pil_image, f'{timestring}.mp4')
                    
    try:
        timestring = time.strftime("%Y%m%d%H%M%S")
        image,video = get_image(timestring)
        return([image,video])
    except KeyboardInterrupt:
        pass

image = gr.outputs.Image(type="pil", label="Your imge")
video = gr.outputs.Video(type="mp4", label="Your video")
css = ".output-image{height: 528px !important} .output-video{height: 528px !important}"
iface = gr.Interface(fn=run, inputs=[
    gr.inputs.Textbox(label="Prompt",default="Hong Kong by Studio Ghibli"),
    gr.inputs.Slider(label="Steps - more steps can increase quality but will take longer to generate",default=300,maximum=500,minimum=10,step=1),
    #gr.inputs.Radio(label="Aspect Ratio", choices=["Square", "Horizontal", "Vertical"],default="Horizontal"),
    gr.inputs.Dropdown(label="Model", choices=["imagenet256","Pokemon256", "ffhq256"], default="imagenet256")
    #gr.inputs.Radio(label="Height", choices=[32,64,128,256,512],default=256),
    #gr.inputs.Slider(label="Images - How many images you wish to generate", default=2, step=1, minimum=1, maximum=4),
    #gr.inputs.Slider(label="Diversity scale - How different from one another you wish the images to be",default=5.0, minimum=1.0, maximum=15.0),
    #gr.inputs.Slider(label="ETA - between 0 and 1. Lower values can provide better quality, higher values can be more diverse",default=0.0,minimum=0.0, maximum=1.0,step=0.1),
    ], 
    outputs=[image,video],
    css=css,
    title="Generate images from text with StyleGAN XL + CLIP",
    description="<div>By typing a prompt and pressing submit you generate images based on it. <a href='https://github.com/autonomousvision/stylegan_xl' target='_blank'>StyleGAN XL</a> is a general purpose StyleGAN, and it is CLIP Guidance notebook was created by <a href='https://github.com/CasualGANPapers/StyleGANXL-CLIP' target='_blank'>ryudrigo and ouhenio</a>, and optimised by <a href='https://twitter.com/rivershavewings' target='_blank'>Katherine Crowson</a> This Spaces Gradio UI to the model was assembled by <a style='color: rgb(99, 102, 241);font-weight:bold' href='https://twitter.com/multimodalart' target='_blank'>@multimodalart</a>, keep up with the <a style='color: rgb(99, 102, 241);' href='https://multimodal.art/news' target='_blank'>latest multimodal ai art news here</a> and consider <a style='color: rgb(99, 102, 241);' href='https://www.patreon.com/multimodalart' target='_blank'>supporting us on Patreon</a></div>",
    article="<h4 style='font-size: 110%;margin-top:.5em'>Biases acknowledgment</h4><div>Despite how impressive being able to turn text into image is, beware to the fact that this model may output content that reinforces or exarcbates societal biases. According to the <a href='https://arxiv.org/abs/2112.10752' target='_blank'>Latent Diffusion paper</a>:<i> \"Deep learning modules tend to reproduce or exacerbate biases that are already present in the data\"</i>. The models are meant to be used for research purposes, such as this one.</div><h4 style='font-size: 110%;margin-top:1em'>Who owns the images produced by this demo?</h4><div>Definetly not me! Probably you do. I say probably because the Copyright discussion about AI generated art is ongoing. So <a href='https://www.theverge.com/2022/2/21/22944335/us-copyright-office-reject-ai-generated-art-recent-entrance-to-paradise' target='_blank'>it may be the case that everything produced here falls automatically into the public domain</a>. But in any case it is either yours or is in the public domain.</div>")
iface.launch(enable_queue=True)