File size: 13,300 Bytes
d4fe1e6
 
 
 
 
 
 
 
 
2d78361
 
d4fe1e6
85c6a20
2d78361
d4fe1e6
 
 
 
 
 
 
5d29bbd
 
 
 
 
eb601c1
8d6462e
 
 
 
d4fe1e6
 
 
 
8712735
241cc6f
8712735
d4fe1e6
8d6462e
 
 
fe5b0a5
241cc6f
8d6462e
d4fe1e6
eb601c1
d4fe1e6
 
eb601c1
d4fe1e6
 
 
 
8712735
 
d4fe1e6
 
 
 
 
eb601c1
d4fe1e6
 
 
 
8712735
 
d4fe1e6
 
eb601c1
d4fe1e6
 
eb601c1
d4fe1e6
 
 
eb601c1
1955fd3
 
 
 
eb601c1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d4fe1e6
eb601c1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d4fe1e6
eb601c1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d4fe1e6
5d29bbd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
bc31fa8
5d29bbd
 
 
eb601c1
5d29bbd
 
 
 
 
eb601c1
b97fb24
 
bc31fa8
5d29bbd
eb601c1
1955fd3
eb601c1
 
 
5d29bbd
eb601c1
1955fd3
 
 
5d29bbd
 
 
c97c9f9
5d29bbd
 
 
 
 
 
 
eb601c1
5d29bbd
bc31fa8
 
 
 
 
38930b8
5d29bbd
eb601c1
5d29bbd
 
 
 
 
 
eb601c1
5d29bbd
 
 
eb601c1
 
5219f50
3f591a2
eb601c1
5d29bbd
 
 
1955fd3
8712735
1955fd3
8d6462e
 
 
1955fd3
c4bd367
 
241cc6f
 
 
 
c4bd367
 
ff2dfb3
241cc6f
 
 
c4bd367
 
ff2dfb3
241cc6f
 
 
c4bd367
5d29bbd
eb601c1
3f1a935
ff2dfb3
3f1a935
ff2dfb3
 
 
 
63a03db
ff2dfb3
 
 
 
 
 
1955fd3
 
2442693
d4fe1e6
 
c52d702
1955fd3
 
 
 
 
 
c4bd367
1955fd3
 
 
eb601c1
c52d702
8d6462e
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
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
# -*- coding: utf-8 -*-
"""Copy of compose_glide.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/19xx6Nu4FeiGj-TzTUFxBf-15IkeuFx_F
"""

# from PIL import Image
# from IPython.display import display
import torch as th
import numpy as np

from glide_text2im.download import load_checkpoint
from glide_text2im.model_creation import (
    create_model_and_diffusion,
    model_and_diffusion_defaults,
    model_and_diffusion_defaults_upsampler
)

from composable_diffusion.download import download_model
from composable_diffusion.model_creation import create_model_and_diffusion as create_model_and_diffusion_for_clevr
from composable_diffusion.model_creation import model_and_diffusion_defaults as model_and_diffusion_defaults_for_clevr


from PIL import Image

from torch import autocast
from diffusers import StableDiffusionPipeline

# This notebook supports both CPU and GPU.
# On CPU, generating one sample may take on the order of 20 minutes.
# On a GPU, it should be under a minute.

has_cuda = False
device = th.device('cpu' if not th.cuda.is_available() else 'cuda')
cpu = th.device('cpu')

# iniatilize stable diffusion model
pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    use_auth_token='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'
).to(cpu)

# Create base model.
timestep_respacing = 100  # @param{type: 'number'}
options = model_and_diffusion_defaults()
options['use_fp16'] = has_cuda
options['timestep_respacing'] = str(timestep_respacing)  # use 100 diffusion steps for fast sampling
model, diffusion = create_model_and_diffusion(**options)
model.eval()
if has_cuda:
    model.convert_to_fp16()
model.to(cpu)
model.load_state_dict(load_checkpoint('base', cpu))
print('total base parameters', sum(x.numel() for x in model.parameters()))

