model.py

import argparse, os, sys, glob
import torch
import PIL
import numpy as np
from omegaconf import OmegaConf
from PIL import Image
from tqdm import tqdm, trange
from itertools import islice
from einops import rearrange, repeat
from torchvision.utils import make_grid
import time
from pytorch_lightning import seed_everything
from torch import autocast
from contextlib import contextmanager, nullcontext

from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler


def chunk(it, size):
    it = iter(it)
    return iter(lambda: tuple(islice(it, size)), ())

def torch_gc():
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()


def load_model_from_config(config, ckpt, verbose=False):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu")
    if "global_step" in pl_sd:
        print(f"Global Step: {pl_sd['global_step']}")
    sd = pl_sd["state_dict"]
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    if len(m) > 0 and verbose:
        print("missing keys:")
        print(m)
    if len(u) > 0 and verbose:
        print("unexpected keys:")
        print(u)

    model.cuda()
    model.half()
    model.eval()
    return model

def load_img(image, W, H):
    w, h = image.size
    print(f"loaded input image of size ({w}, {h})")
    image = image.resize((int(W), int(H)), resample=PIL.Image.LANCZOS)
    print(f"resize input image to size ({W}, {H})")
    image = np.array(image).astype(np.float32) / 255.0
    image = image[None].transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return 2.*image - 1.

class AppModel():
    def __init__(self,):
        self.config = OmegaConf.load("configs/stable-diffusion/v1-inference.yaml")
        self.model = load_model_from_config(self.config, "models/ldm/stable-diffusion-v1/model.ckpt")

        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        self.device = device
        self.model = self.model.to(device)

        self.sampler = PLMSSampler(self.model)
        self.img_sampler = DDIMSampler(self.model)

        self.C = 4 # latent channels
        self.f = 8 # downsampling factors

    def run_with_prompt(self, seed, prompt, n_samples, W, H, scale, ddim_steps,  strength=0., init_img=None):
        torch_gc()
        seed_everything(seed)
        ddim_eta=0.0

        assert prompt is not None
        print(f"Prompt: {prompt}")
        batch_size = n_samples
        data = [batch_size * [prompt]]

        start_code = None

        n_rows = int(n_samples**0.5)

        precision_scope = autocast 
        if init_img is None:
            with torch.no_grad():
                with precision_scope(device_type='cuda', dtype=torch.float16):
                    with self.model.ema_scope():
                        all_samples = list()
                        for prompts in tqdm(data, desc="data"):
                            torch_gc()
                            uc = None
                            if scale != 1.0:
                                uc = self.model.get_learned_conditioning(batch_size * [""])
                            if isinstance(prompts, tuple):
                                prompts = list(prompts)
                            c = self.model.get_learned_conditioning(prompts)
                            shape = [self.C, H // self.f, W // self.f]
                            samples_ddim, _ = self.sampler.sample(S=ddim_steps,
                                                            conditioning=c,
                                                            batch_size=n_samples,
                                                            shape=shape,
                                                            verbose=False,
                                                            unconditional_guidance_scale=scale,
                                                            unconditional_conditioning=uc,
                                                            eta=ddim_eta,
                                                            x_T=start_code)

                            x_samples_ddim = self.model.decode_first_stage(samples_ddim)
                            x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)

                            for x_sample in x_samples_ddim:
                                x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                                image = Image.fromarray(x_sample.astype(np.uint8))
                                all_samples.append(image)

                        # additionally, grid image
                        grid = torch.stack([x_samples_ddim], 0)
                        grid = rearrange(grid, 'n b c h w -> (n b) c h w')
                        grid = make_grid(grid, nrow=n_rows)

                        # to image
                        grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
                        grid = grid.astype(np.uint8)
                        torch_gc()

                        return grid, all_samples
        else:
            init_image = load_img(init_img, W, H).to(self.device)
            init_image = repeat(init_image, '1 ... -> b ...', b=batch_size)
            torch_gc()
            with precision_scope(device_type='cuda', dtype=torch.float16):
                init_latent = self.model.get_first_stage_encoding(self.model.encode_first_stage(init_image))  # move to latent space
            torch_gc()

            sampler = self.img_sampler

            sampler.make_schedule(ddim_num_steps=ddim_steps, ddim_eta=ddim_eta, verbose=False)

            assert 0. <= strength < 1., 'can only work with strength in [0.0, 1.0)'
            t_enc = int(strength * ddim_steps)
            print(f"target t_enc is {t_enc} steps")

            with torch.no_grad():
                with precision_scope(device_type='cuda', dtype=torch.float16):
                    with self.model.ema_scope():
                        all_samples = list()
                        for prompts in tqdm(data, desc="data"):
                            uc = None
                            if scale != 1.0:
                                uc = self.model.get_learned_conditioning(batch_size * [""])
                            if isinstance(prompts, tuple):
                                prompts = list(prompts)
                            c = self.model.get_learned_conditioning(prompts)

                            # encode (scaled latent)
                            z_enc = sampler.stochastic_encode(init_latent, torch.tensor([t_enc]*batch_size).to(self.device))
                            # decode it
                            samples = sampler.decode(z_enc, c, t_enc, unconditional_guidance_scale=scale,
                                                    unconditional_conditioning=uc,)

                            x_samples = self.model.decode_first_stage(samples)
                            x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)

                            for x_sample in x_samples:
                                x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
                                image = Image.fromarray(x_sample.astype(np.uint8))
                                all_samples.append(image)

                        # additionally, save as grid
                        grid = torch.stack([x_samples], 0)
                        grid = rearrange(grid, 'n b c h w -> (n b) c h w')
                        grid = make_grid(grid, nrow=n_rows)

                        # to image
                        grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
                        grid = grid.astype(np.uint8)
                        torch_gc()
                        return grid, all_samples
            



if __name__ == "__main__":
    main()