app.py

import math
import fire
import gradio as gr
import numpy as np
import rich
import torch
from contextlib import nullcontext
from einops import rearrange
from functools import partial
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.util import load_and_preprocess, instantiate_from_config
from omegaconf import OmegaConf
from PIL import Image
from rich import print
from torch import autocast
from torchvision import transforms


_SHOW_INTERMEDIATE = True
_GPU_INDEX = 0
# _GPU_INDEX = 2


def load_model_from_config(config, ckpt, device, verbose=False):
    print(f'Loading model from {ckpt}')
    pl_sd = torch.load(ckpt, map_location=device)
    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.to(device)
    model.eval()
    return model


@torch.no_grad()
def sample_model(input_im, model, sampler, precision, h, w, ddim_steps, n_samples, scale,
                 ddim_eta, x, y, z):
    precision_scope = autocast if precision == 'autocast' else nullcontext
    with precision_scope('cuda'):
        with model.ema_scope():
            c = model.get_learned_conditioning(input_im).tile(n_samples, 1, 1)
            T = torch.tensor([math.radians(x), math.sin(
                math.radians(y)), math.cos(math.radians(y)), z])
            T = T[None, None, :].repeat(n_samples, 1, 1).to(c.device)
            c = torch.cat([c, T], dim=-1)
            c = model.cc_projection(c)
            cond = {}
            cond['c_crossattn'] = [c]
            c_concat = model.encode_first_stage((input_im.to(c.device))).mode().detach()
            cond['c_concat'] = [model.encode_first_stage((input_im.to(c.device))).mode().detach()
                                .repeat(n_samples, 1, 1, 1)]
            if scale != 1.0:
                uc = {}
                uc['c_concat'] = [torch.zeros(n_samples, 4, h // 8, w // 8).to(c.device)]
                uc['c_crossattn'] = [torch.zeros_like(c).to(c.device)]
            else:
                uc = None

            shape = [4, h // 8, w // 8]
            samples_ddim, _ = sampler.sample(S=ddim_steps,
                                             conditioning=cond,
                                             batch_size=n_samples,
                                             shape=shape,
                                             verbose=False,
                                             unconditional_guidance_scale=scale,
                                             unconditional_conditioning=uc,
                                             eta=ddim_eta,
                                             x_T=None)
            print(samples_ddim.shape)
            # samples_ddim = torch.nn.functional.interpolate(samples_ddim, 64, mode='nearest', antialias=False)
            x_samples_ddim = model.decode_first_stage(samples_ddim)
            return torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0).cpu()


def main(
    model,
    device,
    input_im,
    preprocess=True,
    x=0.,
    y=0.,
    z=0.,
    scale=3.0,
    n_samples=4,
    ddim_steps=50,
    ddim_eta=1.0,
    precision='fp32',
    h=256,
    w=256,
):
    # input_im[input_im == [0., 0., 0.]] = [1., 1., 1., 1.]
    print('old input_im:', input_im.size)

    if preprocess:
        input_im = load_and_preprocess(input_im)
        input_im = (input_im / 255.0).astype(np.float32)
        # (H, W, 3) array in [0, 1].

    else:
        input_im = input_im.resize([256, 256], Image.Resampling.LANCZOS)
        input_im = np.asarray(input_im, dtype=np.float32) / 255.0
        # (H, W, 4) array in [0, 1].
        
        # old method: very important, thresholding background
        # input_im[input_im[:, :, -1] <= 0.9] = [1., 1., 1., 1.]

        # new method: apply correct method of compositing to avoid sudden transitions / thresholding
        # (smoothly transition foreground to white background based on alpha values)
        alpha = input_im[:, :, 3:4]
        white_im = np.ones_like(input_im)
        input_im = alpha * input_im + (1.0 - alpha) * white_im
        
        input_im = input_im[:, :, 0:3]
        # (H, W, 3) array in [0, 1].

    print('new input_im:', input_im.shape, input_im.dtype, input_im.min(), input_im.max())
    show_in_im = Image.fromarray((input_im * 255).astype(np.uint8))

    input_im = transforms.ToTensor()(input_im).unsqueeze(0).to(device)
    input_im = input_im * 2 - 1
    input_im = transforms.functional.resize(input_im, [h, w])

    sampler = DDIMSampler(model)
    x_samples_ddim = sample_model(input_im, model, sampler, precision, h, w,
                                  ddim_steps, n_samples, scale, ddim_eta, x, y, z)

    output_ims = []
    for x_sample in x_samples_ddim:
        x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
        output_ims.append(Image.fromarray(x_sample.astype(np.uint8)))
    
    if _SHOW_INTERMEDIATE:
        return (output_ims, show_in_im)
    else:
        return output_ims


description = '''
Generate novel viewpoints of an object depicted in one input image using a fine-tuned version of Stable Diffusion.
'''

article = '''
## How to use this?
TBD
## How does this work?
TBD
'''


def run_demo(
    device_idx=_GPU_INDEX,
    ckpt='last.ckpt',
    config='configs/sd-objaverse-finetune-c_concat-256.yaml',
):

    device = f'cuda:{device_idx}'
    config = OmegaConf.load(config)
    model = load_model_from_config(config, ckpt, device=device)

    inputs = [
        gr.Image(type='pil', image_mode='RGBA'),  # shape=[512, 512]
        gr.Checkbox(True, label='Preprocess image (remove background and center)',
        info='If enabled, the uploaded image will be preprocessed to remove the background and center the object by cropping and/or padding as necessary. '
        'If disabled, the image will be used as-is, *BUT* a fully transparent or white background is required.'),
        # gr.Number(label='polar (between axis z+)'),
        # gr.Number(label='azimuth (between axis x+)'),
        # gr.Number(label='z (distance from center)'),
        gr.Slider(-90, 90, value=0, step=5, label='Polar angle (vertical rotation in degrees)',
        info='Positive values move the camera down, while negative values move the camera up.'),
        gr.Slider(-90, 90, value=0, step=5, label='Azimuth angle (horizontal rotation in degrees)',
        info='Positive values move the camera right, while negative values move the camera left.'),
        gr.Slider(-2, 2, value=0, step=0.5, label='Radius (distance from center)',
        info='Positive values move the camera further away, while negative values move the camera closer.'),
        gr.Slider(0, 30, value=3, step=1, label='cfg scale'),
        gr.Slider(1, 8, value=4, step=1, label='Number of samples to generate'),
        gr.Slider(5, 200, value=100, step=5, label='Number of steps'),
    ]
    output = [gr.Gallery(label='Generated images from specified new viewpoint')]
    output[0].style(grid=2)
    
    if _SHOW_INTERMEDIATE:
        output += [gr.Image(type='pil', image_mode='RGB', label='Preprocessed input image')]

    fn_with_model = partial(main, model, device)
    fn_with_model.__name__ = 'fn_with_model'

    examples = [
        # ['assets/zero-shot/bear.png', 0, 0, 0, 3, 4, 100],
        # ['assets/zero-shot/car.png', 0, 0, 0, 3, 4, 100],
        # ['assets/zero-shot/elephant.png', 0, 0, 0, 3, 4, 100],
        # ['assets/zero-shot/pikachu.png', 0, 0, 0, 3, 4, 100],
        # ['assets/zero-shot/spyro.png', 0, 0, 0, 3, 4, 100],
        # ['assets/zero-shot/taxi.png', 0, 0, 0, 3, 4, 100],
    ]

    demo = gr.Interface(
        fn=fn_with_model,
        title='Demo for Zero-Shot Control of Camera Viewpoints within a Single Image',
        description=description,
        article=article,
        inputs=inputs,
        outputs=output,
        examples=examples,
        allow_flagging='never',
    )
    demo.launch(enable_queue=True, share=True)


if __name__ == '__main__':
    fire.Fire(run_demo)