appiq-solution 1.4.8 → 1.5.0

package/#Tools/SkyReels-V2/skycaptioner_v1/scripts/vllm_struct_caption.py

@@ -1,156 +0,0 @@
-
-import torch
-import decord
-import argparse
-
-import pandas as pd
-import numpy as np
-
-from tqdm import tqdm
-from vllm import LLM, SamplingParams
-from transformers import AutoTokenizer, AutoProcessor
-
-from torch.utils.data import DataLoader
-
-SYSTEM_PROMPT = "I need you to generate a structured and detailed caption for the provided video. The structured output and the requirements for each field are as shown in the following JSON content: {\"subjects\": [{\"appearance\": \"Main subject appearance description\", \"action\": \"Main subject action\", \"expression\": \"Main subject expression  (Only for human/animal categories, empty otherwise)\", \"position\": \"Subject position in the video (Can be relative position to other objects or spatial description)\", \"TYPES\": {\"type\": \"Main category (e.g., Human)\", \"sub_type\": \"Sub-category (e.g., Man)\"}, \"is_main_subject\": true}, {\"appearance\": \"Non-main subject appearance description\", \"action\": \"Non-main subject action\", \"expression\": \"Non-main subject expression (Only for human/animal categories, empty otherwise)\", \"position\": \"Position of non-main subject 1\", \"TYPES\": {\"type\": \"Main category (e.g., Vehicles)\", \"sub_type\": \"Sub-category (e.g., Ship)\"}, \"is_main_subject\": false}], \"shot_type\": \"Shot type(Options: long_shot/full_shot/medium_shot/close_up/extreme_close_up/other)\", \"shot_angle\": \"Camera angle(Options: eye_level/high_angle/low_angle/other)\", \"shot_position\": \"Camera position(Options: front_view/back_view/side_view/over_the_shoulder/overhead_view/point_of_view/aerial_view/overlooking_view/other)\", \"camera_motion\": \"Camera movement description\", \"environment\": \"Video background/environment description\", \"lighting\": \"Lighting information in the video\"}"
-
-
-class VideoTextDataset(torch.utils.data.Dataset):
-    def __init__(self, csv_path, model_path):
-        if isinstance(csv_path, pd.DataFrame):
-            self.meta = csv_path
-        else:
-            self.meta = pd.read_csv(csv_path)
-        self._path = 'path'
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-        self.processor = AutoProcessor.from_pretrained(model_path)
-  
-    def __getitem__(self, index):
-        row = self.meta.iloc[index]
-        path = row[self._path]
-        real_index = self.meta.index[index]
-        vr = decord.VideoReader(path, ctx=decord.cpu(0), width=360, height=420)
-        start = 0
-        end = len(vr)
-        # avg_fps = vr.get_avg_fps()
-        index = self.get_index(end-start, 16, st=start)
-        frames = vr.get_batch(index).asnumpy() # n h w c
-        video_inputs = [torch.from_numpy(frames).permute(0, 3, 1, 2)]
-        conversation = {
-                    "role": "user",
-                    "content": [
-                        {
-                            "type": "video",
-                            "video": row['path'],
-                            "max_pixels": 360 * 420, # 460800
-                            "fps": 2.0,
-                        },
-                        {   
-                            "type": "text", 
-                            "text": SYSTEM_PROMPT
-                        },
-                    ],
-                }
-                
-        # 生成 user_input
-        user_input = self.processor.apply_chat_template(
-            [conversation],
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        results = dict()
-        inputs = {
-            'prompt': user_input,
-            'multi_modal_data': {'video': video_inputs}
-        }
-        results["index"] = real_index
-        results['input'] = inputs
-        return results
-
-    def __len__(self):
-        return len(self.meta)
-
-    def get_index(self, video_size, num_frames, st=0):
-        seg_size = max(0., float(video_size - 1) / num_frames)
-        max_frame = int(video_size) - 1
-        seq = []
-        # index from 1, must add 1
-        for i in range(num_frames):
-            start = int(np.round(seg_size * i))
-            # end = int(np.round(seg_size * (i + 1)))
-            idx = min(start, max_frame)
-            seq.append(idx+st)
-        return seq
-    
-def result_writer(indices_list: list, result_list: list, meta: pd.DataFrame, column):
-    flat_indices = []
-    for x in zip(indices_list):
-        flat_indices.extend(x)
-    flat_results = []
-    for x in zip(result_list):
-        flat_results.extend(x)
-    
-    flat_indices = np.array(flat_indices)
-    flat_results = np.array(flat_results)
-
-    unique_indices, unique_indices_idx = np.unique(flat_indices, return_index=True)
-    meta.loc[unique_indices, column[0]] = flat_results[unique_indices_idx]
-
-    meta = meta.loc[unique_indices]
-    return meta
-
-
-def worker_init_fn(worker_id):
-    # Set different seed for each worker
-    worker_seed = torch.initial_seed() % 2**32
-    np.random.seed(worker_seed)
-    # Prevent deadlocks by setting timeout
-    torch.set_num_threads(1)
-
-def main():
-    parser = argparse.ArgumentParser(description="SkyCaptioner-V1 vllm batch inference")
-    parser.add_argument("--input_csv", default="./examples/test.csv")
-    parser.add_argument("--out_csv", default="./examples/test_result.csv")
-    parser.add_argument("--bs", type=int, default=4)
-    parser.add_argument("--tp", type=int, default=1)
-    parser.add_argument("--model_path", required=True, type=str, help="skycaptioner-v1 model path")
-    args = parser.parse_args()
-    
-    dataset = VideoTextDataset(csv_path=args.input_csv, model_path=args.model_path)
-    dataloader = DataLoader(
-        dataset,
-        batch_size=args.bs,
-        num_workers=4,
-        worker_init_fn=worker_init_fn,
-        persistent_workers=True,
-        timeout=180,
-    )
-
-    sampling_params = SamplingParams(temperature=0.05, max_tokens=2048)
-    
-    llm = LLM(model=args.model_path,
-        gpu_memory_utilization=0.6, 
-        max_model_len=31920,
-        tensor_parallel_size=args.tp)
-    
-    indices_list = []
-    caption_save = []
-    for video_batch in tqdm(dataloader):
-        indices = video_batch["index"]
-        inputs = video_batch["input"]
-        batch_user_inputs = []
-        for prompt, video in zip(inputs['prompt'], inputs['multi_modal_data']['video'][0]):
-            usi={'prompt':prompt, 'multi_modal_data':{'video':video}}
-            batch_user_inputs.append(usi)
-        outputs = llm.generate(batch_user_inputs, sampling_params, use_tqdm=False)
-        struct_outputs = [output.outputs[0].text for output in outputs]
-
-        indices_list.extend(indices.tolist())
-        caption_save.extend(struct_outputs)
-    
-    meta_new = result_writer(indices_list, caption_save, dataset.meta, column=["structural_caption"])
-    meta_new.to_csv(args.out_csv, index=False)
-    print(f'Saved structural_caption to {args.out_csv}')
-
-if __name__ == '__main__':
-    main()

package/#Tools/SkyReels-V2/skyreels_v2_infer/__init__.py

@@ -1 +0,0 @@
-from .pipelines import DiffusionForcingPipeline

package/#Tools/SkyReels-V2/skyreels_v2_infer/distributed/xdit_context_parallel.py

@@ -1,286 +0,0 @@
-import numpy as np
-import torch
-import torch.amp as amp
-from torch.backends.cuda import sdp_kernel
-from xfuser.core.distributed import get_sequence_parallel_rank
-from xfuser.core.distributed import get_sequence_parallel_world_size
-from xfuser.core.distributed import get_sp_group
-from xfuser.core.long_ctx_attention import xFuserLongContextAttention
-
-from ..modules.transformer import sinusoidal_embedding_1d
-
-
-def pad_freqs(original_tensor, target_len):
-    seq_len, s1, s2 = original_tensor.shape
-    pad_size = target_len - seq_len
-    padding_tensor = torch.ones(pad_size, s1, s2, dtype=original_tensor.dtype, device=original_tensor.device)
-    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
-    return padded_tensor
-
-
-@amp.autocast("cuda", enabled=False)
-def rope_apply(x, grid_sizes, freqs):
-    """
-    x:          [B, L, N, C].
-    grid_sizes: [B, 3].
-    freqs:      [M, C // 2].
-    """
-    s, n, c = x.size(1), x.size(2), x.size(3) // 2
-    # split freqs
-    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
-
-    # loop over samples
-    output = []
-    grid = [grid_sizes.tolist()] * x.size(0)
-    for i, (f, h, w) in enumerate(grid):
-        seq_len = f * h * w
-
-        # precompute multipliers
-        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(s, n, -1, 2))
-        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,
-        ).reshape(seq_len, 1, -1)
-
-        # apply rotary embedding
-        sp_size = get_sequence_parallel_world_size()
-        sp_rank = get_sequence_parallel_rank()
-        freqs_i = pad_freqs(freqs_i, s * sp_size)
-        s_per_rank = s
-        freqs_i_rank = freqs_i[(sp_rank * s_per_rank) : ((sp_rank + 1) * s_per_rank), :, :]
-        x_i = torch.view_as_real(x_i * freqs_i_rank.cuda()).flatten(2)
-        x_i = torch.cat([x_i, x[i, s:]])
-
-        # append to collection
-        output.append(x_i)
-    return torch.stack(output).float()
-
-
-def broadcast_should_calc(should_calc: bool) -> bool:
-    import torch.distributed as dist
-
-    device = torch.cuda.current_device()
-    int_should_calc = 1 if should_calc else 0
-    tensor = torch.tensor([int_should_calc], device=device, dtype=torch.int8)
-    dist.broadcast(tensor, src=0)
-    should_calc = tensor.item() == 1
-    return should_calc
-
-
-def usp_dit_forward(self, x, t, context, clip_fea=None, y=None, fps=None):
-    """
-    x:              A list of videos each with shape [C, T, H, W].
-    t:              [B].
-    context:        A list of text embeddings each with shape [L, C].
-    """
-    if self.model_type == "i2v":
-        assert clip_fea is not None and y is not None
-    # params
-    device = self.patch_embedding.weight.device
-    if self.freqs.device != device:
-        self.freqs = self.freqs.to(device)
-
-    if y is not None:
-        x = torch.cat([x, y], dim=1)
-
-    # embeddings
-    x = self.patch_embedding(x)
-    grid_sizes = torch.tensor(x.shape[2:], dtype=torch.long)
-    x = x.flatten(2).transpose(1, 2)
-
-    if self.flag_causal_attention:
-        frame_num = grid_sizes[0]
-        height = grid_sizes[1]
-        width = grid_sizes[2]
-        block_num = frame_num // self.num_frame_per_block
-        range_tensor = torch.arange(block_num).view(-1, 1)
-        range_tensor = range_tensor.repeat(1, self.num_frame_per_block).flatten()
-        casual_mask = range_tensor.unsqueeze(0) <= range_tensor.unsqueeze(1)  # f, f
-        casual_mask = casual_mask.view(frame_num, 1, 1, frame_num, 1, 1).to(x.device)
-        casual_mask = casual_mask.repeat(1, height, width, 1, height, width)
-        casual_mask = casual_mask.reshape(frame_num * height * width, frame_num * height * width)
-        self.block_mask = casual_mask.unsqueeze(0).unsqueeze(0)
-
-    # time embeddings
-    with amp.autocast("cuda", dtype=torch.float32):
-        if t.dim() == 2:
-            b, f = t.shape
-            _flag_df = True
-        else:
-            _flag_df = False
-        e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(self.patch_embedding.weight.dtype)
-        )  # b, dim
-        e0 = self.time_projection(e).unflatten(1, (6, self.dim))  # b, 6, dim
-
-        if self.inject_sample_info:
-            fps = torch.tensor(fps, dtype=torch.long, device=device)
-
-            fps_emb = self.fps_embedding(fps).float()
-            if _flag_df:
-                e0 = e0 + self.fps_projection(fps_emb).unflatten(1, (6, self.dim)).repeat(t.shape[1], 1, 1)
-            else:
-                e0 = e0 + self.fps_projection(fps_emb).unflatten(1, (6, self.dim))
-
-        if _flag_df:
-            e = e.view(b, f, 1, 1, self.dim)
-            e0 = e0.view(b, f, 1, 1, 6, self.dim)
-            e = e.repeat(1, 1, grid_sizes[1], grid_sizes[2], 1).flatten(1, 3)
-            e0 = e0.repeat(1, 1, grid_sizes[1], grid_sizes[2], 1, 1).flatten(1, 3)
-            e0 = e0.transpose(1, 2).contiguous()
-
-        assert e.dtype == torch.float32 and e0.dtype == torch.float32
-
-    # context
-    context = self.text_embedding(context)
-
-    if clip_fea is not None:
-        context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
-        context = torch.concat([context_clip, context], dim=1)
-
-    # arguments
-    if e0.ndim == 4:
-        e0 = torch.chunk(e0, get_sequence_parallel_world_size(), dim=2)[get_sequence_parallel_rank()]
-    kwargs = dict(e=e0, grid_sizes=grid_sizes, freqs=self.freqs, context=context, block_mask=self.block_mask)
-
-    if self.enable_teacache:
-        modulated_inp = e0 if self.use_ref_steps else e
-        # teacache
-        if self.cnt % 2 == 0:  # even -> conditon
-            self.is_even = True
-            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
-                should_calc_even = True
-                self.accumulated_rel_l1_distance_even = 0
-            else:
-                rescale_func = np.poly1d(self.coefficients)
-                self.accumulated_rel_l1_distance_even += rescale_func(
-                    ((modulated_inp - self.previous_e0_even).abs().mean() / self.previous_e0_even.abs().mean())
-                    .cpu()
-                    .item()
-                )
-                if self.accumulated_rel_l1_distance_even < self.teacache_thresh:
-                    should_calc_even = False
-                else:
-                    should_calc_even = True
-                    self.accumulated_rel_l1_distance_even = 0
-            self.previous_e0_even = modulated_inp.clone()
-        else:  # odd -> unconditon
-            self.is_even = False
-            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
-                should_calc_odd = True
-                self.accumulated_rel_l1_distance_odd = 0
-            else:
-                rescale_func = np.poly1d(self.coefficients)
-                self.accumulated_rel_l1_distance_odd += rescale_func(
-                    ((modulated_inp - self.previous_e0_odd).abs().mean() / self.previous_e0_odd.abs().mean())
-                    .cpu()
-                    .item()
-                )
-                if self.accumulated_rel_l1_distance_odd < self.teacache_thresh:
-                    should_calc_odd = False
-                else:
-                    should_calc_odd = True
-                    self.accumulated_rel_l1_distance_odd = 0
-            self.previous_e0_odd = modulated_inp.clone()
-
-    x = torch.chunk(x, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
-    if self.enable_teacache:
-        if self.is_even:
-            should_calc_even = broadcast_should_calc(should_calc_even)
-            if not should_calc_even:
-                x += self.previous_residual_even
-            else:
-                ori_x = x.clone()
-                for block in self.blocks:
-                    x = block(x, **kwargs)
-                ori_x.mul_(-1)
-                ori_x.add_(x)
-                self.previous_residual_even = ori_x
-        else:
-            should_calc_odd = broadcast_should_calc(should_calc_odd)
-            if not should_calc_odd:
-                x += self.previous_residual_odd
-            else:
-                ori_x = x.clone()
-                for block in self.blocks:
-                    x = block(x, **kwargs)
-                ori_x.mul_(-1)
-                ori_x.add_(x)
-                self.previous_residual_odd = ori_x
-        self.cnt += 1
-        if self.cnt >= self.num_steps:
-            self.cnt = 0
-    else:
-        # Context Parallel
-        for block in self.blocks:
-            x = block(x, **kwargs)
-
-    # head
-    if e.ndim == 3:
-        e = torch.chunk(e, get_sequence_parallel_world_size(), dim=1)[get_sequence_parallel_rank()]
-    x = self.head(x, e)
-    # Context Parallel
-    x = get_sp_group().all_gather(x, dim=1)
-    # unpatchify
-    x = self.unpatchify(x, grid_sizes)
-    return x.float()
-
-
-def usp_attn_forward(self, x, grid_sizes, freqs, block_mask):
-
-    r"""
-    Args:
-        x(Tensor): Shape [B, L, num_heads, C / num_heads]
-        seq_lens(Tensor): Shape [B]
-        grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
-        freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
-    """
-    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
-    half_dtypes = (torch.float16, torch.bfloat16)
-
-    def half(x):
-        return x if x.dtype in half_dtypes else x.to(torch.bfloat16)
-
-    # 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
-
-    x = x.to(self.q.weight.dtype)
-    q, k, v = qkv_fn(x)
-
-    if not self._flag_ar_attention:
-        q = rope_apply(q, grid_sizes, freqs)
-        k = rope_apply(k, grid_sizes, freqs)
-    else:
-
-        q = rope_apply(q, grid_sizes, freqs)
-        k = rope_apply(k, grid_sizes, freqs)
-        q = q.to(torch.bfloat16)
-        k = k.to(torch.bfloat16)
-        v = v.to(torch.bfloat16)
-        # x = torch.nn.functional.scaled_dot_product_attention(
-        #     q.transpose(1, 2),
-        #     k.transpose(1, 2),
-        #     v.transpose(1, 2),
-        #    ).transpose(1, 2).contiguous()
-        with sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
-            x = (
-                torch.nn.functional.scaled_dot_product_attention(
-                    q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_mask=block_mask
-                )
-                .transpose(1, 2)
-                .contiguous()
-            )
-    x = xFuserLongContextAttention()(None, query=half(q), key=half(k), value=half(v), window_size=self.window_size)
-
-    # output
-    x = x.flatten(2)
-    x = self.o(x)
-    return x

package/#Tools/SkyReels-V2/skyreels_v2_infer/modules/__init__.py

@@ -1,69 +0,0 @@
-import gc
-import os
-
-import torch
-from safetensors.torch import load_file
-
-from .clip import CLIPModel
-from .t5 import T5EncoderModel
-from .transformer import WanModel
-from .vae import WanVAE
-
-
-def download_model(model_id):
-    if not os.path.exists(model_id):
-        from huggingface_hub import snapshot_download
-
-        model_id = snapshot_download(repo_id=model_id)
-    return model_id
-
-
-def get_vae(model_path, device="cuda", weight_dtype=torch.float32) -> WanVAE:
-    vae = WanVAE(model_path).to(device).to(weight_dtype)
-    vae.vae.requires_grad_(False)
-    vae.vae.eval()
-    gc.collect()
-    torch.cuda.empty_cache()
-    return vae
-
-
-def get_transformer(model_path, device="cuda", weight_dtype=torch.bfloat16) -> WanModel:
-    config_path = os.path.join(model_path, "config.json")
-    transformer = WanModel.from_config(config_path).to(weight_dtype).to(device)
-
-    for file in os.listdir(model_path):
-        if file.endswith(".safetensors"):
-            file_path = os.path.join(model_path, file)
-            state_dict = load_file(file_path)
-            transformer.load_state_dict(state_dict, strict=False)
-            del state_dict
-            gc.collect()
-            torch.cuda.empty_cache()
-
-    transformer.requires_grad_(False)
-    transformer.eval()
-    gc.collect()
-    torch.cuda.empty_cache()
-    return transformer
-
-
-def get_text_encoder(model_path, device="cuda", weight_dtype=torch.bfloat16) -> T5EncoderModel:
-    t5_model = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth")
-    tokenizer_path = os.path.join(model_path, "google", "umt5-xxl")
-    text_encoder = T5EncoderModel(checkpoint_path=t5_model, tokenizer_path=tokenizer_path).to(device).to(weight_dtype)
-    text_encoder.requires_grad_(False)
-    text_encoder.eval()
-    gc.collect()
-    torch.cuda.empty_cache()
-    return text_encoder
-
-
-def get_image_encoder(model_path, device="cuda", weight_dtype=torch.bfloat16) -> CLIPModel:
-    checkpoint_path = os.path.join(model_path, "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth")
-    tokenizer_path = os.path.join(model_path, "xlm-roberta-large")
-    image_enc = CLIPModel(checkpoint_path, tokenizer_path).to(weight_dtype).to(device)
-    image_enc.requires_grad_(False)
-    image_enc.eval()
-    gc.collect()
-    torch.cuda.empty_cache()
-    return image_enc

package/#Tools/SkyReels-V2/skyreels_v2_infer/modules/attention.py

@@ -1,179 +0,0 @@
-# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
-import torch
-
-try:
-    import flash_attn_interface
-
-    FLASH_ATTN_3_AVAILABLE = True
-except ModuleNotFoundError:
-    FLASH_ATTN_3_AVAILABLE = False
-
-try:
-    import flash_attn
-
-    FLASH_ATTN_2_AVAILABLE = True
-except ModuleNotFoundError:
-    FLASH_ATTN_2_AVAILABLE = False
-
-import warnings
-
-__all__ = [
-    "flash_attention",
-    "attention",
-]
-
-
-def flash_attention(
-    q,
-    k,
-    v,
-    q_lens=None,
-    k_lens=None,
-    dropout_p=0.0,
-    softmax_scale=None,
-    q_scale=None,
-    causal=False,
-    window_size=(-1, -1),
-    deterministic=False,
-    dtype=torch.bfloat16,
-    version=None,
-):
-    """
-    q:              [B, Lq, Nq, C1].
-    k:              [B, Lk, Nk, C1].
-    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
-    q_lens:         [B].
-    k_lens:         [B].
-    dropout_p:      float. Dropout probability.
-    softmax_scale:  float. The scaling of QK^T before applying softmax.
-    causal:         bool. Whether to apply causal attention mask.
-    window_size:    (left right). If not (-1, -1), apply sliding window local 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.
-    """
-    half_dtypes = (torch.float16, torch.bfloat16)
-    assert dtype in half_dtypes
-    assert q.device.type == "cuda" and q.size(-1) <= 256
-
-    # params
-    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
-
-    def half(x):
-        return x if x.dtype in half_dtypes else x.to(dtype)
-
-    # preprocess query
-
-    q = half(q.flatten(0, 1))
-    q_lens = torch.tensor([lq] * b, dtype=torch.int32).to(device=q.device, non_blocking=True)
-
-    # preprocess key, value
-
-    k = half(k.flatten(0, 1))
-    v = half(v.flatten(0, 1))
-    k_lens = torch.tensor([lk] * b, dtype=torch.int32).to(device=k.device, non_blocking=True)
-
-    q = q.to(v.dtype)
-    k = k.to(v.dtype)
-
-    if q_scale is not None:
-        q = q * q_scale
-
-    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
-        warnings.warn("Flash attention 3 is not available, use flash attention 2 instead.")
-
-    torch.cuda.nvtx.range_push(f"{list(q.shape)}-{list(k.shape)}-{list(v.shape)}-{q.dtype}-{k.dtype}-{v.dtype}")
-    # apply attention
-    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
-        # Note: dropout_p, window_size are not supported in FA3 now.
-        x = flash_attn_interface.flash_attn_varlen_func(
-            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_k=torch.cat([k_lens.new_zeros([1]), k_lens])
-            .cumsum(0, dtype=torch.int32)
-            .to(q.device, non_blocking=True),
-            seqused_q=None,
-            seqused_k=None,
-            max_seqlen_q=lq,
-            max_seqlen_k=lk,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            deterministic=deterministic,
-        )[0].unflatten(0, (b, lq))
-    else:
-        assert FLASH_ATTN_2_AVAILABLE
-        x = flash_attn.flash_attn_varlen_func(
-            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_k=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_k=lk,
-            dropout_p=dropout_p,
-            softmax_scale=softmax_scale,
-            causal=causal,
-            window_size=window_size,
-            deterministic=deterministic,
-        ).unflatten(0, (b, lq))
-    torch.cuda.nvtx.range_pop()
-
-    # output
-    return x
-
-
-def attention(
-    q,
-    k,
-    v,
-    q_lens=None,
-    k_lens=None,
-    dropout_p=0.0,
-    softmax_scale=None,
-    q_scale=None,
-    causal=False,
-    window_size=(-1, -1),
-    deterministic=False,
-    dtype=torch.bfloat16,
-    fa_version=None,
-):
-    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
-        return flash_attention(
-            q=q,
-            k=k,
-            v=v,
-            q_lens=q_lens,
-            k_lens=k_lens,
-            dropout_p=dropout_p,
-            softmax_scale=softmax_scale,
-            q_scale=q_scale,
-            causal=causal,
-            window_size=window_size,
-            deterministic=deterministic,
-            dtype=dtype,
-            version=fa_version,
-        )
-    else:
-        if q_lens is not None or k_lens is not None:
-            warnings.warn(
-                "Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance."
-            )
-        attn_mask = None
-
-        q = q.transpose(1, 2).to(dtype)
-        k = k.transpose(1, 2).to(dtype)
-        v = v.transpose(1, 2).to(dtype)
-
-        out = torch.nn.functional.scaled_dot_product_attention(
-            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p
-        )
-
-        out = out.transpose(1, 2).contiguous()
-        return out