Spaces:
Runtime error
Runtime error
| import os | |
| os.system("pip install --upgrade torch==1.9.1+cu111 torchvision==0.10.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html") | |
| os.system("git clone https://github.com/NVlabs/stylegan3") | |
| os.system("git clone https://github.com/openai/CLIP") | |
| os.system("pip install -e ./CLIP") | |
| os.system("pip install einops ninja scipy numpy Pillow tqdm") | |
| import sys | |
| sys.path.append('./CLIP') | |
| sys.path.append('./stylegan3') | |
| import io | |
| import os, time | |
| 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 | |
| from tqdm.notebook import tqdm | |
| from torchvision.transforms import Compose, Resize, ToTensor, Normalize | |
| from einops import rearrange | |
| device = torch.device('cuda:0') | |
| 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): | |
| basename = os.path.basename(url_or_path) | |
| if os.path.exists(basename): | |
| return basename | |
| else: | |
| os.system("wget -c '{url_or_path}'") | |
| return basename | |
| 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) | |
| 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]) | |
| 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() | |
| # Load stylegan model | |
| base_url = "https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/" | |
| model_name = "stylegan3-t-ffhqu-1024x1024.pkl" | |
| #model_name = "stylegan3-r-metfacesu-1024x1024.pkl" | |
| #model_name = "stylegan3-t-afhqv2-512x512.pkl" | |
| network_url = base_url + model_name | |
| os.system("wget -c https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhqu-1024x1024.pkl") | |
| with open('stylegan3-t-ffhqu-1024x1024.pkl', 'rb') as fp: | |
| G = pickle.load(fp)['G_ema'].to(device) | |
| zs = torch.randn([10000, G.mapping.z_dim], device=device) | |
| w_stds = G.mapping(zs, None).std(0) | |
| def inference(text): | |
| target = clip_model.embed_text(text) | |
| steps = 600 | |
| seed = 2 | |
| tf = Compose([ | |
| Resize(224), | |
| lambda x: torch.clamp((x+1)/2,min=0,max=1), | |
| ]) | |
| torch.manual_seed(seed) | |
| timestring = time.strftime('%Y%m%d%H%M%S') | |
| with torch.no_grad(): | |
| qs = [] | |
| losses = [] | |
| for _ in range(8): | |
| q = (G.mapping(torch.randn([4,G.mapping.z_dim], device=device), None, truncation_psi=0.7) - G.mapping.w_avg) / w_stds | |
| images = G.synthesis(q * w_stds + G.mapping.w_avg) | |
| embeds = embed_image(images.add(1).div(2)) | |
| loss = spherical_dist_loss(embeds, target).mean(0) | |
| i = torch.argmin(loss) | |
| qs.append(q[i]) | |
| losses.append(loss[i]) | |
| qs = torch.stack(qs) | |
| losses = torch.stack(losses) | |
| print(losses) | |
| print(losses.shape, qs.shape) | |
| i = torch.argmin(losses) | |
| q = qs[i].unsqueeze(0) | |
| q.requires_grad_() | |
| q_ema = q | |
| opt = torch.optim.AdamW([q], lr=0.03, betas=(0.0,0.999)) | |
| loop = tqdm(range(steps)) | |
| for i in loop: | |
| opt.zero_grad() | |
| w = q * w_stds | |
| image = G.synthesis(w + G.mapping.w_avg, noise_mode='const') | |
| embed = embed_image(image.add(1).div(2)) | |
| loss = spherical_dist_loss(embed, target).mean() | |
| loss.backward() | |
| opt.step() | |
| loop.set_postfix(loss=loss.item(), q_magnitude=q.std().item()) | |
| q_ema = q_ema * 0.9 + q * 0.1 | |
| image = G.synthesis(q_ema * w_stds + G.mapping.w_avg, noise_mode='const') | |
| if i % 10 == 0: | |
| display(TF.to_pil_image(tf(image)[0])) | |
| pil_image = TF.to_pil_image(image[0].add(1).div(2).clamp(0,1)) | |
| #os.makedirs(f'samples/{timestring}', exist_ok=True) | |
| #pil_image.save(f'samples/{timestring}/{i:04}.jpg') | |
| return pil_image | |
| title = "StyleGAN+CLIP_with_Latent_Bootstraping" | |
| description = "Gradio demo for StyleGAN+CLIP_with_Latent_Bootstraping. To use it, simply add your text, or click one of the examples to load them. Read more at the links below." | |
| article = "<p style='text-align: center'>colab by https://twitter.com/EricHallahan <a href='https://colab.research.google.com/drive/1br7GP_D6XCgulxPTAFhwGaV-ijFe084X' target='_blank'>Colab</a></p>" | |
| examples = [['elon musk']] | |
| gr.Interface( | |
| inference, | |
| "text", | |
| gr.outputs.Image(type="pil", label="Output"), | |
| title=title, | |
| description=description, | |
| article=article, | |
| enable_queue=True, | |
| examples=examples | |
| ).launch(debug=True) | |