inversion_utils.py

import torch
from tqdm import tqdm
from typing import List, Optional, Tuple
from models import PipelineWrapper
import gradio as gr


def inversion_forward_process(model: PipelineWrapper,
                              x0: torch.Tensor,
                              etas: Optional[float] = None,
                              prompts: List[str] = [""],
                              cfg_scales: List[float] = [3.5],
                              num_inference_steps: int = 50,
                              numerical_fix: bool = False,
                              duration: Optional[float] = None,
                              first_order: bool = False,
                              save_compute: bool = True,
                              progress=gr.Progress()) -> Tuple:
    if len(prompts) > 1 or prompts[0] != "":
        text_embeddings_hidden_states, text_embeddings_class_labels, \
            text_embeddings_boolean_prompt_mask = model.encode_text(prompts)

        # In the forward negative prompts are not supported currently (TODO)
        uncond_embeddings_hidden_states, uncond_embeddings_class_lables, uncond_boolean_prompt_mask = model.encode_text(
            [""], negative=True, save_compute=save_compute, cond_length=text_embeddings_class_labels.shape[1]
            if text_embeddings_class_labels is not None else None)
    else:
        uncond_embeddings_hidden_states, uncond_embeddings_class_lables, uncond_boolean_prompt_mask = model.encode_text(
            [""], negative=True, save_compute=False)

    timesteps = model.model.scheduler.timesteps.to(model.device)
    variance_noise_shape = model.get_noise_shape(x0, num_inference_steps)

    if type(etas) in [int, float]:
        etas = [etas]*model.model.scheduler.num_inference_steps
    xts = model.sample_xts_from_x0(x0, num_inference_steps=num_inference_steps)
    zs = torch.zeros(size=variance_noise_shape, device=model.device)
    extra_info = [None] * len(zs)

    if timesteps[0].dtype == torch.int64:
        t_to_idx = {int(v): k for k, v in enumerate(timesteps)}
    elif timesteps[0].dtype == torch.float32:
        t_to_idx = {float(v): k for k, v in enumerate(timesteps)}
    xt = x0
    op = tqdm(timesteps, desc="Inverting")
    model.setup_extra_inputs(xt, init_timestep=timesteps[0], audio_end_in_s=duration,
                             save_compute=save_compute and prompts[0] != "")
    app_op = progress.tqdm(timesteps, desc="Inverting")
    for t, _ in zip(op, app_op):
        idx = num_inference_steps - t_to_idx[int(t) if timesteps[0].dtype == torch.int64 else float(t)] - 1

        # 1. predict noise residual
        xt = xts[idx+1][None]
        xt_inp = model.model.scheduler.scale_model_input(xt, t)

        with torch.no_grad():
            if save_compute and prompts[0] != "":
                comb_out, _, _ = model.unet_forward(
                    xt_inp.expand(2, -1, -1, -1) if hasattr(model.model, 'unet') else xt_inp.expand(2, -1, -1),
                    timestep=t,
                    encoder_hidden_states=torch.cat([uncond_embeddings_hidden_states, text_embeddings_hidden_states
                                                     ], dim=0)
                    if uncond_embeddings_hidden_states is not None else None,
                    class_labels=torch.cat([uncond_embeddings_class_lables, text_embeddings_class_labels], dim=0)
                    if uncond_embeddings_class_lables is not None else None,
                    encoder_attention_mask=torch.cat([uncond_boolean_prompt_mask, text_embeddings_boolean_prompt_mask
                                                      ], dim=0)
                    if uncond_boolean_prompt_mask is not None else None,
                )
                out, cond_out = comb_out.sample.chunk(2, dim=0)
            else:
                out = model.unet_forward(xt_inp, timestep=t,
                                         encoder_hidden_states=uncond_embeddings_hidden_states,
                                         class_labels=uncond_embeddings_class_lables,
                                         encoder_attention_mask=uncond_boolean_prompt_mask)[0].sample
                if len(prompts) > 1 or prompts[0] != "":
                    cond_out = model.unet_forward(
                        xt_inp,
                        timestep=t,
                        encoder_hidden_states=text_embeddings_hidden_states,
                        class_labels=text_embeddings_class_labels,
                        encoder_attention_mask=text_embeddings_boolean_prompt_mask)[0].sample

        if len(prompts) > 1 or prompts[0] != "":
            # # classifier free guidance
            noise_pred = out + (cfg_scales[0] * (cond_out - out)).sum(axis=0).unsqueeze(0)
        else:
            noise_pred = out

        # xtm1 =  xts[idx+1][None]
        xtm1 = xts[idx][None]
        z, xtm1, extra = model.get_zs_from_xts(xt, xtm1, noise_pred, t,
                                               eta=etas[idx], numerical_fix=numerical_fix,
                                               first_order=first_order)
        zs[idx] = z
        # print(f"Fix Xt-1 distance -  NORM:{torch.norm(xts[idx] - xtm1):.4g}, MSE:{((xts[idx] - xtm1)**2).mean():.4g}")
        xts[idx] = xtm1
        extra_info[idx] = extra

    if zs is not None:
        # zs[-1] = torch.zeros_like(zs[-1])
        zs[0] = torch.zeros_like(zs[0])
        # zs_cycle[0] = torch.zeros_like(zs[0])

    del app_op.iterables[0]
    return xt, zs, xts, extra_info


