diffusers_vdm/pipeline.py

import os
import torch
import einops

from diffusers import DiffusionPipeline
from transformers import CLIPTextModel, CLIPTokenizer
from huggingface_hub import snapshot_download
from diffusers_vdm.vae import VideoAutoencoderKL
from diffusers_vdm.projection import Resampler
from diffusers_vdm.unet import UNet3DModel
from diffusers_vdm.improved_clip_vision import ImprovedCLIPVisionModelWithProjection
from diffusers_vdm.dynamic_tsnr_sampler import SamplerDynamicTSNR


class LatentVideoDiffusionPipeline(DiffusionPipeline):
    def __init__(self, tokenizer, text_encoder, image_encoder, vae, image_projection, unet, fp16=True, eval=True):
        super().__init__()

        self.loading_components = dict(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            unet=unet,
            image_encoder=image_encoder,
            image_projection=image_projection
        )

        for k, v in self.loading_components.items():
            setattr(self, k, v)

        if fp16:
            self.vae.half()
            self.text_encoder.half()
            self.unet.half()
            self.image_encoder.half()
            self.image_projection.half()

        self.vae.requires_grad_(False)
        self.text_encoder.requires_grad_(False)
        self.image_encoder.requires_grad_(False)

        self.vae.eval()
        self.text_encoder.eval()
        self.image_encoder.eval()

        if eval:
            self.unet.eval()
            self.image_projection.eval()
        else:
            self.unet.train()
            self.image_projection.train()

    def to(self, *args, **kwargs):
        for k, v in self.loading_components.items():
            if hasattr(v, 'to'):
                v.to(*args, **kwargs)
        return self

    def save_pretrained(self, save_directory, **kwargs):
        for k, v in self.loading_components.items():
            folder = os.path.join(save_directory, k)
            os.makedirs(folder, exist_ok=True)
            v.save_pretrained(folder)
        return

    @classmethod
    def from_pretrained(cls, repo_id, fp16=True, eval=True, token=None):
        local_folder = snapshot_download(repo_id=repo_id, token=token)
        return cls(
            tokenizer=CLIPTokenizer.from_pretrained(os.path.join(local_folder, "tokenizer")),
            text_encoder=CLIPTextModel.from_pretrained(os.path.join(local_folder, "text_encoder")),
            image_encoder=ImprovedCLIPVisionModelWithProjection.from_pretrained(os.path.join(local_folder, "image_encoder")),
            vae=VideoAutoencoderKL.from_pretrained(os.path.join(local_folder, "vae")),
            image_projection=Resampler.from_pretrained(os.path.join(local_folder, "image_projection")),
            unet=UNet3DModel.from_pretrained(os.path.join(local_folder, "unet")),
            fp16=fp16,
            eval=eval
        )

    @torch.inference_mode()
    def encode_cropped_prompt_77tokens(self, prompt: str):
        cond_ids = self.tokenizer(prompt,
                                  padding="max_length",
                                  max_length=self.tokenizer.model_max_length,
                                  truncation=True,
                                  return_tensors="pt").input_ids.to(self.text_encoder.device)
        cond = self.text_encoder(cond_ids, attention_mask=None).last_hidden_state
        return cond

    @torch.inference_mode()
    def encode_clip_vision(self, frames):
        b, c, t, h, w = frames.shape
        frames = einops.rearrange(frames, 'b c t h w -> (b t) c h w')
        clipvision_embed = self.image_encoder(frames).last_hidden_state
        clipvision_embed = einops.rearrange(clipvision_embed, '(b t) d c -> b t d c', t=t)
        return clipvision_embed

    @torch.inference_mode()
    def encode_latents(self, videos, return_hidden_states=True):
        b, c, t, h, w = videos.shape
        x = einops.rearrange(videos, 'b c t h w -> (b t) c h w')
        encoder_posterior, hidden_states = self.vae.encode(x, return_hidden_states=return_hidden_states)
        z = encoder_posterior.mode() * self.vae.scale_factor
        z = einops.rearrange(z, '(b t) c h w -> b c t h w', b=b, t=t)

        if not return_hidden_states:
            return z

        hidden_states = [einops.rearrange(h, '(b t) c h w -> b c t h w', b=b) for h in hidden_states]
        hidden_states = [h[:, :, [0, -1], :, :] for h in hidden_states]  # only need first and last

        return z, hidden_states

    @torch.inference_mode()
    def decode_latents(self, latents, hidden_states):
        B, C, T, H, W = latents.shape
        latents = einops.rearrange(latents, 'b c t h w -> (b t) c h w')
        latents = latents.to(device=self.vae.device, dtype=self.vae.dtype) / self.vae.scale_factor
        pixels = self.vae.decode(latents, ref_context=hidden_states, timesteps=T)
        pixels = einops.rearrange(pixels, '(b t) c h w -> b c t h w', b=B, t=T)
        return pixels

    @torch.inference_mode()
    def __call__(
            self,
            batch_size: int = 1,
            steps: int = 50,
            guidance_scale: float = 5.0,
            positive_text_cond = None,
            negative_text_cond = None,
            positive_image_cond = None,
            negative_image_cond = None,
            concat_cond = None,
            fs = 3,
            progress_tqdm = None,
    ):
        unet_is_training = self.unet.training

        if unet_is_training:
            self.unet.eval()

        device = self.unet.device
        dtype = self.unet.dtype
        dynamic_tsnr_model = SamplerDynamicTSNR(self.unet)

        # Batch

        concat_cond = concat_cond.repeat(batch_size, 1, 1, 1, 1).to(device=device, dtype=dtype)  # b, c, t, h, w
        positive_text_cond = positive_text_cond.repeat(batch_size, 1, 1).to(concat_cond)  # b, f, c
        negative_text_cond = negative_text_cond.repeat(batch_size, 1, 1).to(concat_cond)  # b, f, c
        positive_image_cond = positive_image_cond.repeat(batch_size, 1, 1, 1).to(concat_cond)  # b, t, l, c
        negative_image_cond = negative_image_cond.repeat(batch_size, 1, 1, 1).to(concat_cond)

        if isinstance(fs, torch.Tensor):
            fs = fs.repeat(batch_size, ).to(dtype=torch.long, device=device)  # b
        else:
            fs = torch.tensor([fs] * batch_size, dtype=torch.long, device=device)  # b

        # Initial latents

        latent_shape = concat_cond.shape

        # Feeds

        sampler_kwargs = dict(
            cfg_scale=guidance_scale,
            positive=dict(
                context_text=positive_text_cond,
                context_img=positive_image_cond,
                fs=fs,
                concat_cond=concat_cond
            ),
            negative=dict(
                context_text=negative_text_cond,
                context_img=negative_image_cond,
                fs=fs,
                concat_cond=concat_cond
            )
        )

        # Sample

        results = dynamic_tsnr_model(latent_shape, steps, extra_args=sampler_kwargs, progress_tqdm=progress_tqdm)

        if unet_is_training:
            self.unet.train()

        return results