SubstanceSHIFT/wan-2-2-first-last-frame

1"""2"""3import contextlib4from contextvars import ContextVar5from io import BytesIO6from typing import Any7from typing import cast8from unittest.mock import patch9 10import torch11from torch._inductor.package.package import package_aoti12from torch.export.pt2_archive._package import AOTICompiledModel13from torch.export.pt2_archive._package_weights import Weights14 15 16INDUCTOR_CONFIGS_OVERRIDES = {17    'aot_inductor.package_constants_in_so': False,18    'aot_inductor.package_constants_on_disk': True,19    'aot_inductor.package': True,20}21 22 23class ZeroGPUWeights:24    def __init__(self, constants_map: dict[str, torch.Tensor], to_cuda: bool = False):25        if to_cuda:26            self.constants_map = {name: tensor.to('cuda') for name, tensor in constants_map.items()}27        else:28            self.constants_map = constants_map29    def __reduce__(self):30        constants_map: dict[str, torch.Tensor] = {}31        for name, tensor in self.constants_map.items():32            tensor_ = torch.empty_like(tensor, device='cpu').pin_memory()33            constants_map[name] = tensor_.copy_(tensor).detach().share_memory_()34        return ZeroGPUWeights, (constants_map, True)35 36 37class ZeroGPUCompiledModel:38    def __init__(self, archive_file: torch.types.FileLike, weights: ZeroGPUWeights):39        self.archive_file = archive_file40        self.weights = weights41        self.compiled_model: ContextVar[AOTICompiledModel | None] = ContextVar('compiled_model', default=None)42    def __call__(self, *args, **kwargs):43        if (compiled_model := self.compiled_model.get()) is None:44            compiled_model = cast(AOTICompiledModel, torch._inductor.aoti_load_package(self.archive_file))45            compiled_model.load_constants(self.weights.constants_map, check_full_update=True, user_managed=True)46            self.compiled_model.set(compiled_model)47        return compiled_model(*args, **kwargs)48    def __reduce__(self):49        return ZeroGPUCompiledModel, (self.archive_file, self.weights)50 51 52def aoti_compile(53    exported_program: torch.export.ExportedProgram,54    inductor_configs: dict[str, Any] | None = None,55):56    inductor_configs = (inductor_configs or {}) | INDUCTOR_CONFIGS_OVERRIDES57    gm = cast(torch.fx.GraphModule, exported_program.module())58    assert exported_program.example_inputs is not None59    args, kwargs = exported_program.example_inputs60    artifacts = torch._inductor.aot_compile(gm, args, kwargs, options=inductor_configs)61    archive_file = BytesIO()62    files: list[str | Weights] = [file for file in artifacts if isinstance(file, str)]63    package_aoti(archive_file, files)64    weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))65    zerogpu_weights = ZeroGPUWeights({name: weights.get_weight(name)[0] for name in weights})66    return ZeroGPUCompiledModel(archive_file, zerogpu_weights)67 68 69@contextlib.contextmanager70def capture_component_call(71    pipeline: Any,72    component_name: str,73    component_method='forward',74):75 76    class CapturedCallException(Exception):77        def __init__(self, *args, **kwargs):78            super().__init__()79            self.args = args80            self.kwargs = kwargs81 82    class CapturedCall:83        def __init__(self):84            self.args: tuple[Any, ...] = ()85            self.kwargs: dict[str, Any] = {}86 87    component = getattr(pipeline, component_name)88    captured_call = CapturedCall()89 90    def capture_call(*args, **kwargs):91        raise CapturedCallException(*args, **kwargs)92 93    with patch.object(component, component_method, new=capture_call):94        try:95            yield captured_call96        except CapturedCallException as e:97            captured_call.args = e.args98            captured_call.kwargs = e.kwargs99 100 101def drain_module_parameters(module: torch.nn.Module):102    state_dict_meta = {name: {'device': tensor.device, 'dtype': tensor.dtype} for name, tensor in module.state_dict().items()}103    state_dict = {name: torch.nn.Parameter(torch.empty_like(tensor, device='cpu')) for name, tensor in module.state_dict().items()}104    module.load_state_dict(state_dict, assign=True)105    for name, param in state_dict.items():106        meta = state_dict_meta[name]107        param.data = torch.Tensor([]).to(**meta)108