beta version

wan/image2video.py

@@ -39,6 +39,9 @@ class WanI2V:
         use_usp=False,
         t5_cpu=False,
         init_on_cpu=True,
+        i2v720p= True,
+        model_filename ="",
+        text_encoder_filename="",
     ):
         r"""
         Initializes the image-to-video generation model components.
@@ -77,7 +80,7 @@ class WanI2V:
             text_len=config.text_len,
             dtype=config.t5_dtype,
             device=torch.device('cpu'),
-            checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
+            checkpoint_path=text_encoder_filename,
             tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
             shard_fn=shard_fn if t5_fsdp else None,
         )
@@ -95,8 +98,10 @@ class WanI2V:
                                          config.clip_checkpoint),
             tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
 
-        logging.info(f"Creating WanModel from {checkpoint_dir}")
-        self.model = WanModel.from_pretrained(checkpoint_dir)
+        logging.info(f"Creating WanModel from {model_filename}")
+        from mmgp import offload
+
+        self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel)
         self.model.eval().requires_grad_(False)
 
         if t5_fsdp or dit_fsdp or use_usp:
@@ -116,28 +121,30 @@ class WanI2V:
         else:
             self.sp_size = 1
 
-        if dist.is_initialized():
-            dist.barrier()
-        if dit_fsdp:
-            self.model = shard_fn(self.model)
-        else:
-            if not init_on_cpu:
-                self.model.to(self.device)
+        # if dist.is_initialized():
+        #     dist.barrier()
+        # if dit_fsdp:
+        #     self.model = shard_fn(self.model)
+        # else:
+        #     if not init_on_cpu:
+        #         self.model.to(self.device)
 
         self.sample_neg_prompt = config.sample_neg_prompt
 
     def generate(self,
-                 input_prompt,
-                 img,
-                 max_area=720 * 1280,
-                 frame_num=81,
-                 shift=5.0,
-                 sample_solver='unipc',
-                 sampling_steps=40,
-                 guide_scale=5.0,
-                 n_prompt="",
-                 seed=-1,
-                 offload_model=True):
+                input_prompt,
+                img,
+                max_area=720 * 1280,
+                frame_num=81,
+                shift=5.0,
+                sample_solver='unipc',
+                sampling_steps=40,
+                guide_scale=5.0,
+                n_prompt="",
+                seed=-1,
+                offload_model=True,
+                callback = None         
+                ):
         r"""
         Generates video frames from input image and text prompt using diffusion process.
 
@@ -197,14 +204,14 @@ class WanI2V:
         seed_g.manual_seed(seed)
         noise = torch.randn(
             16,
-            21,
+            int((frame_num - 1)/4 + 1), #21,
             lat_h,
             lat_w,
             dtype=torch.float32,
             generator=seed_g,
             device=self.device)
 
-        msk = torch.ones(1, 81, lat_h, lat_w, device=self.device)
+        msk = torch.ones(1, frame_num, lat_h, lat_w, device=self.device)
         msk[:, 1:] = 0
         msk = torch.concat([
             torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]
@@ -218,7 +225,7 @@ class WanI2V:
 
         # preprocess
         if not self.t5_cpu:
-            self.text_encoder.model.to(self.device)
+            # self.text_encoder.model.to(self.device)
             context = self.text_encoder([input_prompt], self.device)
             context_null = self.text_encoder([n_prompt], self.device)
             if offload_model:
@@ -229,20 +236,23 @@ class WanI2V:
             context = [t.to(self.device) for t in context]
             context_null = [t.to(self.device) for t in context_null]
 
-        self.clip.model.to(self.device)
+        # self.clip.model.to(self.device)
         clip_context = self.clip.visual([img[:, None, :, :]])
         if offload_model:
             self.clip.model.cpu()
 
-        y = self.vae.encode([
-            torch.concat([
-                torch.nn.functional.interpolate(
-                    img[None].cpu(), size=(h, w), mode='bicubic').transpose(
-                        0, 1),
-                torch.zeros(3, 80, h, w)
-            ],
-                         dim=1).to(self.device)
-        ])[0]
+        from mmgp import offload
+
+        offload.last_offload_obj.unload_all()
+        enc= torch.concat([
+            torch.nn.functional.interpolate(
+                img[None].cpu(), size=(h, w), mode='bicubic').transpose(
+                    0, 1).to(torch.bfloat16),
+            torch.zeros(3, frame_num-1, h, w, device="cpu", dtype= torch.bfloat16)
+        ], dim=1).to(self.device)
+        # enc = None
+
+        y = self.vae.encode([enc])[0]
         y = torch.concat([msk, y])
 
         @contextmanager
@@ -283,6 +293,7 @@ class WanI2V:
                 'clip_fea': clip_context,
                 'seq_len': max_seq_len,
                 'y': [y],
+                'pipeline' : self
             }
 
             arg_null = {
@@ -290,30 +301,39 @@ class WanI2V:
                 'clip_fea': clip_context,
                 'seq_len': max_seq_len,
                 'y': [y],
+                'pipeline' : self
             }
 
             if offload_model:
                 torch.cuda.empty_cache()
 