# Create upsampler model.
options_up = model_and_diffusion_defaults_upsampler()
options_up['use_fp16'] = has_cuda
options_up['timestep_respacing'] = 'fast27'  # use 27 diffusion steps for very fast sampling
model_up, diffusion_up = create_model_and_diffusion(**options_up)
model_up.eval()
if has_cuda:
    model_up.convert_to_fp16()
model_up.to(cpu)
model_up.load_state_dict(load_checkpoint('upsample', cpu))
print('total upsampler parameters', sum(x.numel() for x in model_up.parameters()))


def show_images(batch: th.Tensor):
    """ Display a batch of images inline. """
    scaled = ((batch + 1) * 127.5).round().clamp(0, 255).to(th.uint8).cpu()
    reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])
    display(Image.fromarray(reshaped.numpy()))


def compose_language_descriptions(prompt, guidance_scale, steps):
    options['timestep_respacing'] = str(steps)
    _, diffusion = create_model_and_diffusion(**options)

    # @markdown `prompt`: when composing  multiple sentences, using `|` as the delimiter.
    prompts = [x.strip() for x in prompt.split('|')]

    batch_size = 1
    # Tune this parameter to control the sharpness of 256x256 images.
    # A value of 1.0 is sharper, but sometimes results in grainy artifacts.
    upsample_temp = 0.980  # @param{type: 'number'}

    masks = [True] * len(prompts) + [False]
    # coefficients = th.tensor([0.5, 0.5], device=device).reshape(-1, 1, 1, 1)
    masks = th.tensor(masks, dtype=th.bool, device=device)

    # sampling function
    def model_fn(x_t, ts, **kwargs):
        half = x_t[:1]
        combined = th.cat([half] * x_t.size(0), dim=0)
        model_out = model(combined, ts, **kwargs)
        eps, rest = model_out[:, :3], model_out[:, 3:]
        cond_eps = eps[masks].mean(dim=0, keepdim=True)
        # cond_eps = (coefficients * eps[masks]).sum(dim=0)[None]
        uncond_eps = eps[~masks].mean(dim=0, keepdim=True)
        half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
        eps = th.cat([half_eps] * x_t.size(0), dim=0)
        return th.cat([eps, rest], dim=1)

    ##############################
    # Sample from the base model #
    ##############################

    # Create the text tokens to feed to the model.
    def sample_64(prompts):
        tokens_list = [model.tokenizer.encode(prompt) for prompt in prompts]
        outputs = [model.tokenizer.padded_tokens_and_mask(
            tokens, options['text_ctx']
        ) for tokens in tokens_list]

        cond_tokens, cond_masks = zip(*outputs)
        cond_tokens, cond_masks = list(cond_tokens), list(cond_masks)

        full_batch_size = batch_size * (len(prompts) + 1)
        uncond_tokens, uncond_mask = model.tokenizer.padded_tokens_and_mask(
            [], options['text_ctx']
        )

        # Pack the tokens together into model kwargs.
        model_kwargs = dict(
            tokens=th.tensor(
                cond_tokens + [uncond_tokens], device=device
            ),
            mask=th.tensor(
                cond_masks + [uncond_mask],
                dtype=th.bool,
                device=device,
            ),
        )

        # Sample from the base model.
        model.del_cache()
        samples = diffusion.p_sample_loop(
            model_fn,
            (full_batch_size, 3, options["image_size"], options["image_size"]),
            device=device,
            clip_denoised=True,
            progress=True,
            model_kwargs=model_kwargs,
            cond_fn=None,
        )[:batch_size]
        model.del_cache()

        # Show the output
        return samples

    ##############################
    # Upsample the 64x64 samples #
    ##############################

    def upsampling_256(prompts, samples):
        tokens = model_up.tokenizer.encode("".join(prompts))
        tokens, mask = model_up.tokenizer.padded_tokens_and_mask(
            tokens, options_up['text_ctx']
        )

        # Create the model conditioning dict.
        model_kwargs = dict(
            # Low-res image to upsample.
            low_res=((samples + 1) * 127.5).round() / 127.5 - 1,

            # Text tokens
            tokens=th.tensor(
                [tokens] * batch_size, device=device
            ),
            mask=th.tensor(
                [mask] * batch_size,
                dtype=th.bool,
                device=device,
            ),
        )

        # Sample from the base model.
        model_up.del_cache()
        up_shape = (batch_size, 3, options_up["image_size"], options_up["image_size"])
        up_samples = diffusion_up.ddim_sample_loop(
            model_up,
            up_shape,
            noise=th.randn(up_shape, device=device) * upsample_temp,
            device=device,
            clip_denoised=True,
            progress=True,
            model_kwargs=model_kwargs,
            cond_fn=None,
        )[:batch_size]
        model_up.del_cache()

        # Show the output
        return up_samples

    # sampling 64x64 images
    samples = sample_64(prompts)
    # show_images(samples)

    # upsample from 64x64 to 256x256
    upsamples = upsampling_256(prompts, samples)
    # show_images(upsamples)

    out_img = upsamples[0].permute(1, 2, 0)
    out_img = (out_img + 1) / 2
    out_img = (out_img.detach().cpu() * 255.).to(th.uint8)
    out_img = out_img.numpy()
    return out_img


