aoti-blocks-load (#23)

- AoTI blocks load (627bf16d600bbe23b2c9c25603c5afee24964694)


Co-authored-by: Charles Bensimon <[email protected]>

aoti.py

@@ -0,0 +1,35 @@
+"""
+"""
+
+from typing import cast
+
+import torch
+from huggingface_hub import hf_hub_download
+from spaces.zero.torch.aoti import ZeroGPUCompiledModel
+from spaces.zero.torch.aoti import ZeroGPUWeights
+from torch._functorch._aot_autograd.subclass_parametrization import unwrap_tensor_subclass_parameters
+
+
+def _shallow_clone_module(module: torch.nn.Module) -> torch.nn.Module:
+    clone = object.__new__(module.__class__)
+    clone.__dict__ = module.__dict__.copy()
+    clone._parameters = module._parameters.copy()
+    clone._buffers = module._buffers.copy()
+    clone._modules = {k: _shallow_clone_module(v) for k, v in module._modules.items() if v is not None}
+    return clone
+
+
+def aoti_blocks_load(module: torch.nn.Module, repo_id: str, variant: str | None = None):
+    repeated_blocks = cast(list[str], module._repeated_blocks)
+    aoti_files = {name: hf_hub_download(
+        repo_id=repo_id,
+        filename='package.pt2',
+        subfolder=name if variant is None else f'{name}.{variant}',
+    ) for name in repeated_blocks}
+    for block_name, aoti_file in aoti_files.items():
+        for block in module.modules():
+            if block.__class__.__name__ == block_name:
+                block_ = _shallow_clone_module(block)
+                unwrap_tensor_subclass_parameters(block_)
+                weights = ZeroGPUWeights(block_.state_dict())
+                block.forward = ZeroGPUCompiledModel(aoti_file, weights)

app.py

@@ -9,7 +9,12 @@ import numpy as np
 from PIL import Image
 import random
 import gc
-from optimization import optimize_pipeline_
+
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
+from torchao.quantization import Int8WeightOnlyConfig
+
+import aoti
 
 
 MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
@@ -23,7 +28,7 @@ MAX_SEED = np.iinfo(np.int32).max
 
 FIXED_FPS = 16
 MIN_FRAMES_MODEL = 8
-MAX_FRAMES_MODEL = 81
+MAX_FRAMES_MODEL = 80
 
 MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS,1)
 MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS,1)
@@ -43,21 +48,29 @@ pipe = WanImageToVideoPipeline.from_pretrained(MODEL_ID,
     torch_dtype=torch.bfloat16,
 ).to('cuda')
 
-for i in range(3): 
-    gc.collect()
-    torch.cuda.synchronize() 
-    torch.cuda.empty_cache()
+pipe.load_lora_weights(
+    "Kijai/WanVideo_comfy", 
+    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors", 
+    adapter_name="lightx2v"
+)
+kwargs_lora = {}
+kwargs_lora["load_into_transformer_2"] = True
+pipe.load_lora_weights(
+    "Kijai/WanVideo_comfy", 
+    weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors", 
+    adapter_name="lightx2v_2", **kwargs_lora
+)
+pipe.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
+pipe.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
+pipe.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
+pipe.unload_lora_weights()
 
-OPTIMIZE_WIDTH = 832 
-OPTIMIZE_HEIGHT = 624
+quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
 
-optimize_pipeline_(pipe,
-    image=Image.new('RGB', (OPTIMIZE_WIDTH, OPTIMIZE_HEIGHT)),
-    prompt='prompt',
-    height=OPTIMIZE_HEIGHT,
-    width=OPTIMIZE_WIDTH,
-    num_frames=MAX_FRAMES_MODEL,
-)
+aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
+aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
 
 
 default_prompt_i2v = "make this image come alive, cinematic motion, smooth animation"
@@ -109,6 +122,14 @@ def resize_image(image: Image.Image) -> Image.Image:
     return image_to_resize.resize((final_w, final_h), Image.LANCZOS)
 
 