def inversion_reverse_process(model: PipelineWrapper,
                              xT: torch.Tensor,
                              tstart: torch.Tensor,
                              etas: float = 0,
                              prompts: List[str] = [""],
                              neg_prompts: List[str] = [""],
                              cfg_scales: Optional[List[float]] = None,
                              zs: Optional[List[torch.Tensor]] = None,
                              duration: Optional[float] = None,
                              first_order: bool = False,
                              extra_info: Optional[List] = None,
                              save_compute: bool = True,
                              progress=gr.Progress()) -> Tuple[torch.Tensor, torch.Tensor]:

    text_embeddings_hidden_states, text_embeddings_class_labels, \
        text_embeddings_boolean_prompt_mask = model.encode_text(prompts)
    uncond_embeddings_hidden_states, uncond_embeddings_class_lables, \
        uncond_boolean_prompt_mask = model.encode_text(neg_prompts,
                                                       negative=True,
                                                       save_compute=save_compute,
                                                       cond_length=text_embeddings_class_labels.shape[1]
                                                       if text_embeddings_class_labels is not None else None)

    xt = xT[tstart.max()].unsqueeze(0)

    if etas is None:
        etas = 0
    if type(etas) in [int, float]:
        etas = [etas]*model.model.scheduler.num_inference_steps
    assert len(etas) == model.model.scheduler.num_inference_steps
    timesteps = model.model.scheduler.timesteps.to(model.device)

    op = tqdm(timesteps[-zs.shape[0]:], desc="Editing")
    if timesteps[0].dtype == torch.int64:
        t_to_idx = {int(v): k for k, v in enumerate(timesteps[-zs.shape[0]:])}
    elif timesteps[0].dtype == torch.float32:
        t_to_idx = {float(v): k for k, v in enumerate(timesteps[-zs.shape[0]:])}
    model.setup_extra_inputs(xt, extra_info=extra_info, init_timestep=timesteps[-zs.shape[0]],
                             audio_end_in_s=duration, save_compute=save_compute)
    app_op = progress.tqdm(timesteps[-zs.shape[0]:], desc="Editing")
    for it, (t, _) in enumerate(zip(op, app_op)):
        idx = model.model.scheduler.num_inference_steps - t_to_idx[
            int(t) if timesteps[0].dtype == torch.int64 else float(t)] - \
                (model.model.scheduler.num_inference_steps - zs.shape[0] + 1)

        xt_inp = model.model.scheduler.scale_model_input(xt, t)

        # # Unconditional embedding
        with torch.no_grad():
            # print(f'xt_inp.shape: {xt_inp.shape}')
            # print(f't.shape: {t.shape}')
            # print(f'uncond_embeddings_hidden_states.shape: {uncond_embeddings_hidden_states.shape}')
            # print(f'uncond_embeddings_class_lables.shape: {uncond_embeddings_class_lables.shape}')
            # print(f'uncond_boolean_prompt_mask.shape: {uncond_boolean_prompt_mask.shape}')
            # print(f'text_embeddings_hidden_states.shape: {text_embeddings_hidden_states.shape}')
            # print(f'text_embeddings_class_labels.shape: {text_embeddings_class_labels.shape}')
            # print(f'text_embeddings_boolean_prompt_mask.shape: {text_embeddings_boolean_prompt_mask.shape}')

            if save_compute:
                comb_out, _, _ = model.unet_forward(
                    xt_inp.expand(2, -1, -1, -1) if hasattr(model.model, 'unet') else xt_inp.expand(2, -1, -1),
                    timestep=t,
                    encoder_hidden_states=torch.cat([uncond_embeddings_hidden_states, text_embeddings_hidden_states
                                                     ], dim=0)
                    if uncond_embeddings_hidden_states is not None else None,
                    class_labels=torch.cat([uncond_embeddings_class_lables, text_embeddings_class_labels], dim=0)
                    if uncond_embeddings_class_lables is not None else None,
                    encoder_attention_mask=torch.cat([uncond_boolean_prompt_mask, text_embeddings_boolean_prompt_mask
                                                      ], dim=0)
                    if uncond_boolean_prompt_mask is not None else None,
                )
                uncond_out, cond_out = comb_out.sample.chunk(2, dim=0)
            else:
                uncond_out = model.unet_forward(
                    xt_inp, timestep=t,
                    encoder_hidden_states=uncond_embeddings_hidden_states,
                    class_labels=uncond_embeddings_class_lables,
                    encoder_attention_mask=uncond_boolean_prompt_mask,
                    )[0].sample

                # Conditional embedding
                cond_out = model.unet_forward(
                    xt_inp,
                    timestep=t,
                    encoder_hidden_states=text_embeddings_hidden_states,
                    class_labels=text_embeddings_class_labels,
                    encoder_attention_mask=text_embeddings_boolean_prompt_mask,
                    )[0].sample

        z = zs[idx] if zs is not None else None
        z = z.unsqueeze(0)
        # classifier free guidance
        noise_pred = uncond_out + (cfg_scales[0] * (cond_out - uncond_out)).sum(axis=0).unsqueeze(0)

        # 2. compute less noisy image and set x_t -> x_t-1
        xt = model.reverse_step_with_custom_noise(noise_pred, t, xt, variance_noise=z,
                                                  eta=etas[idx], first_order=first_order)

    del app_op.iterables[0]
    return xt, zs