Add contrib models: LongCat-Image-Edit, Qwen-Image-Edit, Wan2.2-TI2V-5B

[whn09]

Description

Add NxDI contrib models for three diffusion/generation models: LongCat-Image-Edit (FLUX-style image editing), Qwen-Image-Edit (Qwen2.5-VL-based image editing), and Wan2.2-TI2V-5B (text/image-to-video generation). All three run on trn2.48xlarge with CFG Parallel or Context Parallel strategies and include NKI Flash Attention kernels for Trainium2.

Model Information

Model Name: LongCat-Image-Edit, Qwen-Image-Edit, Wan2.2-TI2V-5B

Model Architecture: Multi-component diffusion models (Vision Encoder + Language Model + DiT/FLUX Transformer + VAE)

Purpose: Image editing (LongCat, Qwen), text/image-to-video generation (Wan2.2)

Checklist

Required Components

• Accuracy Test (test/integration/test_model.py)

  • LongCat-Image-Edit: Full end-to-end CFG inference verified (1024x1024, 50 steps, 18.17s)
  • Qwen-Image-Edit: Full end-to-end V3 CFG inference verified (1024x1024, 50 steps, 37.62s total)
  • Wan2.2-TI2V-5B: Full end-to-end T2V inference verified (512x384, 81 frames, 15.77s CFG)

• README.md with the following sections:

  • Usage Example: Complete compile + inference workflows for each model
  • Compatibility Matrix: Tested on Trn2 (trn2.48xlarge) with Neuron SDK 2.22+
  • Example Checkpoints: Links to HuggingFace model hubs
  • Testing Instructions: pytest and manual test commands

• Source Code (src/)

  • Compilation scripts for each component (transformer, VAE, text/vision encoders)
  • Inference runner scripts with CFG Parallel / Context Parallel support
  • NKI Flash Attention kernels, TP sharding utilities, Neuron-compatible RoPE

Optional Components

• Unit Tests — Integration tests in test/integration/ for each model

Folder Structure

contrib/models/LongCat-Image-Edit/
  README.md
  requirements.txt
  assets/test.png
  src/
    run_longcat_image_edit.py           # Main inference script
    neuron_commons.py                   # NKI attention, encoder wrappers
    neuron_parallel_utils.py            # FLUX TP sharding
    neuron_rope.py                      # 3-axis RoPE
    compile_transformer.py              # FLUX transformer (TP=4, CP=2)
    compile_transformer_cfg.py          # FLUX transformer (TP=4, DP=2, CFG)
    compile_vae.py                      # AutoencoderKL (1024x1024)
    compile_vision_encoder.py           # Qwen2.5-VL ViT (TP=4)
    compile_language_model.py           # Qwen2.5-VL LM (TP=4)
    compile.sh, cache_hf_model.py, setup_nvme.sh
  test/integration/test_model.py

contrib/models/Qwen-Image-Edit/
  README.md
  requirements.txt
  assets/image1.png, image2.png
  src/
    run_qwen_image_edit.py              # Main inference (6 compilation APIs)
    neuron_commons.py                   # SDPA implementations
    autoencoder_kl_qwenimage_neuron.py  # Neuron-compatible 3D VAE
    compile_transformer_v3_cfg.py       # V3 CFG (fastest, recommended)
    compile_transformer_v3_cp.py        # V3 Context Parallel
    compile_transformer_v{1,2}*.py      # V1/V2 variants
    compile_vae.py, compile_text_encoder.py, ...
  test/integration/ (6 test files)

contrib/models/Wan2.2-TI2V-5B/
  README.md
  requirements.txt
  assets/cat.png
  src/
    run_wan2.2_ti2v.py                  # T2V and I2V inference
    neuron_commons.py                   # Rolling cache decoder, attention utils
    compile_transformer.py              # DiT transformer (TP=4, CFG/CP)
    compile_decoder_rolling.py          # Stateful rolling cache VAE decoder
    compile_text_encoder.py             # UMT5 (ModelBuilder, TP=4)
    compile.sh, cache_hf_model.py, setup_nvme.sh
  test/integration/test_model.py

Testing

How did you test this change?

All models compiled and tested end-to-end on trn2.48xlarge (Neuron SDK 2.22+, PyTorch 2.9).

LongCat-Image-Edit:

Config	Total Time	Per Step	Notes
CFG Parallel (TP=4, DP=2)	18.17s	0.36s	Recommended
H100 reference (single GPU)	23.61s	0.47s	Baseline

Qwen-Image-Edit:

Version	Per Step	Total (50 steps)	Notes
V3 CFG (TP=4, DP=2, NKI Flash)	~0.75s	37.62s	Fastest, recommended
V3 CP (TP=4, CP=2, NKI Flash)	~0.77s	38.67s	Context Parallel
V1 Flash / V2 Flash (TP=8)	~1.2s	~60s	NKI kernel
V1 baseline (TP=8)	~2.4s	~120s	Standard SDPA

Wan2.2-TI2V-5B:

Resolution	Frames	CFG (s)	H100 (s)
512x384	81	15.77	16.13
640x480	121	45.15	39.67
1280x704	81	155.06	87.66

Compatibility

Tested with:

• Neuron SDK Version(s): 2.22+
• Instance Type(s): Trn2 (trn2.48xlarge)
• PyTorch Version: 2.9
• Python Version: 3.12

Additional Information

Key design decisions shared across all models:

• CFG Parallel vs Context Parallel: CFG batches cond+uncond into single forward (no K/V all-gather), ~10% faster for guided generation. CP is better for no-guidance scenarios.
• NKI Flash Attention: Custom NKI kernels for Trainium2 attention computation, requires XLA_DISABLE_FUNCTIONALIZATION=1.
• Neuron-compatible RoPE: Uses (cos, sin) tuples instead of complex numbers (C64 unsupported on Neuron).

LongCat-specific:

• M-RoPE from Qwen2.5-VL for 3D position IDs
• VAE compiled at full 1024x1024 (no tiling) to avoid seam artifacts

Qwen-Image-Edit-specific:

• 6 compilation APIs with progressive optimization (V1→V3)
• Modulation layer sharding reduces memory from ~17GB to ~5.2GB per shard
• 3D causal VAE with tiled spatial processing

Wan2.2-specific:

• Stateful rolling cache: 34 feat_cache tensors on-device (HBM) via input-output aliasing, eliminates ~960MB roundtrip per call
• VAE encoder on CPU due to Neuron compiler bug (NCC_IBIR158)
• Multi-resolution support up to 720P with tiled spatial decode

Related Issues

N/A

vLLM Integration

• This model/feature is intended for use with vLLM
• Documentation includes vLLM registration instructions

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included