+def get_num_frames(duration_seconds: float):
+    return 1 + int(np.clip(
+        int(round(duration_seconds * FIXED_FPS)),
+        MIN_FRAMES_MODEL,
+        MAX_FRAMES_MODEL,
+    ))
+
+
 def get_duration(
     input_image,
     prompt,
@@ -121,7 +142,13 @@ def get_duration(
     randomize_seed,
     progress,
 ):
-    return int(steps) * 15
+    BASE_FRAMES_HEIGHT_WIDTH = 81 * 832 * 624
+    BASE_STEP_DURATION = 15
+    width, height = resize_image(input_image).size
+    frames = get_num_frames(duration_seconds)
+    factor = frames * width * height / BASE_FRAMES_HEIGHT_WIDTH
+    step_duration = BASE_STEP_DURATION * factor ** 1.5
+    return 10 + int(steps) * step_duration
 
 @spaces.GPU(duration=get_duration)
 def generate_video(
@@ -179,7 +206,7 @@ def generate_video(
     if input_image is None:
         raise gr.Error("Please upload an input image.")
     
-    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+    num_frames = get_num_frames(duration_seconds)
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
     resized_image = resize_image(input_image)
 

optimization.py

@@ -1,106 +0,0 @@
-"""
-"""
-
-from typing import Any
-from typing import Callable
-from typing import ParamSpec
-
-import spaces
-import torch
-from torch.utils._pytree import tree_map_only
-from torchao.quantization import quantize_
-from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
-from torchao.quantization import Int8WeightOnlyConfig
-
-from optimization_utils import capture_component_call
-from optimization_utils import aoti_compile
-from optimization_utils import drain_module_parameters
-
-
-P = ParamSpec('P')
-
-LATENT_FRAMES_DIM = torch.export.Dim('num_latent_frames', min=8, max=81)
-
-LATENT_PATCHED_HEIGHT_DIM = torch.export.Dim('latent_patched_height', min=30, max=52)
-LATENT_PATCHED_WIDTH_DIM = torch.export.Dim('latent_patched_width', min=30, max=52)
-
-TRANSFORMER_DYNAMIC_SHAPES = {
-    'hidden_states': {
-        2: LATENT_FRAMES_DIM,
-        3: 2 * LATENT_PATCHED_HEIGHT_DIM,
-        4: 2 * LATENT_PATCHED_WIDTH_DIM,
-    },
-}
-
-INDUCTOR_CONFIGS = {
-    'conv_1x1_as_mm': True,
-    'epilogue_fusion': False,
-    'coordinate_descent_tuning': True,
-    'coordinate_descent_check_all_directions': True,
-    'max_autotune': True,
-    'triton.cudagraphs': True,
-}
-
-
-def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-
-    @spaces.GPU(duration=1500)
-    def compile_transformer():
-        
-        # This LoRA fusion part remains the same
-        pipeline.load_lora_weights(
-            "Kijai/WanVideo_comfy", 
-            weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors", 
-            adapter_name="lightx2v"
-        )
-        kwargs_lora = {}
-        kwargs_lora["load_into_transformer_2"] = True
-        pipeline.load_lora_weights(
-            "Kijai/WanVideo_comfy", 
-            weight_name="Lightx2v/lightx2v_I2V_14B_480p_cfg_step_distill_rank128_bf16.safetensors", 
-            adapter_name="lightx2v_2", **kwargs_lora
-        )
-        pipeline.set_adapters(["lightx2v", "lightx2v_2"], adapter_weights=[1., 1.])
-        pipeline.fuse_lora(adapter_names=["lightx2v"], lora_scale=3., components=["transformer"])
-        pipeline.fuse_lora(adapter_names=["lightx2v_2"], lora_scale=1., components=["transformer_2"])
-        pipeline.unload_lora_weights()
-        
-        with capture_component_call(pipeline, 'transformer') as call:
-            pipeline(*args, **kwargs)
-        
-        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
-
-        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-        quantize_(pipeline.transformer_2, Float8DynamicActivationFloat8WeightConfig())
-        
-        
-        exported_1 = torch.export.export(
-            mod=pipeline.transformer,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
-        
-        exported_2 = torch.export.export(
-            mod=pipeline.transformer_2,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
-
-        compiled_1 = aoti_compile(exported_1, INDUCTOR_CONFIGS)
-        compiled_2 = aoti_compile(exported_2, INDUCTOR_CONFIGS)
-        
-        return compiled_1, compiled_2
-
-
-    quantize_(pipeline.text_encoder, Int8WeightOnlyConfig())
-    
-    compiled_transformer_1, compiled_transformer_2 = compile_transformer()
-
-    pipeline.transformer.forward = compiled_transformer_1
-    drain_module_parameters(pipeline.transformer)
-
-    pipeline.transformer_2.forward = compiled_transformer_2
-    drain_module_parameters(pipeline.transformer_2)

optimization_utils.py

@@ -1,107 +0,0 @@
-"""
-"""
-import contextlib
-from contextvars import ContextVar
-from io import BytesIO
-from typing import Any
-from typing import cast
-from unittest.mock import patch
-
-import torch
-from torch._inductor.package.package import package_aoti
-from torch.export.pt2_archive._package import AOTICompiledModel
-from torch.export.pt2_archive._package_weights import Weights
-
-
-INDUCTOR_CONFIGS_OVERRIDES = {
-    'aot_inductor.package_constants_in_so': False,
-    'aot_inductor.package_constants_on_disk': True,
-    'aot_inductor.package': True,
-}
-
-
-class ZeroGPUWeights:
-    def __init__(self, constants_map: dict[str, torch.Tensor], to_cuda: bool = False):
-        if to_cuda:
-            self.constants_map = {name: tensor.to('cuda') for name, tensor in constants_map.items()}
-        else:
-            self.constants_map = constants_map
-    def __reduce__(self):
-        constants_map: dict[str, torch.Tensor] = {}
-        for name, tensor in self.constants_map.items():
-            tensor_ = torch.empty_like(tensor, device='cpu').pin_memory()
-            constants_map[name] = tensor_.copy_(tensor).detach().share_memory_()
-        return ZeroGPUWeights, (constants_map, True)
-
-
-class ZeroGPUCompiledModel:
-    def __init__(self, archive_file: torch.types.FileLike, weights: ZeroGPUWeights):
-        self.archive_file = archive_file
-        self.weights = weights
-        self.compiled_model: ContextVar[AOTICompiledModel | None] = ContextVar('compiled_model', default=None)
-    def __call__(self, *args, **kwargs):
-        if (compiled_model := self.compiled_model.get()) is None:
-            compiled_model = cast(AOTICompiledModel, torch._inductor.aoti_load_package(self.archive_file))
-            compiled_model.load_constants(self.weights.constants_map, check_full_update=True, user_managed=True)
-            self.compiled_model.set(compiled_model)
-        return compiled_model(*args, **kwargs)
-    def __reduce__(self):
-        return ZeroGPUCompiledModel, (self.archive_file, self.weights)
-
-
-def aoti_compile(
-    exported_program: torch.export.ExportedProgram,
-    inductor_configs: dict[str, Any] | None = None,
-):
-    inductor_configs = (inductor_configs or {}) | INDUCTOR_CONFIGS_OVERRIDES
-    gm = cast(torch.fx.GraphModule, exported_program.module())
-    assert exported_program.example_inputs is not None
-    args, kwargs = exported_program.example_inputs
-    artifacts = torch._inductor.aot_compile(gm, args, kwargs, options=inductor_configs)
-    archive_file = BytesIO()
-    files: list[str | Weights] = [file for file in artifacts if isinstance(file, str)]
-    package_aoti(archive_file, files)
-    weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))
-    zerogpu_weights = ZeroGPUWeights({name: weights.get_weight(name)[0] for name in weights})
-    return ZeroGPUCompiledModel(archive_file, zerogpu_weights)
-
-
-@contextlib.contextmanager
-def capture_component_call(
-    pipeline: Any,
-    component_name: str,
-    component_method='forward',
-):
-
-    class CapturedCallException(Exception):
-        def __init__(self, *args, **kwargs):
-            super().__init__()
-            self.args = args
-            self.kwargs = kwargs
-
-    class CapturedCall:
-        def __init__(self):
-            self.args: tuple[Any, ...] = ()
-            self.kwargs: dict[str, Any] = {}
-
-    component = getattr(pipeline, component_name)
-    captured_call = CapturedCall()
-
-    def capture_call(*args, **kwargs):
-        raise CapturedCallException(*args, **kwargs)
-
-    with patch.object(component, component_method, new=capture_call):
-        try:
-            yield captured_call
-        except CapturedCallException as e:
-            captured_call.args = e.args
-            captured_call.kwargs = e.kwargs
-
-
-def drain_module_parameters(module: torch.nn.Module):
-    state_dict_meta = {name: {'device': tensor.device, 'dtype': tensor.dtype} for name, tensor in module.state_dict().items()}
-    state_dict = {name: torch.nn.Parameter(torch.empty_like(tensor, device='cpu')) for name, tensor in module.state_dict().items()}
-    module.load_state_dict(state_dict, assign=True)
-    for name, param in state_dict.items():
-        meta = state_dict_meta[name]
-        param.data = torch.Tensor([]).to(**meta)