-            self.model.to(self.device)
-            for _, t in enumerate(tqdm(timesteps)):
+            # self.model.to(self.device)
+            if callback != None:
+                callback(-1, None)
+
+            self._interrupt = False
+            for i, t in enumerate(tqdm(timesteps)):
                 latent_model_input = [latent.to(self.device)]
                 timestep = [t]
 
                 timestep = torch.stack(timestep).to(self.device)
 
                 noise_pred_cond = self.model(
-                    latent_model_input, t=timestep, **arg_c)[0].to(
-                        torch.device('cpu') if offload_model else self.device)
+                    latent_model_input, t=timestep, **arg_c)[0]
+                if self._interrupt:
+                    return None                
                 if offload_model:
                     torch.cuda.empty_cache()
                 noise_pred_uncond = self.model(
-                    latent_model_input, t=timestep, **arg_null)[0].to(
-                        torch.device('cpu') if offload_model else self.device)
+                    latent_model_input, t=timestep, **arg_null)[0]
+                if self._interrupt:
+                    return None                
+                del latent_model_input
                 if offload_model:
                     torch.cuda.empty_cache()
                 noise_pred = noise_pred_uncond + guide_scale * (
                     noise_pred_cond - noise_pred_uncond)
+                del noise_pred_uncond
 
                 latent = latent.to(
                     torch.device('cpu') if offload_model else self.device)
@@ -325,9 +345,14 @@ class WanI2V:
                     return_dict=False,
                     generator=seed_g)[0]
                 latent = temp_x0.squeeze(0)
+                del temp_x0
+                del timestep
 
-                x0 = [latent.to(self.device)]
-                del latent_model_input, timestep
+                if callback is not None:
+                    callback(i, latent)         
+
+
+            x0 = [latent.to(self.device)]
 
             if offload_model:
                 self.model.cpu()

wan/modules/__init__.py

@@ -1,4 +1,4 @@
-from .attention import flash_attention
+from .attention import pay_attention
 from .model import WanModel
 from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
 from .tokenizers import HuggingfaceTokenizer
@@ -12,5 +12,5 @@ __all__ = [
     'T5Decoder',
     'T5EncoderModel',
     'HuggingfaceTokenizer',
-    'flash_attention',
+    'pay_attention',
 ]

wan/modules/attention.py

@@ -1,5 +1,9 @@
 # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
 import torch
+from importlib.metadata import version
+from mmgp import offload
+import torch.nn.functional as F
+
 
 try:
     import flash_attn_interface
@@ -12,19 +16,99 @@ try:
     FLASH_ATTN_2_AVAILABLE = True
 except ModuleNotFoundError:
     FLASH_ATTN_2_AVAILABLE = False
+    flash_attn = None
+
+try:
+    from sageattention import sageattn_varlen
+    def sageattn_varlen_wrapper(
+            q,
+            k,
+            v,
+            cu_seqlens_q,
+            cu_seqlens_kv,
+            max_seqlen_q,
+            max_seqlen_kv,
+        ):
+        return sageattn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv)
+except ImportError:
+    sageattn_varlen_wrapper = None
+
 
 import warnings
 
+try:
+    from sageattention import sageattn
+    @torch.compiler.disable()
+    def sageattn_wrapper(
+            qkv_list,
+            attention_length
+        ):
+        q,k, v = qkv_list
+        padding_length = q.shape[0] -attention_length
+        q = q[:attention_length, :, : ].unsqueeze(0)
+        k = k[:attention_length, :, : ].unsqueeze(0)
+        v = v[:attention_length, :, : ].unsqueeze(0)
+
+        o = sageattn(q, k, v, tensor_layout="NHD").squeeze(0)
+        del q, k ,v
+        qkv_list.clear()
+
+        if padding_length > 0:
+            o = torch.cat([o, torch.empty( (padding_length, *o.shape[-2:]), dtype= o.dtype, device=o.device  ) ], 0)
+
+        return o
+except ImportError:
+    sageattn = None
+
+
+@torch.compiler.disable()
+def sdpa_wrapper(
+        qkv_list,
+        attention_length
+    ):
+    q,k, v = qkv_list
+    padding_length = q.shape[0] -attention_length
+    q = q[:attention_length, :].transpose(0,1).unsqueeze(0)
+    k = k[:attention_length, :].transpose(0,1).unsqueeze(0)
+    v = v[:attention_length, :].transpose(0,1).unsqueeze(0)
+
+    o = F.scaled_dot_product_attention(
+        q, k, v, attn_mask=None, is_causal=False
+    ).squeeze(0).transpose(0,1)
+    del q, k ,v
+    qkv_list.clear()
+
+    if padding_length > 0:
+        o = torch.cat([o, torch.empty( (padding_length, *o.shape[-2:]), dtype= o.dtype, device=o.device  ) ], 0)
+
+    return o
+
+
+def get_attention_modes():
+    ret = ["sdpa", "auto"]
+    if flash_attn != None:
+        ret.append("flash")
+    # if memory_efficient_attention != None:
+    #     ret.append("xformers")
+    if sageattn_varlen_wrapper != None:
+        ret.append("sage")
+    if sageattn != None and version("sageattention").startswith("2") :
+        ret.append("sage2")
+
+    return ret
+
+
 __all__ = [
-    'flash_attention',
+    'pay_attention',
     'attention',
 ]
 
 