# create model for CLEVR Objects
clevr_options = model_and_diffusion_defaults_for_clevr()

flags = {
    "image_size": 128,
    "num_channels": 192,
    "num_res_blocks": 2,
    "learn_sigma": True,
    "use_scale_shift_norm": False,
    "raw_unet": True,
    "noise_schedule": "squaredcos_cap_v2",
    "rescale_learned_sigmas": False,
    "rescale_timesteps": False,
    "num_classes": '2',
    "dataset": "clevr_pos",
    "use_fp16": has_cuda,
    "timestep_respacing": '100'
}

for key, val in flags.items():
    clevr_options[key] = val

clevr_model, clevr_diffusion = create_model_and_diffusion_for_clevr(**clevr_options)
clevr_model.eval()
if has_cuda:
    clevr_model.convert_to_fp16()

clevr_model.to(th.device('cpu'))
clevr_model.load_state_dict(th.load(download_model('clevr_pos'), th.device('cpu')))
print('total clevr_pos parameters', sum(x.numel() for x in clevr_model.parameters()))


def compose_clevr_objects(prompt, guidance_scale, steps):
    coordinates = [[float(x.split(',')[0].strip()), float(x.split(',')[1].strip())]
                   for x in prompt.split('|')]
    coordinates += [[-1, -1]]  # add unconditional score label
    batch_size = 1

    clevr_options['timestep_respacing'] = str(int(steps))
    _, clevr_diffusion = create_model_and_diffusion_for_clevr(**clevr_options)

    def model_fn(x_t, ts, **kwargs):
        half = x_t[:1]
        combined = th.cat([half] * kwargs['y'].size(0), dim=0)
        model_out = clevr_model(combined, ts, **kwargs)
        eps, rest = model_out[:, :3], model_out[:, 3:]
        masks = kwargs.get('masks')
        cond_eps = eps[masks].mean(dim=0, keepdim=True)
        uncond_eps = eps[~masks].mean(dim=0, keepdim=True)
        half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)
        eps = th.cat([half_eps] * x_t.size(0), dim=0)
        return th.cat([eps, rest], dim=1)

    def sample(coordinates):
        masks = [True] * (len(coordinates) - 1) + [False]
        model_kwargs = dict(
            y=th.tensor(coordinates, dtype=th.float, device=device),
            masks=th.tensor(masks, dtype=th.bool, device=device)
        )
        samples = clevr_diffusion.p_sample_loop(
            model_fn,
            (len(coordinates), 3, clevr_options["image_size"], clevr_options["image_size"]),
            device=device,
            clip_denoised=True,
            progress=True,
            model_kwargs=model_kwargs,
            cond_fn=None,
        )[:batch_size]

        return samples

    samples = sample(coordinates)
    out_img = samples[0].permute(1, 2, 0)
    out_img = (out_img + 1) / 2
    out_img = (out_img.detach().cpu() * 255.).to(th.uint8)
    out_img = out_img.numpy()

    return out_img


