Spaces:
Sleeping
Sleeping
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__() | |
def set_eval_and_device(device: torch.device, *modules): | |
for module in modules: | |
module.eval() | |
module.to(device) | |
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) | |
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 | |
class VideoSysPipelineOutput(BaseOutput): | |
video: torch.Tensor | |