File size: 13,298 Bytes
b18ff48
ddad699
 
b18ff48
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a9dfd5a
ddad699
 
56049a8
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4a368e8
b56dcdb
4a368e8
ddad699
 
 
 
 
 
 
 
 
 
 
 
b18ff48
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56049a8
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d69d28e
ddad699
 
 
 
22e38c7
b56dcdb
 
7c17938
 
ddad699
 
 
 
 
 
 
22e38c7
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7c17938
 
 
 
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
56049a8
7c17938
ddad699
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7c17938
ddad699
ce39e0b
ddad699
 
7c17938
ddad699
 
 
 
 
 
 
 
ce39e0b
ddad699
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
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)