drizzymedia/StreamDiffusionV2-Realtime

1from models.scheduler import SchedulerInterface2from abc import abstractmethod, ABC3from typing import List, Optional4import torch5import types6 7 8class DiffusionModelInterface(ABC, torch.nn.Module):9    scheduler: SchedulerInterface10 11    @abstractmethod12    def forward(13        self, noisy_image_or_video: torch.Tensor, conditional_dict: dict,14        timestep: torch.Tensor, kv_cache: Optional[List[dict]] = None,15        crossattn_cache: Optional[List[dict]] = None,16        current_start: Optional[int] = None,17        current_end: Optional[int] = None18    ) -> torch.Tensor:19        """20        A method to run diffusion model.21        Input:22            - noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.23            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).24            - timestep: a tensor with shape [B, F]  where the number of frame is 1 for images.25            all data should be on the same device as the model.26            - kv_cache: a list of dictionaries containing the key and value tensors for each attention layer.27            - current_start: the start index of the current frame in the sequence.28            - current_end: the end index of the current frame in the sequence.29        Output: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.30        We always expect a X0 prediction form for the output.31        """32        pass33 34    def get_scheduler(self) -> SchedulerInterface:35        """36        Update the current scheduler with the interface's static method37        """38        scheduler = self.scheduler39        scheduler.convert_x0_to_noise = types.MethodType(40            SchedulerInterface.convert_x0_to_noise, scheduler)41        scheduler.convert_noise_to_x0 = types.MethodType(42            SchedulerInterface.convert_noise_to_x0, scheduler)43        scheduler.convert_velocity_to_x0 = types.MethodType(44            SchedulerInterface.convert_velocity_to_x0, scheduler)45        self.scheduler = scheduler46        return scheduler47 48    def post_init(self):49        """50        A few custom initialization steps that should be called after the object is created.51        Currently, the only one we have is to bind a few methods to scheduler.52        We can gradually add more methods here if needed.53        """54        self.get_scheduler()55 56    def set_module_grad(self, module_grad: dict) -> None:57        """58        Adjusts the state of each module in the object.59 60        Parameters:61        - module_grad (dict): A dictionary where each key is the name of a module (as an attribute of the object),62          and each value is a bool indicating whether the module's parameters require gradients.63 64        Functionality:65        For each module name in the dictionary:66        - Updates whether its parameters require gradients based on 'is_trainable'.67        """68        for k, is_trainable in module_grad.items():69            getattr(self, k).requires_grad_(is_trainable)70 71    @abstractmethod72    def enable_gradient_checkpointing(self) -> None:73        """74        Activates gradient checkpointing for the current model (may be referred to as *activation checkpointing* or75        *checkpoint activations* in other frameworks).76        """77        pass78 79 80class VAEInterface(ABC, torch.nn.Module):81    @abstractmethod82    def decode_to_pixel(self, latent: torch.Tensor) -> torch.Tensor:83        """84        A method to decode a latent representation to an image or video.85        Input: a tensor with shape [B, F // T, C, H // S, W // S] where T and S are temporal and spatial compression factors.86        Output: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.87        """88        pass89 90 91class TextEncoderInterface(ABC, torch.nn.Module):92    @abstractmethod93    def forward(self, text_prompts: List[str]) -> dict:94        """95        A method to tokenize text prompts with a tokenizer and encode them into a latent representation.96        Input: a list of strings.97        Output: a dictionary containing the encoded representation of the text prompts.98        """99        pass100 101 102class InferencePipelineInterface(ABC):103    @abstractmethod104    def inference_with_trajectory(self, noise: torch.Tensor, conditional_dict: dict) -> torch.Tensor:105        """106        Run inference with the given diffusion / distilled generators.107        Input:108            - noise: a tensor sampled from N(0, 1) with shape [B, F, C, H, W] where the number of frame is 1 for images.109            - conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).110        Output:111            - output: a tensor with shape [B, T, F, C, H, W].112            T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0113            represents the x0 prediction at each timestep.114        """115