File size: 1,755 Bytes
a28e78a 07c6a04 a28e78a 07c6a04 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import inspect
from abc import abstractmethod
from dataclasses import dataclass
import torch
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.utils import BaseOutput
class VideoSysPipeline(DiffusionPipeline):
def __init__(self):
super().__init__()
@staticmethod
def set_eval_and_device(device: torch.device, *modules):
for module in modules:
module.eval()
module.to(device)
@abstractmethod
def generate(self, *args, **kwargs):
pass
def __call__(self, *args, **kwargs):
"""
In diffusers, it is a convention to call the pipeline object.
But in VideoSys, we will use the generate method for better prompt.
This is a wrapper for the generate method to support the diffusers usage.
"""
return self.generate(*args, **kwargs)
@classmethod
def _get_signature_keys(cls, obj):
parameters = inspect.signature(obj.__init__).parameters
required_parameters = {k: v for k, v in parameters.items() if v.default == inspect._empty}
optional_parameters = set({k for k, v in parameters.items() if v.default != inspect._empty})
expected_modules = set(required_parameters.keys()) - {"self"}
# modify: remove the config module from the expected modules
expected_modules = expected_modules - {"config"}
optional_names = list(optional_parameters)
for name in optional_names:
if name in cls._optional_components:
expected_modules.add(name)
optional_parameters.remove(name)
return expected_modules, optional_parameters
@dataclass
class VideoSysPipelineOutput(BaseOutput):
video: torch.Tensor
|