ulysses enabling in native attention path

[sywangyi]

fix the corrupted output when enable ulysses with native attention. since native attention is widely use in no-cuda plaform. and ulysses does not rely on lse. so ulysses attention still could be used in native attention path.

• General functionalities: @sayakpaul @yiyixuxu @DN6

[sywangyi]

from diffusers.hooks.group_offloading import apply_group_offloading
from diffusers.utils import export_to_video


model_id="Wan-AI/Wan2.2-T2V-A14B-Diffusers"

def setup_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    device = torch.device(f"cuda:{dist.get_rank()}")
    torch.cuda.set_device(device)
    return device


def set_seed_for_all_ranks(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    generator = torch.Generator(device="cuda")
    generator.manual_seed(seed)
    return generator


device = setup_distributed()
generator = set_seed_for_all_ranks(42)
onload_device = device
offload_device = torch.device("cpu")

vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
# group-offloading
pipe = WanPipeline.from_pretrained(
    model_id,
    vae=vae,
    torch_dtype=torch.bfloat16,
)
ulysses_degree = torch.distributed.get_world_size()
pipe.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=ulysses_degree))
pipe.transformer_2.enable_parallelism(config=ContextParallelConfig(ulysses_degree=ulysses_degree))
pipe.to("cuda")

pipe.vae.enable_tiling(tile_sample_min_height=480,tile_sample_min_width=960,tile_sample_stride_height=352,tile_sample_stride_width=640)
height = 704
width = 1280
num_frames = 24
num_inference_steps = 50
guidance_scale = 5.0


prompt = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩>残留，丑陋的，残缺的>，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手>指融合，静止不动的画面，杂>乱的背景，三条腿，背>景人很多，倒着走"

output = pipe(
    prompt=prompt,
    negative_prompt=negative_prompt,
    height=height,
    width=width,
    num_frames=num_frames,
    guidance_scale=guidance_scale,
    num_inference_steps=num_inference_steps,
    generator=generator,
).frames[0]
if torch.distributed.get_rank() == 0:
    export_to_video(output, "5bit2v_output_bad.mp4", fps=24)
if dist.is_initialized():
    torch.distributed.destroy_process_group()

[sywangyi]

torchrun --nproc-per-node 4 test.py
corrupted video output:

video output after fix:

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DefTruth]

Also worked for me, thanks

[DefTruth]

Here, 'supports_context_parallel=True' should be also added to register @sywangyi @sayakpaul

[sywangyi]

thanks, added

[sayakpaul]

@DefTruth that should also fix #12446 (comment) right? Could you give this a check?

[DefTruth]

@DefTruth that should also fix #12446 (comment) right? Could you give this a check?

confirm fixed