def stable_diffusion_compose(prompt, scale, steps):
    with autocast('cpu' if not th.cuda.is_available() else 'cuda'):
        image = pipe(prompt, guidance_scale=scale, num_inference_steps=steps)["sample"][0]
        return image


def compose(prompt, version, guidance_scale, steps):
    with th.no_grad():
        if version == 'GLIDE':
            clevr_model.to(cpu)
            pipe.to(cpu)
            model.to(device)
            model_up.to(device)
            return compose_language_descriptions(prompt, guidance_scale, steps)
        elif version == 'Stable_Diffusion_1v_4':
            clevr_model.to(cpu)
            model.to(cpu)
            model_up.to(cpu)
            pipe.to(device)
            return stable_diffusion_compose(prompt, guidance_scale, steps)
        else:
            pipe.to(cpu)
            model.to(cpu)
            model_up.to(cpu)
            clevr_model.to(device)
            return compose_clevr_objects(prompt, guidance_scale, steps)


examples_1 = 'a camel | a forest'
examples_2 = 'A blue sky  | A mountain in the horizon | Cherry Blossoms in front of the mountain'
examples_3 = '0.1, 0.5 | 0.3, 0.5 | 0.5, 0.5 | 0.7, 0.5 | 0.9, 0.5'
examples_4 = 'red trees | a blue house'
examples_5 = 'a white church | lightning in the background'
examples_6 = 'a camel | arctic'
examples_7 = 'A lake | A mountain  | Cherry Blossoms next to the lake'
examples = [
            [examples_7, 'Stable_Diffusion_1v_4', 15, 50],
            [examples_5, 'Stable_Diffusion_1v_4', 15, 50],
            [examples_4, 'Stable_Diffusion_1v_4', 20, 50],
            [examples_6, 'Stable_Diffusion_1v_4', 15, 50],
            [examples_1, 'GLIDE', 15, 100],
            [examples_2, 'GLIDE', 15, 100],
            [examples_3, 'CLEVR Objects', 10, 100]
]

import gradio as gr

title = 'Compositional Visual Generation with Composable Diffusion Models'
description = '<p>Demo for Composable Diffusion<ul><li>~30s per GLIDE/Stable-Diffusion example</li><li>~10s per CLEVR Object example</li>(<b>Note</b>: time is varied depending on what gpu is used.)</ul></p><p>See more information from our <a href="https://energy-based-model.github.io/Compositional-Visual-Generation-with-Composable-Diffusion-Models/">Project Page</a>.</p><ul><li>One version is based on the released <a href="https://github.com/openai/glide-text2im">GLIDE</a> and <a href="https://github.com/CompVis/stable-diffusion/">Stable Diffusion</a> for composing natural language description.</li><li>Another is based on our pre-trained CLEVR Object Model for composing objects. <br>(<b>Note</b>: We recommend using <b><i>x</i></b> in range <b><i>[0.1, 0.9]</i></b> and <b><i>y</i></b> in range <b><i>[0.25, 0.7]</i></b>, since the training dataset labels are in given ranges.)</li></ul><p>When composing  multiple sentences, use `|` as the delimiter, see given examples below.</p><p><b>Note</b>: When using more steps, the results can improve.</p>'

iface = gr.Interface(compose,
                     inputs=[
                         "text",
                         gr.Radio(['Stable_Diffusion_1v_4', 'GLIDE', 'CLEVR Objects'], type="value", label='version'),
                         gr.Slider(2, 30),
                         gr.Slider(10, 200)
                     ],
                     outputs='image',
                     title=title, description=description, examples=examples)

iface.launch()