-def flash_attention(
-    q,
-    k,
-    v,
+def pay_attention(
+    qkv_list,
+    # q,
+    # k,
+    # v,
     q_lens=None,
     k_lens=None,
     dropout_p=0.,
@@ -49,6 +133,10 @@ def flash_attention(
     deterministic:  bool. If True, slightly slower and uses more memory.
     dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
     """
+    attn = offload.shared_state["_attention"]
+    q,k,v = qkv_list
+    qkv_list.clear()
+
     half_dtypes = (torch.float16, torch.bfloat16)
     assert dtype in half_dtypes
     assert q.device.type == 'cuda' and q.size(-1) <= 256
@@ -91,7 +179,27 @@ def flash_attention(
         )
 
     # apply attention
-    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+    if attn=="sage":
+        x = sageattn_varlen_wrapper(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_kv=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_kv=lk,
+        ).unflatten(0, (b, lq))
+    elif attn=="sage2":
+        qkv_list = [q,k,v]
+        del q,k,v
+        x = sageattn_wrapper(qkv_list, lq).unsqueeze(0)
+    elif attn=="sdpa":
+        qkv_list = [q, k, v]
+        del q, k , v
+        x = sdpa_wrapper( qkv_list, lq).unsqueeze(0)
+    elif attn=="flash" and (version is None or version == 3):
         # Note: dropout_p, window_size are not supported in FA3 now.
         x = flash_attn_interface.flash_attn_varlen_func(
             q=q,
@@ -108,8 +216,7 @@ def flash_attention(
             softmax_scale=softmax_scale,
             causal=causal,
             deterministic=deterministic)[0].unflatten(0, (b, lq))
-    else:
-        assert FLASH_ATTN_2_AVAILABLE
+    elif attn=="flash":
         x = flash_attn.flash_attn_varlen_func(
             q=q,
             k=k,
@@ -146,7 +253,7 @@ def attention(
     fa_version=None,
 ):
     if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
-        return flash_attention(
+        return pay_attention(
             q=q,
             k=k,
             v=v,

wan/modules/clip.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torchvision.transforms as T
 
-from .attention import flash_attention
+from .attention import pay_attention
 from .tokenizers import HuggingfaceTokenizer
 from .xlm_roberta import XLMRoberta
 
@@ -82,7 +82,7 @@ class SelfAttention(nn.Module):
 
         # compute attention
         p = self.attn_dropout if self.training else 0.0
-        x = flash_attention(q, k, v, dropout_p=p, causal=self.causal, version=2)
+        x = pay_attention([q, k, v], dropout_p=p, causal=self.causal, version=2)
         x = x.reshape(b, s, c)
 
         # output
@@ -194,7 +194,7 @@ class AttentionPool(nn.Module):
         k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)
 
         # compute attention
-        x = flash_attention(q, k, v, version=2)
+        x = pay_attention(q, k, v, version=2)
         x = x.reshape(b, 1, c)
 
         # output
@@ -441,11 +441,12 @@ def _clip(pretrained=False,
           device='cpu',
           **kwargs):
     # init a model on device
+    device ="cpu"
     with torch.device(device):
         model = model_cls(**kwargs)
 
     # set device
-    model = model.to(dtype=dtype, device=device)
+    # model = model.to(dtype=dtype, device=device)
     output = (model,)
 
     # init transforms
@@ -507,16 +508,19 @@ class CLIPModel:
         self.tokenizer_path = tokenizer_path
 
         # init model
-        self.model, self.transforms = clip_xlm_roberta_vit_h_14(
-            pretrained=False,
-            return_transforms=True,
-            return_tokenizer=False,
-            dtype=dtype,
-            device=device)
+        from accelerate import init_empty_weights
+
+        with init_empty_weights():
+            self.model, self.transforms = clip_xlm_roberta_vit_h_14(
+                pretrained=False,
+                return_transforms=True,
+                return_tokenizer=False,
+                dtype=dtype,
+                device=device)
         self.model = self.model.eval().requires_grad_(False)
         logging.info(f'loading {checkpoint_path}')
         self.model.load_state_dict(
-            torch.load(checkpoint_path, map_location='cpu'))
+            torch.load(checkpoint_path, map_location='cpu'), assign= True)
 
         # init tokenizer
         self.tokenizer = HuggingfaceTokenizer(

wan/modules/model.py

@@ -7,7 +7,7 @@ import torch.nn as nn
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
 
-from .attention import flash_attention
+from .attention import pay_attention
 
 __all__ = ['WanModel']
 
@@ -16,7 +16,7 @@ def sinusoidal_embedding_1d(dim, position):
     # preprocess
     assert dim % 2 == 0
     half = dim // 2
-    position = position.type(torch.float64)
+    position = position.type(torch.float32)
 
     # calculation
     sinusoid = torch.outer(
@@ -25,18 +25,47 @@ def sinusoidal_embedding_1d(dim, position):
     return x
 
 
-@amp.autocast(enabled=False)
+# @amp.autocast(enabled=False)
 def rope_params(max_seq_len, dim, theta=10000):
     assert dim % 2 == 0
     freqs = torch.outer(
         torch.arange(max_seq_len),
         1.0 / torch.pow(theta,
-                        torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+                        torch.arange(0, dim, 2).to(torch.float32).div(dim)))
     freqs = torch.polar(torch.ones_like(freqs), freqs)
     return freqs
 
 
-@amp.autocast(enabled=False)
+def rope_apply_(x, grid_sizes, freqs):
+    assert x.shape[0]==1
+
+    n, c = x.size(2), x.size(3) // 2
+
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    f, h, w = grid_sizes[0]
+    seq_len = f * h * w
+    x_i = x[0, :seq_len, :, :]
+
+    x_i = x_i.to(torch.float32)
+    x_i = x_i.reshape(seq_len, n, -1, 2)        
+    x_i = torch.view_as_complex(x_i)
+    freqs_i = torch.cat([
+        freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+        freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+        freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+    ], dim=-1)
+    freqs_i= freqs_i.reshape(seq_len, 1, -1)
+
+    # apply rotary embedding
+    x_i *= freqs_i
+    x_i = torch.view_as_real(x_i).flatten(2)
+    x[0, :seq_len, :, :] = x_i.to(torch.bfloat16)
+    # x_i = torch.cat([x_i, x[0, seq_len:]])
+    return x
+
+# @amp.autocast(enabled=False)
 def rope_apply(x, grid_sizes, freqs):
     n, c = x.size(2), x.size(3) // 2
 
@@ -45,12 +74,17 @@ def rope_apply(x, grid_sizes, freqs):
 
     # loop over samples
     output = []
-    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+    for i, (f, h, w) in enumerate(grid_sizes):
         seq_len = f * h * w
 
         # precompute multipliers
-        x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
-            seq_len, n, -1, 2))
+        # x_i = x[i, :seq_len]
+        x_i = x[i]
+        x_i = x_i[:seq_len, :, :]
+
+        x_i = x_i.to(torch.float32)
+        x_i = x_i.reshape(seq_len, n, -1, 2)        
+        x_i = torch.view_as_complex(x_i)
         freqs_i = torch.cat([
             freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
             freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
@@ -59,12 +93,14 @@ def rope_apply(x, grid_sizes, freqs):
                             dim=-1).reshape(seq_len, 1, -1)
 
         # apply rotary embedding
-        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+        x_i *= freqs_i
+        x_i = torch.view_as_real(x_i).flatten(2)
+        x_i = x_i.to(torch.bfloat16)
         x_i = torch.cat([x_i, x[i, seq_len:]])
 
         # append to collection
         output.append(x_i)
-    return torch.stack(output).float()
+    return torch.stack(output) #.float()
 
 
 class WanRMSNorm(nn.Module):
@@ -80,11 +116,31 @@ class WanRMSNorm(nn.Module):
         Args:
             x(Tensor): Shape [B, L, C]
         """
-        return self._norm(x.float()).type_as(x) * self.weight
+        y = x.float()
+        y.pow_(2)
+        y = y.mean(dim=-1, keepdim=True)
+        y += self.eps
+        y.rsqrt_()
+        x *=  y
+        x *= self.weight
+        return x
+        # return self._norm(x).type_as(x) * self.weight
 
     def _norm(self, x):
         return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
 
+def my_LayerNorm(norm, x):
+        y = x.float()
+        y_m = y.mean(dim=-1, keepdim=True)
+        y -= y_m 
+        del y_m
+        y.pow_(2)
+        y = y.mean(dim=-1, keepdim=True)
+        y += norm.eps
+        y.rsqrt_()
+        x = x *  y
+        return x
+
 
 class WanLayerNorm(nn.LayerNorm):
 
@@ -96,7 +152,13 @@ class WanLayerNorm(nn.LayerNorm):
         Args:
             x(Tensor): Shape [B, L, C]
         """
-        return super().forward(x.float()).type_as(x)
+        # return F.layer_norm(
+        #     input, self.normalized_shape, self.weight, self.bias, self.eps
+        # )
+        y = super().forward(x)
+        x = y.type_as(x)
+        return x
+        # return super().forward(x).type_as(x)
 
 
 class WanSelfAttention(nn.Module):
@@ -124,7 +186,7 @@ class WanSelfAttention(nn.Module):
         self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
-    def forward(self, x, seq_lens, grid_sizes, freqs):
+    def forward(self, xlist, seq_lens, grid_sizes, freqs):
         r"""
         Args:
             x(Tensor): Shape [B, L, num_heads, C / num_heads]
@@ -132,24 +194,31 @@ class WanSelfAttention(nn.Module):
             grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
             freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
         """
+        x = xlist[0]
+        xlist.clear()
+
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
 
         # query, key, value function
-        def qkv_fn(x):
-            q = self.norm_q(self.q(x)).view(b, s, n, d)
-            k = self.norm_k(self.k(x)).view(b, s, n, d)
-            v = self.v(x).view(b, s, n, d)
-            return q, k, v
-
-        q, k, v = qkv_fn(x)
-
-        x = flash_attention(
-            q=rope_apply(q, grid_sizes, freqs),
-            k=rope_apply(k, grid_sizes, freqs),
-            v=v,
-            k_lens=seq_lens,
+        q = self.q(x)
+        self.norm_q(q)
+        q = q.view(b, s, n, d) # !!!
+        k = self.k(x)
+        self.norm_k(k)
+        k = k.view(b, s, n, d) 
+        v = self.v(x).view(b, s, n, d)
+        del x
+        rope_apply_(q, grid_sizes, freqs)
+        rope_apply_(k, grid_sizes, freqs)
+        qkv_list = [q,k,v]
+        del q,k,v
+        x = pay_attention(
+            qkv_list,
+            # q=q,
+            # k=k,
+            # v=v,
+            # k_lens=seq_lens,
             window_size=self.window_size)
-
         # output
         x = x.flatten(2)
         x = self.o(x)
@@ -158,22 +227,31 @@ class WanSelfAttention(nn.Module):
 
 class WanT2VCrossAttention(WanSelfAttention):
 
-    def forward(self, x, context, context_lens):
+    def forward(self, xlist, context, context_lens):
         r"""
         Args:
             x(Tensor): Shape [B, L1, C]
             context(Tensor): Shape [B, L2, C]
             context_lens(Tensor): Shape [B]
         """
+        x = xlist[0]
+        xlist.clear()
         b, n, d = x.size(0), self.num_heads, self.head_dim
 
         # compute query, key, value
-        q = self.norm_q(self.q(x)).view(b, -1, n, d)
-        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        q = self.q(x)
+        del x
+        self.norm_q(q)
+        q= q.view(b, -1, n, d)
+        k = self.k(context)
+        self.norm_k(k)
+        k = k.view(b, -1, n, d)
         v = self.v(context).view(b, -1, n, d)
 
         # compute attention
-        x = flash_attention(q, k, v, k_lens=context_lens)
+        qvl_list=[q, k, v]
+        del q, k, v
+        x = pay_attention(qvl_list, k_lens=context_lens)
 
         # output
         x = x.flatten(2)
@@ -196,31 +274,54 @@ class WanI2VCrossAttention(WanSelfAttention):
         # self.alpha = nn.Parameter(torch.zeros((1, )))
         self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
-    def forward(self, x, context, context_lens):
+    def forward(self, xlist, context, context_lens):
         r"""
         Args:
             x(Tensor): Shape [B, L1, C]
             context(Tensor): Shape [B, L2, C]
             context_lens(Tensor): Shape [B]
         """
+
+        ##### Enjoy this spagheti VRAM optimizations done by DeepBeepMeep !
+        # I am sure you are a nice person and as you copy this code, you will give me officially proper credits:
+        # Please link to https://github.com/deepbeepmeep/Wan2GP and @deepbeepmeep on twitter  
+
+        x = xlist[0]
+        xlist.clear()
+
         context_img = context[:, :257]
         context = context[:, 257:]
         b, n, d = x.size(0), self.num_heads, self.head_dim
 
         # compute query, key, value
-        q = self.norm_q(self.q(x)).view(b, -1, n, d)
-        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        q = self.q(x)
+        del x
+        self.norm_q(q)
+        q= q.view(b, -1, n, d)
+        k = self.k(context)
+        self.norm_k(k)
+        k = k.view(b, -1, n, d)
         v = self.v(context).view(b, -1, n, d)
-        k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
+
+        qkv_list = [q, k, v]
+        del k,v
+        x = pay_attention(qkv_list, k_lens=context_lens)
+
+        k_img = self.k_img(context_img)
+        self.norm_k_img(k_img)
+        k_img = k_img.view(b, -1, n, d)
         v_img = self.v_img(context_img).view(b, -1, n, d)
-        img_x = flash_attention(q, k_img, v_img, k_lens=None)
+        qkv_list = [q, k_img, v_img]
+        del q, k_img, v_img
+        img_x = pay_attention(qkv_list, k_lens=None)
         # compute attention
-        x = flash_attention(q, k, v, k_lens=context_lens)
+
 
         # output
         x = x.flatten(2)
         img_x = img_x.flatten(2)
-        x = x + img_x
+        x += img_x
+        del img_x
         x = self.o(x)
         return x
 
@@ -289,27 +390,46 @@ class WanAttentionBlock(nn.Module):
             grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
             freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
         """
-        assert e.dtype == torch.float32
-        with amp.autocast(dtype=torch.float32):
-            e = (self.modulation + e).chunk(6, dim=1)
-        assert e[0].dtype == torch.float32
-
+        e = (self.modulation + e).chunk(6, dim=1)
+ 
         # self-attention
-        y = self.self_attn(
-            self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes,
-            freqs)
-        with amp.autocast(dtype=torch.float32):
-            x = x + y * e[2]
+        x_mod = self.norm1(x)
+        x_mod *= 1 + e[1]
+        x_mod += e[0]
+        xlist = [x_mod]
+        del x_mod
+        y = self.self_attn( xlist, seq_lens, grid_sizes,freqs)
+        x.addcmul_(y, e[2])
+        del y
+        y = self.norm3(x)
+        ylist= [y]
+        del y
+        x += self.cross_attn(ylist, context, context_lens)
+        y = self.norm2(x)
 
-        # cross-attention & ffn function
-        def cross_attn_ffn(x, context, context_lens, e):
-            x = x + self.cross_attn(self.norm3(x), context, context_lens)
-            y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
-            with amp.autocast(dtype=torch.float32):
-                x = x + y * e[5]
-            return x
+        y *= 1 + e[4]
+        y += e[3]
+
+
+        ffn = self.ffn[0]
+        gelu = self.ffn[1]
+        ffn2= self.ffn[2]
+
+        y_shape = y.shape
+        y = y.view(-1, y_shape[-1])
+        chunk_size = int(y_shape[1]/2.7)
+        chunks =torch.split(y, chunk_size)
+        for y_chunk  in chunks:
+            mlp_chunk = ffn(y_chunk)
+            mlp_chunk = gelu(mlp_chunk)
+            y_chunk[...] = ffn2(mlp_chunk)
+            del mlp_chunk 
+        y = y.view(y_shape)
+
+        x.addcmul_(y, e[5])
+
+       
 
-        x = cross_attn_ffn(x, context, context_lens, e)
         return x
 
 
@@ -336,10 +456,13 @@ class Head(nn.Module):
             x(Tensor): Shape [B, L1, C]
             e(Tensor): Shape [B, C]
         """
-        assert e.dtype == torch.float32
-        with amp.autocast(dtype=torch.float32):
-            e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
-            x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+        # assert e.dtype == torch.float32
+
+        e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+        x = self.norm(x).to(torch.bfloat16)
+        x *= (1 + e[1])
+        x += e[0]
+        x = self.head(x)
         return x
 
 
@@ -384,7 +507,8 @@ class WanModel(ModelMixin, ConfigMixin):
                  window_size=(-1, -1),
                  qk_norm=True,
                  cross_attn_norm=True,
-                 eps=1e-6):
+                 eps=1e-6,
+                 ):
         r"""
         Initialize the diffusion model backbone.
 
@@ -466,7 +590,7 @@ class WanModel(ModelMixin, ConfigMixin):
         # buffers (don't use register_buffer otherwise dtype will be changed in to())
         assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
         d = dim // num_heads
-        self.freqs = torch.cat([
+        self.freqs = torch.cat([    
             rope_params(1024, d - 4 * (d // 6)),
             rope_params(1024, 2 * (d // 6)),
             rope_params(1024, 2 * (d // 6))
@@ -487,6 +611,7 @@ class WanModel(ModelMixin, ConfigMixin):
         seq_len,
         clip_fea=None,
         y=None,
+        pipeline = None,
     ):
         r"""
         Forward pass through the diffusion model
@@ -521,8 +646,11 @@ class WanModel(ModelMixin, ConfigMixin):
 
         # embeddings
         x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
-        grid_sizes = torch.stack(
-            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        # grid_sizes = torch.stack(
+        #     [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+
+        grid_sizes = [ list(u.shape[2:]) for u in x]
+
         x = [u.flatten(2).transpose(1, 2) for u in x]
         seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
         assert seq_lens.max() <= seq_len
@@ -532,11 +660,10 @@ class WanModel(ModelMixin, ConfigMixin):
         ])
 
         # time embeddings
-        with amp.autocast(dtype=torch.float32):
-            e = self.time_embedding(
-                sinusoidal_embedding_1d(self.freq_dim, t).float())
-            e0 = self.time_projection(e).unflatten(1, (6, self.dim))
-            assert e.dtype == torch.float32 and e0.dtype == torch.float32
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, t))
+        e0 = self.time_projection(e).unflatten(1, (6, self.dim)).to(torch.bfloat16)
+            # assert e.dtype == torch.float32 and e0.dtype == torch.float32
 
         # context
         context_lens = None
@@ -561,6 +688,9 @@ class WanModel(ModelMixin, ConfigMixin):
             context_lens=context_lens)
 
         for block in self.blocks:
+            if pipeline._interrupt:
+                return [None]
+
             x = block(x, **kwargs)
 
         # head
@@ -588,7 +718,7 @@ class WanModel(ModelMixin, ConfigMixin):
 
         c = self.out_dim
         out = []
-        for u, v in zip(x, grid_sizes.tolist()):
+        for u, v in zip(x, grid_sizes):
             u = u[:math.prod(v)].view(*v, *self.patch_size, c)
             u = torch.einsum('fhwpqrc->cfphqwr', u)
             u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])

wan/modules/t5.py

@@ -442,7 +442,7 @@ def _t5(name,
         model = model_cls(**kwargs)
 
     # set device
-    model = model.to(dtype=dtype, device=device)
+    # model = model.to(dtype=dtype, device=device)
 
     # init tokenizer
     if return_tokenizer:
@@ -486,20 +486,25 @@ class T5EncoderModel:
         self.checkpoint_path = checkpoint_path
         self.tokenizer_path = tokenizer_path
 
+        from accelerate import init_empty_weights
         # init model
-        model = umt5_xxl(
-            encoder_only=True,
-            return_tokenizer=False,
-            dtype=dtype,
-            device=device).eval().requires_grad_(False)
+        with init_empty_weights():
+            model = umt5_xxl(
+                encoder_only=True,
+                return_tokenizer=False,
+                dtype=dtype,
+                device=device).eval().requires_grad_(False)
         logging.info(f'loading {checkpoint_path}')
-        model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
+        from mmgp import offload
+        offload.load_model_data(model,checkpoint_path )
+
         self.model = model
         if shard_fn is not None:
             self.model = shard_fn(self.model, sync_module_states=False)
         else:
             self.model.to(self.device)
         # init tokenizer
+        tokenizer_path= "google/umt5-xxl"
         self.tokenizer = HuggingfaceTokenizer(
             name=tokenizer_path, seq_len=text_len, clean='whitespace')
 

wan/modules/vae.py

@@ -1,6 +1,6 @@
 # Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
 import logging
-
+from mmgp import offload
 import torch
 import torch.cuda.amp as amp
 import torch.nn as nn
@@ -31,9 +31,16 @@ class CausalConv3d(nn.Conv3d):
             cache_x = cache_x.to(x.device)
             x = torch.cat([cache_x, x], dim=2)
             padding[4] -= cache_x.shape[2]
+            cache_x = None
         x = F.pad(x, padding)
+        x = super().forward(x)
 
-        return super().forward(x)
+        mem_threshold = offload.shared_state.get("_vae_threshold",0)
+        vae_config = offload.shared_state.get("_vae",1)
+        
+        if vae_config == 0  and torch.cuda.memory_reserved()  > mem_threshold  or vae_config == 2:  
+            torch.cuda.empty_cache()
+        return x
 
 
 class RMS_norm(nn.Module):
@@ -49,10 +56,11 @@ class RMS_norm(nn.Module):
         self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
 
     def forward(self, x):
-        return F.normalize(
+        x = F.normalize(
             x, dim=(1 if self.channel_first else
                     -1)) * self.scale * self.gamma + self.bias
-
+        x = x.to(torch.bfloat16)
+        return x 
 
 class Upsample(nn.Upsample):
 
@@ -107,11 +115,12 @@ class Resample(nn.Module):
                     feat_cache[idx] = 'Rep'
                     feat_idx[0] += 1
                 else:
-
-                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    clone = True
+                    cache_x = x[:, :, -CACHE_T:, :, :]#.clone()
                     if cache_x.shape[2] < 2 and feat_cache[
                             idx] is not None and feat_cache[idx] != 'Rep':
                         # cache last frame of last two chunk
+                        clone = False
                         cache_x = torch.cat([
                             feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
                                 cache_x.device), cache_x
@@ -119,11 +128,14 @@ class Resample(nn.Module):
                                             dim=2)
                     if cache_x.shape[2] < 2 and feat_cache[
                             idx] is not None and feat_cache[idx] == 'Rep':
+                        clone = False
                         cache_x = torch.cat([
                             torch.zeros_like(cache_x).to(cache_x.device),
                             cache_x
                         ],
                                             dim=2)
+                    if clone:
+                        cache_x = cache_x.clone()
                     if feat_cache[idx] == 'Rep':
                         x = self.time_conv(x)
                     else:
@@ -144,7 +156,7 @@ class Resample(nn.Module):
             if feat_cache is not None:
                 idx = feat_idx[0]
                 if feat_cache[idx] is None:
-                    feat_cache[idx] = x.clone()
+                    feat_cache[idx] = x #.to("cpu") #x.clone() yyyy
                     feat_idx[0] += 1
                 else:
 
@@ -155,7 +167,7 @@ class Resample(nn.Module):
 
                     x = self.time_conv(
                         torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
-                    feat_cache[idx] = cache_x
+                    feat_cache[idx] = cache_x#.to("cpu") #yyyyy
                     feat_idx[0] += 1
         return x
 
@@ -212,11 +224,11 @@ class ResidualBlock(nn.Module):
                             cache_x.device), cache_x
                     ],
                                         dim=2)
-                x = layer(x, feat_cache[idx])
-                feat_cache[idx] = cache_x
+                x = layer(x, feat_cache[idx]).to(torch.bfloat16)
+                feat_cache[idx] = cache_x#.to("cpu")
                 feat_idx[0] += 1
             else:
-                x = layer(x)
+                x = layer(x).to(torch.bfloat16)
         return x + h
 
 
@@ -326,12 +338,16 @@ class Encoder3d(nn.Module):
                         cache_x.device), cache_x
                 ],
                                     dim=2)
-            x = self.conv1(x, feat_cache[idx])
+            x = self.conv1(x, feat_cache[idx]).to(torch.bfloat16)
             feat_cache[idx] = cache_x
+            del cache_x
             feat_idx[0] += 1
         else:
             x = self.conv1(x)
 
+
+        # torch.cuda.empty_cache()
+
         ## downsamples
         for layer in self.downsamples:
             if feat_cache is not None:
@@ -339,6 +355,8 @@ class Encoder3d(nn.Module):
             else:
                 x = layer(x)
 
+        # torch.cuda.empty_cache()
+
         ## middle
         for layer in self.middle:
             if isinstance(layer, ResidualBlock) and feat_cache is not None:
@@ -346,6 +364,8 @@ class Encoder3d(nn.Module):
             else:
                 x = layer(x)
 
+        # torch.cuda.empty_cache()
+
         ## head
         for layer in self.head:
             if isinstance(layer, CausalConv3d) and feat_cache is not None:
@@ -360,9 +380,13 @@ class Encoder3d(nn.Module):
                                         dim=2)
                 x = layer(x, feat_cache[idx])
                 feat_cache[idx] = cache_x
+                del cache_x
                 feat_idx[0] += 1
             else:
                 x = layer(x)
+
+        # torch.cuda.empty_cache()
+
         return x
 
 
@@ -433,10 +457,12 @@ class Decoder3d(nn.Module):
                 ],
                                     dim=2)
             x = self.conv1(x, feat_cache[idx])
-            feat_cache[idx] = cache_x
+            feat_cache[idx] = cache_x#.to("cpu")
+            del cache_x
             feat_idx[0] += 1
         else:
             x = self.conv1(x)
+        cache_x = None
 
         ## middle
         for layer in self.middle:
@@ -456,7 +482,7 @@ class Decoder3d(nn.Module):
         for layer in self.head:
             if isinstance(layer, CausalConv3d) and feat_cache is not None:
                 idx = feat_idx[0]
-                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                cache_x = x[:, :, -CACHE_T:, :, :] .clone()
                 if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
                     # cache last frame of last two chunk
                     cache_x = torch.cat([
@@ -465,7 +491,8 @@ class Decoder3d(nn.Module):
                     ],
                                         dim=2)
                 x = layer(x, feat_cache[idx])
-                feat_cache[idx] = cache_x
+                feat_cache[idx] = cache_x#.to("cpu")
+                del cache_x
                 feat_idx[0] += 1
             else:
                 x = layer(x)
@@ -532,6 +559,8 @@ class WanVAE_(nn.Module):
                     feat_cache=self._enc_feat_map,
                     feat_idx=self._enc_conv_idx)
                 out = torch.cat([out, out_], 2)
+
+
         mu, log_var = self.conv1(out).chunk(2, dim=1)
         if isinstance(scale[0], torch.Tensor):
             mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(

wan/text2video.py

@@ -35,6 +35,8 @@ class WanT2V:
         dit_fsdp=False,
         use_usp=False,
         t5_cpu=False,
+        model_filename = None,
+        text_encoder_filename = None
     ):
         r"""
         Initializes the Wan text-to-video generation model components.
@@ -70,18 +72,26 @@ class WanT2V:
             text_len=config.text_len,
             dtype=config.t5_dtype,
             device=torch.device('cpu'),
-            checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
+            checkpoint_path=text_encoder_filename,
             tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
             shard_fn=shard_fn if t5_fsdp else None)
 
         self.vae_stride = config.vae_stride
         self.patch_size = config.patch_size
+
+        
         self.vae = WanVAE(
             vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
             device=self.device)
 
-        logging.info(f"Creating WanModel from {checkpoint_dir}")
-        self.model = WanModel.from_pretrained(checkpoint_dir)
+        logging.info(f"Creating WanModel from {model_filename}")
+        from mmgp import offload
+
+
+        self.model = offload.fast_load_transformers_model(model_filename, modelClass=WanModel)
+
+
+
         self.model.eval().requires_grad_(False)
 
         if use_usp:
@@ -98,12 +108,12 @@ class WanT2V:
         else:
             self.sp_size = 1
 
-        if dist.is_initialized():
-            dist.barrier()
-        if dit_fsdp:
-            self.model = shard_fn(self.model)
-        else:
-            self.model.to(self.device)
+        # if dist.is_initialized():
+        #     dist.barrier()
+        # if dit_fsdp:
+        #     self.model = shard_fn(self.model)
+        # else:
+        #     self.model.to(self.device)
 
         self.sample_neg_prompt = config.sample_neg_prompt
 
@@ -117,7 +127,9 @@ class WanT2V:
                  guide_scale=5.0,
                  n_prompt="",
                  seed=-1,
-                 offload_model=True):
+                 offload_model=True,
+                 callback = None
+                 ):
         r"""
         Generates video frames from text prompt using diffusion process.
 
@@ -168,7 +180,7 @@ class WanT2V:
         seed_g.manual_seed(seed)
 
         if not self.t5_cpu:
-            self.text_encoder.model.to(self.device)
+            # self.text_encoder.model.to(self.device)
             context = self.text_encoder([input_prompt], self.device)
             context_null = self.text_encoder([n_prompt], self.device)
             if offload_model:
@@ -223,23 +235,32 @@ class WanT2V:
             # sample videos
             latents = noise
 
-            arg_c = {'context': context, 'seq_len': seq_len}
-            arg_null = {'context': context_null, 'seq_len': seq_len}
+            arg_c = {'context': context, 'seq_len': seq_len, 'pipeline': self}
+            arg_null = {'context': context_null, 'seq_len': seq_len, 'pipeline': self}
 
-            for _, t in enumerate(tqdm(timesteps)):
+            if callback != None:
+                callback(-1, None)
+            self._interrupt = False
+            for i, t in enumerate(tqdm(timesteps)):
                 latent_model_input = latents
                 timestep = [t]
 
                 timestep = torch.stack(timestep)
 
-                self.model.to(self.device)
+                # self.model.to(self.device)
                 noise_pred_cond = self.model(
                     latent_model_input, t=timestep, **arg_c)[0]
+                if self._interrupt:
+                    return None               
                 noise_pred_uncond = self.model(
                     latent_model_input, t=timestep, **arg_null)[0]
+                if self._interrupt:
+                    return None
 
+                del latent_model_input
                 noise_pred = noise_pred_uncond + guide_scale * (
                     noise_pred_cond - noise_pred_uncond)
+                del noise_pred_uncond
 
                 temp_x0 = sample_scheduler.step(
                     noise_pred.unsqueeze(0),
@@ -248,6 +269,10 @@ class WanT2V:
                     return_dict=False,
                     generator=seed_g)[0]
                 latents = [temp_x0.squeeze(0)]
+                del temp_x0
+
+                if callback is not None:
+                    callback(i, latents)         
 
             x0 = latents
             if offload_model:
@@ -256,6 +281,7 @@ class WanT2V:
             if self.rank == 0:
                 videos = self.vae.decode(x0)
 
+
         del noise, latents
         del sample_scheduler
         if offload_model: