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="
By typing a prompt and pressing submit you generate images based on it. StyleGAN XL is a general purpose StyleGAN, and it is CLIP Guidance notebook was created by ryudrigo and ouhenio, and optimised by Katherine Crowson This Spaces Gradio UI to the model was assembled by @multimodalart, keep up with the latest multimodal ai art news here and consider supporting us on Patreon
", article="

Biases acknowledgment

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 Latent Diffusion paper: \"Deep learning modules tend to reproduce or exacerbate biases that are already present in the data\". The models are meant to be used for research purposes, such as this one.

Who owns the images produced by this demo?

Definetly not me! Probably you do. I say probably because the Copyright discussion about AI generated art is ongoing. So it may be the case that everything produced here falls automatically into the public domain. But in any case it is either yours or is in the public domain.
") iface.launch(enable_queue=True)