valhalla/glide-text2im

1import base642from io import BytesIO3from fastapi import FastAPI4 5from PIL import Image6import torch as th7 8from glide_text2im.download import load_checkpoint9from glide_text2im.model_creation import (10    create_model_and_diffusion,11    model_and_diffusion_defaults,12    model_and_diffusion_defaults_upsampler13)14 15print("Loading models...")16app = FastAPI()17 18# This notebook supports both CPU and GPU.19# On CPU, generating one sample may take on the order of 20 minutes.20# On a GPU, it should be under a minute.21 22has_cuda = th.cuda.is_available()23device = th.device('cpu' if not has_cuda else 'cuda')24 25# Create base model.26options = model_and_diffusion_defaults()27options['use_fp16'] = has_cuda28options['timestep_respacing'] = '100' # use 100 diffusion steps for fast sampling29model, diffusion = create_model_and_diffusion(**options)30model.eval()31if has_cuda:32    model.convert_to_fp16()33model.to(device)34model.load_state_dict(load_checkpoint('base', device))35print('total base parameters', sum(x.numel() for x in model.parameters()))36 37# Create upsampler model.38options_up = model_and_diffusion_defaults_upsampler()39options_up['use_fp16'] = has_cuda40options_up['timestep_respacing'] = 'fast27' # use 27 diffusion steps for very fast sampling41model_up, diffusion_up = create_model_and_diffusion(**options_up)42model_up.eval()43if has_cuda:44    model_up.convert_to_fp16()45model_up.to(device)46model_up.load_state_dict(load_checkpoint('upsample', device))47print('total upsampler parameters', sum(x.numel() for x in model_up.parameters()))48 49 50def get_images(batch: th.Tensor):51    """ Display a batch of images inline. """52    scaled = ((batch + 1)*127.5).round().clamp(0,255).to(th.uint8).cpu()53    reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])54    Image.fromarray(reshaped.numpy())55 56 57# Create a classifier-free guidance sampling function58guidance_scale = 3.059 60def model_fn(x_t, ts, **kwargs):61    half = x_t[: len(x_t) // 2]62    combined = th.cat([half, half], dim=0)63    model_out = model(combined, ts, **kwargs)64    eps, rest = model_out[:, :3], model_out[:, 3:]65    cond_eps, uncond_eps = th.split(eps, len(eps) // 2, dim=0)66    half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)67    eps = th.cat([half_eps, half_eps], dim=0)68    return th.cat([eps, rest], dim=1)69 70 71@app.get("/")72def read_root():73    return {"glide!"}74 75@app.get("/{generate}")76def sample(prompt):77    # Sampling parameters78    batch_size = 179 80    # Tune this parameter to control the sharpness of 256x256 images.81    # A value of 1.0 is sharper, but sometimes results in grainy artifacts.82    upsample_temp = 0.99783 84    ##############################85    # Sample from the base model #86    ##############################87 88    # Create the text tokens to feed to the model.89    tokens = model.tokenizer.encode(prompt)90    tokens, mask = model.tokenizer.padded_tokens_and_mask(91        tokens, options['text_ctx']92    )93 94    # Create the classifier-free guidance tokens (empty)95    full_batch_size = batch_size * 296    uncond_tokens, uncond_mask = model.tokenizer.padded_tokens_and_mask(97        [], options['text_ctx']98    )99 100    # Pack the tokens together into model kwargs.101    model_kwargs = dict(102        tokens=th.tensor(103            [tokens] * batch_size + [uncond_tokens] * batch_size, device=device104        ),105        mask=th.tensor(106            [mask] * batch_size + [uncond_mask] * batch_size,107            dtype=th.bool,108            device=device,109        ),110    )111 112    # Sample from the base model.113    model.del_cache()114    samples = diffusion.p_sample_loop(115        model_fn,116        (full_batch_size, 3, options["image_size"], options["image_size"]),117        device=device,118        clip_denoised=True,119        progress=True,120        model_kwargs=model_kwargs,121        cond_fn=None,122    )[:batch_size]123    model.del_cache()124 125 126    ##############################127    # Upsample the 64x64 samples #128    ##############################129 130    tokens = model_up.tokenizer.encode(prompt)131    tokens, mask = model_up.tokenizer.padded_tokens_and_mask(132        tokens, options_up['text_ctx']133    )134 135    # Create the model conditioning dict.136    model_kwargs = dict(137        # Low-res image to upsample.138        low_res=((samples+1)*127.5).round()/127.5 - 1,139 140        # Text tokens141        tokens=th.tensor(142            [tokens] * batch_size, device=device143        ),144        mask=th.tensor(145            [mask] * batch_size,146            dtype=th.bool,147            device=device,148        ),149    )150 151    # Sample from the base model.152    model_up.del_cache()153    up_shape = (batch_size, 3, options_up["image_size"], options_up["image_size"])154    up_samples = diffusion_up.ddim_sample_loop(155        model_up,156        up_shape,157        noise=th.randn(up_shape, device=device) * upsample_temp,158        device=device,159        clip_denoised=True,160        progress=True,161        model_kwargs=model_kwargs,162        cond_fn=None,163    )[:batch_size]164    model_up.del_cache()165 166    # Show the output167    image = get_images(up_samples)168    image = to_base64(image)169    return {"image": image}170 171 172def to_base64(pil_image):173    buffered = BytesIO()174    pil_image.save(buffered, format="JPEG")175    return base64.b64encode(buffered.getvalue())