From fa4e3659c79aa03430043cc2e57064920f22c4c2 Mon Sep 17 00:00:00 2001 From: William Lin Date: Wed, 28 Oct 2026 02:17:34 +0000 Subject: [PATCH] [new-model] Re-add PAN2 text-to-video and image-to-video (reverts #2025) This reverts commit 6743f4b4b, restoring #1961. --- docs-site/content/cookbook/pan2/pan2.mdx | 7 + docs-site/lib/cookbook-pages.ts | 12 + docs-site/navigation.json | 23 + docs-site/tests/cookbook.spec.ts | 26 +- docs/assets/cookbook-recipes.json | 27 + docs/assets/logos/SOURCES.md | 1 + docs/assets/logos/ifm.webp | Bin 0 -> 4346 bytes docs/cookbook/index.md | 1 + docs/cookbook/pan2.md | 76 +++ docs/cookbook/recipes/pan2.yaml | 38 ++ docs/cookbook/serving-coverage.md | 1 + .../inference_schema_parity_inventory.yaml | 12 + docs/generate_examples.py | 1 + docs/inference/pan2.md | 86 +++ docs/inference/support_matrix.md | 1 + examples/inference/basic/basic_pan2.py | 40 ++ examples/serving/README.md | 1 + examples/serving/openai_pan2.yaml | 36 ++ fastvideo/configs/models/dits/__init__.py | 3 +- fastvideo/configs/models/dits/pan2.py | 86 +++ fastvideo/configs/models/encoders/__init__.py | 3 + fastvideo/configs/models/encoders/qwen3_5.py | 119 ++++ fastvideo/configs/pipelines/__init__.py | 3 +- fastvideo/configs/pipelines/pan2.py | 46 ++ fastvideo/models/dits/pan2.py | 398 ++++++++++++ fastvideo/models/encoders/qwen3_5.py | 574 ++++++++++++++++++ fastvideo/models/registry.py | 2 + fastvideo/pipelines/basic/pan2/__init__.py | 1 + .../pipelines/basic/pan2/pan2_pipeline.py | 49 ++ fastvideo/pipelines/basic/pan2/presets.py | 48 ++ fastvideo/pipelines/basic/pan2/stages.py | 210 +++++++ fastvideo/pipelines/stages/denoising.py | 11 + .../pipelines/stages/latent_preparation.py | 9 +- fastvideo/registry.py | 21 + .../test_serving_family_baselines.py | 28 + .../stages/test_pan2_first_latent_cfg.py | 57 ++ .../stages/test_pan2_latent_preparation.py | 67 ++ .../encoders/test_pan2_qwen3_5_parity.py | 339 +++++++++++ tests/local_tests/pan2/PORT_STATUS.md | 71 +++ tests/local_tests/pan2/README.md | 67 ++ tests/local_tests/pan2/reference.py | 302 +++++++++ .../pipelines/test_pan2_pipeline_parity.py | 159 +++++ .../schedulers/test_pan2_scheduler_parity.py | 200 ++++++ tests/local_tests/test_cookbook_serving.py | 20 + .../test_pan2_transformer_parity.py | 328 ++++++++++ .../local_tests/vaes/test_pan2_vae_parity.py | 387 ++++++++++++ 46 files changed, 3990 insertions(+), 7 deletions(-) create mode 100644 docs-site/content/cookbook/pan2/pan2.mdx create mode 100644 docs/assets/logos/ifm.webp create mode 100644 docs/cookbook/pan2.md create mode 100644 docs/cookbook/recipes/pan2.yaml create mode 100644 docs/inference/pan2.md create mode 100644 examples/inference/basic/basic_pan2.py create mode 100644 examples/serving/openai_pan2.yaml create mode 100644 fastvideo/configs/models/dits/pan2.py create mode 100644 fastvideo/configs/models/encoders/qwen3_5.py create mode 100644 fastvideo/configs/pipelines/pan2.py create mode 100644 fastvideo/models/dits/pan2.py create mode 100644 fastvideo/models/encoders/qwen3_5.py create mode 100644 fastvideo/pipelines/basic/pan2/__init__.py create mode 100644 fastvideo/pipelines/basic/pan2/pan2_pipeline.py create mode 100644 fastvideo/pipelines/basic/pan2/presets.py create mode 100644 fastvideo/pipelines/basic/pan2/stages.py create mode 100644 fastvideo/tests/stages/test_pan2_first_latent_cfg.py create mode 100644 fastvideo/tests/stages/test_pan2_latent_preparation.py create mode 100644 tests/local_tests/encoders/test_pan2_qwen3_5_parity.py create mode 100644 tests/local_tests/pan2/PORT_STATUS.md create mode 100644 tests/local_tests/pan2/README.md create mode 100644 tests/local_tests/pan2/reference.py create mode 100644 tests/local_tests/pipelines/test_pan2_pipeline_parity.py create mode 100644 tests/local_tests/schedulers/test_pan2_scheduler_parity.py create mode 100644 tests/local_tests/transformers/test_pan2_transformer_parity.py create mode 100644 tests/local_tests/vaes/test_pan2_vae_parity.py diff --git a/docs-site/content/cookbook/pan2/pan2.mdx b/docs-site/content/cookbook/pan2/pan2.mdx new file mode 100644 index 000000000..07836529a --- /dev/null +++ b/docs-site/content/cookbook/pan2/pan2.mdx @@ -0,0 +1,7 @@ +--- +title: PAN 2 +description: Configure text-to-video and first-frame image-to-video generation with PAN 2. +mode: frame +--- + + diff --git a/docs-site/lib/cookbook-pages.ts b/docs-site/lib/cookbook-pages.ts index af4e8e077..650f9c460 100644 --- a/docs-site/lib/cookbook-pages.ts +++ b/docs-site/lib/cookbook-pages.ts @@ -75,6 +75,15 @@ export const recipePages = [ description: "Turn a reference image into video with Wan2.1.", tags: ["Image to video", "14B", "480p"], }, + { + id: "pan2", + family: "PAN 2", + title: "PAN 2", + mark: "PAN", + path: "/cookbook/pan2/pan2", + description: "Text-to-video and first-frame image-to-video from the 29B PAN 2 model.", + tags: ["Text or image to video", "29B"], + }, { id: "ltx2-distilled", family: "LTX", @@ -120,6 +129,7 @@ const familyInfo: Record = { "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": "fastwan21", "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": "wan21-i2v", "FastVideo/LTX2-Distilled-Diffusers": "ltx2-distilled", + "IFM/PAN2": "pan2", }; const hiddenBackends = new Set(["fasth3-preview"]); @@ -185,6 +196,7 @@ const mappedModels = [...new Set(sourceRecipes.map((recipe) => recipe.model))].m kandinsky: "Kandinsky Lab", "black-forest-labs": "Black Forest Labs", qwen: "Qwen", + ifm: "IFM", stabilityai: "Stability AI", nvidia: "NVIDIA", }; diff --git a/docs-site/navigation.json b/docs-site/navigation.json index 489ed306f..68bd730b3 100644 --- a/docs-site/navigation.json +++ b/docs-site/navigation.json @@ -127,6 +127,10 @@ ], "display": "group" }, + { + "label": "PAN 2 text/image to video", + "href": "/inference/pan2" + }, { "label": "Examples", "href": "/inference/examples/examples_inference_index" @@ -249,6 +253,25 @@ } ] }, + { + "label": "PAN 2", + "display": "group", + "collapsed": true, + "items": [ + { + "label": "PAN 2", + "href": "/cookbook/pan2/pan2" + }, + { + "label": "Family guide", + "href": "/cookbook/pan2" + }, + { + "label": "Model reference", + "href": "/inference/pan2" + } + ] + }, { "label": "FLUX", "href": "/cookbook/flux" diff --git a/docs-site/tests/cookbook.spec.ts b/docs-site/tests/cookbook.spec.ts index d1906df92..e7da547ac 100644 --- a/docs-site/tests/cookbook.spec.ts +++ b/docs-site/tests/cookbook.spec.ts @@ -223,7 +223,7 @@ test("inherited values, explicit false, zero and null retain their meanings", as test("catalog discovery opens dedicated model recipes and filters by family and task", async ({ page }) => { await page.goto("/FastVideo/cookbook"); const cards = page.locator(".recipe-card"); - await expect(cards).toHaveCount(38); + await expect(cards).toHaveCount(39); await expect(cards.first()).toContainText("FastH3 Omni v1"); await page.getByLabel("Model family", { exact: true }).selectOption("Wan"); await page.getByLabel("Search recipes").fill("Wan2.1 I2V"); @@ -276,7 +276,7 @@ test("output tabs support keyboard navigation and precision chips use the backen await page.getByLabel("Search recipes").fill("does not exist"); await expect(page.getByRole("heading", { name: "No matching recipes" })).toBeVisible(); await page.getByRole("button", { name: "Clear filters" }).click(); - await expect(page.locator(".recipe-card")).toHaveCount(38); + await expect(page.locator(".recipe-card")).toHaveCount(39); }); test("cookbook cards and controls remain readable in dark mode and reflow at 320px", async ({ page }) => { @@ -471,6 +471,26 @@ test("source-backed serving pages expose their actual server settings", async ({ await expect(page.getByRole("button", { name: "Copy config + command", exact: true })).toBeEnabled(); }); +test("PAN 2 has a family card with the IFM logo and an editable serving deployment", async ({ page }) => { + await page.goto("/FastVideo/cookbook"); + await page.getByLabel("Search recipes").fill("PAN 2"); + const pan2 = page.locator(".recipe-card"); + await expect(pan2).toHaveCount(1); + await expect(pan2.locator("img")).toHaveAttribute("src", /\/ifm\.webp$/); + await expect(pan2.locator(".recipe-card-kicker")).toHaveText("IFM"); + await pan2.click(); + await expect(page).toHaveURL(/cookbook\/pan2\/pan2$/); + await expect(page.locator(".recipe-matrix")).toBeVisible({ timeout: 30000 }); + await expect(page.getByRole("heading", { level: 1 })).toHaveText("PAN 2"); + await settings(page, "Serving"); + await page.getByLabel("Port", { exact: true }).fill("9010"); + await output(page, "config.yaml"); + await expect(page.getByLabel("Generated YAML", { exact: true })).toContainText(/"port": 9010/); + await expect(page.getByLabel("Generated YAML", { exact: true })).toContainText(/"model_path": "IFM\/PAN2"/); + await expect(page.getByLabel("Generated YAML", { exact: true })).toContainText(/"dit_layerwise": false/); + await expect(page.getByRole("button", { name: "Download YAML", exact: false })).toBeEnabled(); +}); + test("source model hardware preferences survive hydration and reload", async ({ page }) => { const errors: string[] = []; page.on("pageerror", (error) => errors.push(error.message)); @@ -789,7 +809,7 @@ test("MLX Omni discovery and its original route expose the MLX launch command", test("category counts match visible cards and exclude retired Preview", async ({ page }) => { await page.goto("/FastVideo/cookbook"); - await expect(page.locator(".recipe-card")).toHaveCount(38); + await expect(page.locator(".recipe-card")).toHaveCount(39); for (const category of ["All models", "Video", "Image", "Audio", "World"]) { const button = page.getByRole("button", { name: category, exact: true }); const count = Number(await button.locator(".recipe-filter-count").innerText()); diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 0c13d83fd..b58221bf6 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -381,6 +381,33 @@ "expected_artifact": "Upscaled MP4 under outputs_video/kandinsky6_sr/", "related": ["kandinsky6-vsr", "kandinsky6-ti2va-piflow"] }, + { + "id": "pan2-ti2v", + "family": "pan2", + "stage": "inference", + "task": "Text or image to video", + "label": "PAN 2", + "summary": "Generate a 13-second, 704 × 1248 video at 24 fps from text, or set IMAGE_PATH in the maintained example to condition on a first frame.", + "model": "IFM/PAN2", + "source": "examples/inference/basic/basic_pan2.py", + "serving": { + "source": "examples/serving/openai_pan2.yaml", + "install": "UV_TORCH_BACKEND=cu130 uv pip install -e .", + "task": "Text to video", + "limitations": [ + "This server walkthrough sends a text prompt only. The server API also accepts an image reference (image_reference or input_reference) and then runs image-to-video from that first frame, but the checked-in clients do not upload one." + ] + }, + "command": "python examples/inference/basic/basic_pan2.py", + "gpu_types": ["NVIDIA"], + "hardware": {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"}, + "evidence": "Source-backed", + "expected_artifact": "MP4 video under video_samples_pan2/", + "limitations": [ + "The full IFM/PAN2 checkpoint is not public yet; numerical parity is checked with a public random-weight checkpoint.", + "The example keeps the 29B DiT resident in bf16 and offloads only the text encoder. GPU memory requirements have not been established for this recipe; set PAN2_NUM_GPUS to shard video tokens with sequence parallelism, which does not shard the DiT weights." + ] + }, { "id": "flux2-klein-t2i", "family": "flux", diff --git a/docs/assets/logos/SOURCES.md b/docs/assets/logos/SOURCES.md index 0a2a1f9c5..9e6f78f82 100644 --- a/docs/assets/logos/SOURCES.md +++ b/docs/assets/logos/SOURCES.md @@ -13,6 +13,7 @@ plain typographic tile instead — that tile is a UI placeholder, not a logo. | `nvidia.webp` | [Official NVIDIA Hugging Face organization avatar](https://huggingface.co/nvidia) | Cosmos and GEN3C cards, Cosmos page header | | `kandinsky.webp` | [Official Kandinsky Lab Hugging Face organization avatar](https://huggingface.co/kandinskylab) | Kandinsky 5 and Kandinsky 6 cards and page headers | | `black-forest-labs.webp` | [Official Black Forest Labs Hugging Face organization avatar](https://huggingface.co/black-forest-labs) | FLUX family card and page header | +| `ifm.webp` | [Official Institute of Foundation Models (IFM) Hugging Face organization avatar](https://huggingface.co/IFM) | PAN 2 family card and page header | | `qwen.webp` | [Official Qwen Hugging Face organization avatar](https://huggingface.co/Qwen) | Qwen-Image-2.1 family card and page header | | `minimax.webp` | [Official MiniMax Hugging Face organization avatar](https://huggingface.co/MiniMaxAI) | MiniMax H3 family card and page header | | `pika.webp` | [Official Pika-Labs Hugging Face organization avatar](https://huggingface.co/Pika-Labs) | Pika H3 Flex card and page header | diff --git a/docs/assets/logos/ifm.webp b/docs/assets/logos/ifm.webp new file mode 100644 index 0000000000000000000000000000000000000000..fe4602a9f7e18ab2b61ad9638fa33432e6cc649c GIT binary patch literal 4346 zcmVphx(QvSR8`PUER|6{*)Ij3}2{GUSqhkauEZ}v~_|LwoMe(e0C z{vW{Ki+|Jlfqo_Yd;4GX|JzUO|Fe8-Km+pM>L28PI6uPrIsY-}U)b;Tp5PzJ|D=Dm z`+(;#^uJ&a=ilA`>HmxORsL80hxq^WKeXP}zd!#8|5N;r z|NG4LFZ*lkoBLV+)aVI>rUeIvq{0R!>KV8ylP)Ief?qP+%k0k?a`dFy4n63Y{=Fa+ zvM0bE`z7-)EWEs^BjG6>p-`_aaro|`Nlt)B(P6Ss;CM8s{HwT+>j6+&z##%>qGfs1 z2e1Wy%qAnk3v4gy>fDuZD8kt#TFh3BUZ*=@lECcsl7DxD3bIVQvkk6kF|5d543UjPvLkUH25h;`*b%dab8CR*lqM*8PCGeUhr`Hn*&>sK}d4 zz+3AZ%FV_!jUPp*HJTN}t0}kSPydNB0@M)GBuTGrY%?g#Vsh;ir}InQ#Q@n{v;*aN z_!rkgOjROVh9r@?{q`-BupLT*Wa_u<$M8x5i2=9~0pp*J$XvFekJQBt4D}13)>-~u z_t#e)fv3VQ`u~^@jn{f4XGy?l>ez%*-Mh^`!k^&rqbdoGN0*XUrCGo7?=`MF8P=vl^ z<)-52y9c@`lD=AQs3r3+7ytnN|JH;6ycFpL%zyDM9TPc8tiREFqHfB`+WpMFAb`3) z)?sXEfx+Cg{@Uu49>LaTwGZ2K%65z3JvqR7j;F8iLd}Ffjan_3d>bK`j_r<)z9#DI z@vH;AsyUZxP-3B;Eu@kfT;R-Gqf&xGgS?w*4jns1pYcg!S0ZaDu5-)b@%1; zaC}a@1Y4?+Hp>1PDMloPamFR|3q1r2!E>Y-AyM(g{Xd{bn=n;Q)ZAZr#$l_hviEER zPvKN1n)BS{V>QWVqW9|L*&@FN%m!sdO%EMql%=x(5G2@^*KQD4D(5Lg=`^&mX}qJ% z;DdS{X7Bt*;ewt~?2lr`UbM&ke>INM4qMLz70X6BG=goV#L15~ICdtpnL8e3{$rrX zpA2;o(H}+T?U0c6SjVD6SwA{1rt<2jUzFL_d!k)Z0^!8KroM;Ka@=UO) zHPHjgrd%B2v5_zrc!qejoFI*O5mx(=U!4ys)YAhuQ;$5k^w#)W+oXd$=Re9k)cDFb zSV0!Nfj$StSawWn5@xerPR%v4M$o{-nFRc`p9I!Z9sSi0n+4Qj8cmh?vbwe-} z^R5JdlwE#0!2Cb8M)&d~`+ITPh0R%nw@OB`g;R6Uo`t>LJyT*B7@7hjqz%%{`^wJLi03?fVE{(&?Vs}QE{%O=&*;Rq?ck?BPK!E39A$%FhrN4f&LD!>*S zs&jd}N{aMT$g0U1y;Q*XS8^iu<1a2}Hxl;x{YfF)>95`R0xLh^&Q?WO4bL`@ilxB7 zuIXq(`v!;Cc+N(3&^!8bxUWNAp)PvKYxfHm)V~o_ z<_JsWFA!ci%duS*Mcbr+iTd09W#M?Irx*_9Buv^Gzw;QsA=0)mJwIOS-13?=Xa$jU zCm-QU)wejKu$q#s2}`b?Wx2 zo*Di(fc)$$3%3Z*R-cOVj;u%_pyE7vO`5oO)fr2icR#~AkQb@-g^B1}*sD&$7$l7h zlE&k2x9}yebIt5u`0|Ru5QGmA?K>rfVUk>^a|vfBQ^hgwJr$MKK&IQi^<2Su9RHLr zU#=~Mtr$jJxWc-?B?$r%tiF)BCzL@uZFYn z72;LHtsg)ZniyFjci?QOyQwMQ(NSwqq6*hNczmDXGe=a061tU{=n6n%U&_ca63$8! znKkbYkdC2RICEg)#Lv5WiQv3mBN@dZ*KB!4EDDN{5$*GgV+pEWYg``_x^nWr%i=>E z?O0LC!y=~f5*e%AeEZIFDs^O%6nurg=dar^vzNEqg&O(zib4lYsm=|_lF70Fgq z^{RvM{diKP=$$AjmJGt;pEK;_<1@{UV9@8q(=2zbKtb!Q+cx}=ovg9Vj#vcY)MXH8 zSWP$0q0pA;;hzaVooQ4&60aH4(NCq3-*Sosq9`T}BcC_lFP7uVsU&&LqlL6{TU}wx zDCT*ae~$W=q@cs%EnQDnEXs@5h0{GcitnX9f83h8iG15zXaR$aS&u_qgVFu4dUuA) z6;yz^>IzpTT8Nf;IPA3&SCXJzpkLOTx;3rRJ6AiGV!*OKlKqw8{mz#+ttnSmE!<4i zhR=eNpsm6DP77nar_v3yD=mRtT=~-H8Vp5I20&OH`H|e70r_~sV1E`b%lnGLt`&!< z*h42-2^}~}sHn6-FJNG2fux8$oqB$ZZbf6n?NET*@?lbw?~co!tdSnxv;wXEz~_)P zW79aSgI@P7flO{AUT-+a{)(utpKo?{#ylFj?R@YfsaxTMUr(tdu`*ULBMI&%qW)2V zdcVn6Zc5T%h7FfibEMfN^oGSMP(LI@3dJZZ2?7jjmVF)KuQyHc)eTaQ;jm^clfcoo z#3tRY@!@}FmuiaKG%JC+bE1U$CcY30+wx~`J`zn)r_E@@c{Y@F=Y+Lt!_qrv^WOO- z_ANk|`&*@g2FT9~!`p|BzvuvF=t1Qx-)*ai-|Q1;F~`H_Z#<8;p@#7&=KJ%?GjE^z z35NehI2F`XWSr6LLCId@A#1z{^pFgWaD{$j<2!R{)PejFZBrLK_;&UQu=_6LL}@UU zhCNyqvUX)z()8_n=}ro7b_P|2NvCD~UVC)6v%o$&+t}RTH@y=T`7B*o&h`Yoj8GbM zC0cWCp!j_Qkvlu>(4#l570Sg0zmg_v#9%=NN6f;vN^P0#;eEn&7WZ~}Dj$*`REwyT z$H$`XWPwKHMxE32P#DUlZ)@~P*4gdwD?rY?J*Z!>Vl<}GYF9P2qq$^bkkn5TBF$NN zsuAmnjoQG`MH8_%1;NB2=Rf<)sP6Kwiyp@sX`0_2N1M&lwd|3*Ns3kcCSz`krkjJj z5bH7fk0$r{4y=E&1`iW}C5x*-_2qG2O?8qi-;urbo>-IkFB^)t4CIf}BTAiYo@{4F zou4o;P5hLDa3E-nR|5Z{lbp*Km~BGZ&M7Y(?b_8VW${V=~^i8%5x87!;f`yYHZ z`_$SOD(~T;5Uq+vhDjmT8Eb}#yl8{7h_9cKOa~iUx2&+Z1V-WO8MhF-CMCVP``6o` z8c>y-dMLBZc`@_@5QwVLrs*(-)&O6ydU=_ zfV?m_Ei%8P7xx0-LqN0A@@5b&V-5mN@~g8%zh=2e-lS^i37sx}QNx-sQ)UZ*0Bs-N z?b+-H)1%J7$rU%Gxa1Q{2lNII+|<0E2fa@ptIs2K^V#y~-dHM_=hUs?D4#`G5(bS0 z0}nww$kzwolkqb&;9l)oLw7Gw^c-gh(_TcwJ?jS!p;TOrb9n3ezzVc@04@nJ3@tZv zwNxW?59=@n*8GL0l=;t^|JZ+&*8p?`RHrqyANO*@uCw8F!(cV)b*SqP+~@c_*G4@q zM}1_bb>*3P?%-Pdj!yn4ba^8lus7k$YTAWfRJOV3hYu>y^E=GIE-hGcBI8j6FQOKp6=9MhA_; zI&w7pK6iU7pj}!5^(Ox2Uq%BcW>Muja*5Azy;>u=7qI_zy-fS|%&FVPU1{>vm$6@6 z_K5Wud?Vg&HGYT7RXRH;V3^(bG?vUe72cp|?uWzVRNJKdnMn!Qo)9?JG2g$!=w0+; zM-wUI60*Axg-&A~H~k`C_z2Ru4e&Ng+k>_lyLyo(<2`rF0U9`iD)V2dN~SL5Zf|@P zM#vvtPY#mP>Y_iQ6zE+rcGAxSc+{tWZLD=I&$b3Zl zhe@NBN}rbEd&X=`HOZ|cR_Uf>%_#S|$r0@PuL!-WT~RSqZ>>>ZkNt%dE|pJxVyyvT oPTuQ4FQ8O1RkC{Y#&F4QY$oCf^=C#kA+|SW4;DpKzyJUM0J}7{{{R30 literal 0 HcmV?d00001 diff --git a/docs/cookbook/index.md b/docs/cookbook/index.md index 10a75432c..f1cb374b1 100644 --- a/docs/cookbook/index.md +++ b/docs/cookbook/index.md @@ -6,6 +6,7 @@ The interactive catalog groups models by family and generates YAML for supported - [FastH3](minimax-h3.md) - [Wan and FastMetal](wan.md) - [LTX](ltx.md) +- [PAN 2](pan2.md) - [FLUX](flux.md) - [Qwen-Image-2.1](qwen-image21.md) - [Stable Diffusion](stable-diffusion.md) diff --git a/docs/cookbook/pan2.md b/docs/cookbook/pan2.md new file mode 100644 index 000000000..c19f53e58 --- /dev/null +++ b/docs/cookbook/pan2.md @@ -0,0 +1,76 @@ +# PAN 2 recipes + + + +Generate a video from text, or from a first frame, with `IFM/PAN2`. One pipeline +serves both modes: leave out the first frame for text-to-video and supply one +for image-to-video. The defaults are 704×1248, 313 frames (13 s at 24 fps), +50 steps and guidance 3.0. See [PAN 2 Text/Image to Video](../inference/pan2.md) +for the model, sampler and parallelism details. + +## Run on a GPU + +Follow the [CUDA installation guide](../getting_started/installation/gpu.md), then run +from the repository root. The full `IFM/PAN2` checkpoint is not public yet; pass a +local directory or private repo you can access with `PAN2_MODEL_PATH`. + +```bash +python examples/inference/basic/basic_pan2.py +``` + +Set `IMAGE_PATH` in the script to a first frame for image-to-video. With more +GPUs, `PAN2_NUM_GPUS=4` shards the video tokens with sequence parallelism; the +count must divide PAN 2's 32 attention heads. Each rank still holds the full +29B DiT, so extra GPUs speed up generation but do not reduce per-device weight +memory. See [CPU offloading](../inference/offloading.md) for devices that cannot +hold the DiT. + +## Serve with the OpenAI-compatible API + +`examples/serving/openai_pan2.yaml` serves PAN 2 through the shared +`/v1/videos` routes. It mirrors the Python example: one GPU, the text encoder +offloaded between requests, and the pipeline's own sampling defaults. + +```bash +fastvideo serve --config examples/serving/openai_pan2.yaml --server.host 127.0.0.1 +``` + +Text-to-video uses the same request as any other served video model, with the +alias `pan2`: + +```bash +curl --fail-with-body http://127.0.0.1:8000/v1/videos \ + -H 'Content-Type: application/json' \ + -d '{"model": "pan2", "prompt": "A red fox trots through fresh snow in a pine forest at dawn."}' +``` + +For image-to-video, add a first frame in `image_reference` (or the legacy +`input_reference`). The server then runs image-to-video; no second deployment or +workload is needed. + +```bash +curl --fail-with-body http://127.0.0.1:8000/v1/videos \ + -H 'Content-Type: application/json' \ + -d '{"model": "pan2", "prompt": "The fox turns toward the camera.", + "image_reference": {"image_url": "https://example.com/first-frame.png"}}' +``` + +Local image paths must sit under `server.media_roots`, and private or loopback +URLs need `server.allow_private_media_urls`. Poll `/v1/videos/{id}` and download +`/v1/videos/{id}/content` as in the [serving guide](openai-api.md). Generation +time for the default 13-second clip is not measured, so raise the client's +polling deadline for long jobs instead of treating a slow job as failed. + +The cookbook's generated serving clients send a text prompt only. The image +request above is the supported way to supply a first frame over HTTP, and no +recorded generation backs it yet. + +## Evidence + +Both recipes are **Source-backed**: the model ID, example and serving YAML are +checked against each other, and the serving defaults are tested against the +registered `pan2_t2v` preset. No GPU generation, memory measurement or HTTP +client run is recorded, so hardware fit and throughput are unknown. Numerical +parity with the reference implementation is checked on CPU with a public +random-weight checkpoint, which verifies the port and produces noise rather than +useful video. diff --git a/docs/cookbook/recipes/pan2.yaml b/docs/cookbook/recipes/pan2.yaml new file mode 100644 index 000000000..bfd1a0540 --- /dev/null +++ b/docs/cookbook/recipes/pan2.yaml @@ -0,0 +1,38 @@ +title: PAN 2 +summary: Text-to-video and first-frame image-to-video from the 29B PAN 2 mixture-of-transformers model. + +deployments: + cuda-rest: + runtime: fastvideo-cuda-rest + workload: t2v + defaults: examples/serving/openai_pan2.yaml + requirements: + - >- + Requests without a reference run text-to-video. Add image_reference or input_reference + to condition on a first frame; the same server and pipeline handle both. + - >- + To use more than one GPU, set sp_size equal to num_gpus (Ulysses sequence parallelism). + It must divide PAN 2's 32 attention heads, and each rank still holds the full DiT weights. + options: + - server.* + - generator.engine.num_gpus + - generator.engine.parallelism.* + - "!generator.engine.parallelism.hsdp_*" + - generator.engine.offload.text_encoder + - generator.engine.offload.vae + - default_request.sampling.num_frames + - default_request.sampling.height + - default_request.sampling.width + - default_request.sampling.fps + overrides: + default_request.sampling.num_frames: + default: 313 + description: num_frames - 1 must be a multiple of 4. + default_request.sampling.height: + default: 704 + minimum: 32 + multipleOf: 32 + default_request.sampling.width: + default: 1248 + minimum: 32 + multipleOf: 32 diff --git a/docs/cookbook/serving-coverage.md b/docs/cookbook/serving-coverage.md index 23280315f..197b90e10 100644 --- a/docs/cookbook/serving-coverage.md +++ b/docs/cookbook/serving-coverage.md @@ -49,6 +49,7 @@ never deleted. See `examples/serving/README.md`. | `openai_hunyuan15_i2v_480p.yaml` | HunyuanVideo 1.5 I2V 480p | `i2v` | CUDA REST | `hunyuan15-i2v-480p` | Source-backed; preset renders 720x1280 | | `openai_kandinsky6_ti2va.yaml` | Kandinsky 6 Pro TI2VA | `t2v` | CUDA REST | `kandinsky6-ti2va` | Source-backed | | `openai_kandinsky6_ti2va_distilled.yaml` | Kandinsky 6 Pro-distill TI2VA | `t2v` | CUDA REST | `kandinsky6-ti2va-distilled` | Source-backed | +| `openai_pan2.yaml` | PAN 2 | `t2v` (a reference switches to image-to-video) | CUDA REST | `pan2/cuda-rest` | Source-backed; checkpoint not public, no recorded run | | `openai_ltx2_distilled.yaml` | LTX2 Distilled | `t2v` | CUDA REST | `ltx2-distilled/cuda-rest` | Source-backed | | `streaming_demo.yaml` | LTX2 Distilled | `t2v` | CUDA streaming | `ltx2-distilled/cuda-streaming` | Source-backed | | `mlx_fasth3.yaml` | FastH3 V1 converted | `t2v` | MLX REST | none (legacy catalog only) | Source runtime; T2VA only | diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 8348e5af9..a614eb694 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -295,6 +295,18 @@ surfaces: - fastvideo.configs.pipelines.hunyuan15.Hunyuan15T2V720PConfig - fastvideo.configs.pipelines.hyworld.HYWorldConfig - fastvideo.configs.pipelines.hyworld.Hunyuan15T2V480PConfig + text_system_prompt: + sources: + - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig + text_encoder_skip_layers: + sources: + - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig + text_max_sequence_length: + sources: + - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig + i2v_skip_first_latent_cfg: + sources: + - fastvideo.configs.pipelines.pan2.PAN2PipelineConfig text_encoder_max_lengths: sources: - fastvideo.configs.pipelines.hunyuan15.Hunyuan15I2V480PStepDistilledConfig diff --git a/docs/generate_examples.py b/docs/generate_examples.py index 8203bead7..f7ee16b5e 100644 --- a/docs/generate_examples.py +++ b/docs/generate_examples.py @@ -35,6 +35,7 @@ COOKBOOK_FAMILIES = { "cosmos", "kandinsky5", "kandinsky6", + "pan2", "flux", "qwen_image21", "sd35", diff --git a/docs/inference/pan2.md b/docs/inference/pan2.md new file mode 100644 index 000000000..570594a6c --- /dev/null +++ b/docs/inference/pan2.md @@ -0,0 +1,86 @@ +# PAN 2 Text/Image to Video + +PAN 2 is a 29B mixture-of-transformers video DiT with separate video and text streams, a Qwen3.5 text encoder and the +Wan2.2 VAE. One pipeline serves text-to-video and image-to-video: pass `image_path` to condition on a first frame and +leave it out for text only. + +## Model + +| Hub repo | Pipeline | Workloads | Example | +|---|---|---|---| +| `IFM/PAN2` | `PAN2ModularPipeline` | T2V, I2V | [`basic_pan2.py`](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pan2.py) | + +PAN2 checkpoints use the Modular Diffusers layout (`modular_model_index.json`). Any repo id or local directory whose +name contains `pan2` (or whose manifest declares `PAN2ModularPipeline`) resolves to this pipeline. The defaults come +from the `pan2_t2v` and `pan2_i2v` presets: + +| Setting | Default | +|---|---| +| Resolution | 704x1248 (height and width must be multiples of 32) | +| Frames | 313, 13 s at 24 fps (`num_frames - 1` must be a multiple of 4) | +| Steps | 50, uniform sigmas from 1 to 0 with the checkpoint's flow shift (9 for `IFM/PAN2-TEST`, the reference sampler's value) | +| Guidance | 3.0, classifier-free guidance. Image-to-video keeps the first latent frame out of guidance (`i2v_skip_first_latent_cfg`, on by default, as in the reference sampler) | +| Negative prompt | None unless the request carries one. The reference sampler always sends the release's structured negative caption (`negative_prompt.json` of the PAN2 repository) with its three `time_range_sec` values set to the clip length, `0 to 20` for 481 frames at 24 fps | + +PAN 2 is trained on the structured captions its prompt rewriter produces, serialized as one-line JSON. Short raw +prompts work, with weaker results. + +## Usage + +```bash +python examples/inference/basic/basic_pan2.py +``` + +```python +from fastvideo import VideoGenerator + +generator = VideoGenerator.from_pretrained( + "IFM/PAN2", + num_gpus=1, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=True, +) +generator.generate_video( + "A red fox trots through fresh snow in a pine forest at dawn.", + image_path=None, # or the path of a first frame for image-to-video + output_path="video_samples_pan2", +) +``` + +For image-to-video, the first frame is scaled to cover the requested resolution and center cropped, then encoded by the +VAE into the first latent frame of the conditioning input. + +## Parallelism + +With `num_gpus > 1`, FastVideo runs Ulysses sequence parallelism: the video tokens are sharded across ranks and the +text tokens are replicated. The number of GPUs must divide the 32 attention heads. Sequence-parallel padding of the +video tokens is trimmed before attention, so no attention mask is needed. + +Padding is always excluded, without an opt-in setting. At 704x1248 with 313 frames, the 67,782 video tokens need +two padding tokens at SP=4. They are removed before joining the text tokens in every attention block. This provides +the same exclusion as vLLM-Omni's unconditional padding mask in +[revision `e0c4c731`](https://github.com/vllm-project/vllm-omni/commit/e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74). +The trimmed sequence can use either FlashAttention or Torch SDPA. + +```bash +PAN2_NUM_GPUS=4 python examples/inference/basic/basic_pan2.py +``` + +## Memory and offloading + +The example keeps the DiT on the GPUs and offloads the text encoder between requests. Sequence parallelism shards +video tokens; each rank still holds the DiT weights. Increasing `num_gpus` alone does not shard those weights. +For devices that cannot hold the DiT, see [CPU offloading](offloading.md) for `dit_cpu_offload` and +`dit_layerwise_offload`. Hardware recipes and timings reported by vLLM-Omni use its own offload implementation; +they have not been validated for FastVideo. + +The full `IFM/PAN2` checkpoint is not public yet. Local parity tests use the public random-weight +`wuqing157/tiny-pan2-modular-pipe` checkpoint; it verifies the integration and produces noise rather than useful video. + +## Differences from the vLLM-Omni integration + +- The optional prompt and video guardrails (`pan2-guardrail`) are not included. +- The Cache-DiT `quality` policy and CFG-parallel execution are vLLM-Omni serving features and are not ported. +- The text encoder is a FastVideo-native Qwen3.5 port. Its Gated DeltaNet layers use a PyTorch chunked delta rule + instead of the `fla` kernels. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index e6b40458f..a6d68eee2 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -70,6 +70,7 @@ column links a runnable script in `examples/inference/basic/` where one exists. | minimax_h3 | `MiniMaxAI/MiniMax-H3` | T2V, I2V | [T2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_t2v.py)
[FL2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_fl2va.py)
[Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_minimax_h3_ref2va.py) | | pika_h3_flex | `Pika-Labs/Pika-H3-Flex` | T2V, I2V | [T2VA/I2VA/L2VA/FL2VA/Ref2VA](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pika_h3_flex.py) | | open_h3_regenerate_2k | `Efficient-Large-Model/Open-H3-Regenerate-2K` | — | [basic_open_h3_regenerate_2k.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_open_h3_regenerate_2k.py) | +| pan2 | `IFM/PAN2` | T2V, I2V | [basic_pan2.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_pan2.py) | | sd35 | `stabilityai/stable-diffusion-3.5-medium` | T2I | [basic_sd35_t2i.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_sd35_t2i.py) | | qwen_image21 | `Qwen/Qwen-Image-2.1` | T2I, I2I | [basic_qwen_image21.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_qwen_image21.py) | | stable_audio | `FastVideo/stable-audio-open-1.0-Diffusers` | T2V | [basic_stable_audio.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_stable_audio.py) | diff --git a/examples/inference/basic/basic_pan2.py b/examples/inference/basic/basic_pan2.py new file mode 100644 index 000000000..10600ca96 --- /dev/null +++ b/examples/inference/basic/basic_pan2.py @@ -0,0 +1,40 @@ +import os + +from fastvideo import VideoGenerator + +# PAN2 is trained on the structured captions its prompt rewriter produces (serialized as one-line JSON). Plain +# prompts like this one work, with weaker results. +PROMPT = ("A red fox trots through fresh snow in a pine forest at dawn, its breath visible in the cold air. " + "The camera tracks it from the side at ground level as low golden light filters through the trees.") + +OUTPUT_PATH = "video_samples_pan2" +MODEL_PATH = os.environ.get("PAN2_MODEL_PATH", "IFM/PAN2") + +# Set to a first frame (e.g. "assets/images/girl.png") for image-to-video; the same pipeline runs both modes. +IMAGE_PATH = None + + +def main(): + # The 29B DiT runs in bf16; with more GPUs, num_gpus enables Ulysses sequence parallelism (the 32 attention + # heads must be divisible by it). + generator = VideoGenerator.from_pretrained( + MODEL_PATH, + num_gpus=int(os.environ.get("PAN2_NUM_GPUS", "1")), + use_fsdp_inference=False, + dit_cpu_offload=False, + dit_layerwise_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + ) + + generator.generate_video( + PROMPT, + image_path=IMAGE_PATH, + output_path=OUTPUT_PATH, + save_video=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/serving/README.md b/examples/serving/README.md index b95910345..7cbd01ec3 100644 --- a/examples/serving/README.md +++ b/examples/serving/README.md @@ -90,6 +90,7 @@ uncovered registry families. | `openai_hunyuan15_i2v_480p.yaml` | HunyuanVideo 1.5 I2V 480p | image-to-video | 1 | | `openai_kandinsky6_ti2va.yaml` | Kandinsky 6 Pro TI2VA | text/image-to-video with audio | 1 | | `openai_kandinsky6_ti2va_distilled.yaml` | Kandinsky 6 Pro-distill TI2VA | text/image-to-video with audio | 1 | +| `openai_pan2.yaml` | PAN 2 | text/image-to-video | 1 | | `openai_ltx2_distilled.yaml` | LTX2 Distilled | text-to-video with audio | 4 | | `openai_flux1_dev.yaml` | FLUX.1-dev | text-to-image | 1 | | `openai_flux2_dev.yaml` | FLUX.2-dev | text-to-image | 1 | diff --git a/examples/serving/openai_pan2.yaml b/examples/serving/openai_pan2.yaml new file mode 100644 index 000000000..17bc68793 --- /dev/null +++ b/examples/serving/openai_pan2.yaml @@ -0,0 +1,36 @@ +# OpenAI-compatible server for PAN 2 text-to-video and first-frame image-to-video. +# fastvideo serve --config examples/serving/openai_pan2.yaml +# Omit image_reference for text-to-video; send it (or input_reference) for image-to-video. +generator: + model_path: IFM/PAN2 + engine: + num_gpus: 1 + use_fsdp_inference: false + parallelism: + tp_size: 1 + sp_size: 1 + offload: + dit: false + dit_layerwise: false + vae: false + text_encoder: true + pin_cpu_memory: true + pipeline: + workload_type: t2v + +server: + host: 0.0.0.0 + port: 8000 + output_dir: outputs/openai_pan2 + served_model_name: pan2 + +default_request: + sampling: + height: 704 + width: 1248 + num_frames: 313 + fps: 24 + num_inference_steps: 50 + guidance_scale: 3.0 + output: + return_frames: false diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index ab4a4e0f7..4490c2cc5 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -8,6 +8,7 @@ from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig from fastvideo.configs.models.dits.magi_human import MagiHumanVideoConfig from fastvideo.configs.models.dits.minimax_h3 import MiniMaxH3Config +from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig from fastvideo.configs.models.dits.mmaudio import ( MMAUDIO_44K_TRAINING_VARIANTS, MMAUDIO_TRAINING_VARIANTS, @@ -32,7 +33,7 @@ __all__ = [ "FluxDiTConfig", "Flux2Config", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig", "Kandinsky6ArchConfig", "Kandinsky6VideoAudioConfig", "Kandinsky6SRArchConfig", "Kandinsky6SRConfig", "MagiHumanVideoConfig", "StableAudioConfig", "HeliosConfig", "LingBotWorld2CausalFastVideoConfig", "LingBotWorldFastVideoConfig", - "LingBotVideoConfig", "MiniMaxH3Config", "MMAudioArchConfig", "MMAudioTransformerConfig", + "LingBotVideoConfig", "MiniMaxH3Config", "PAN2VideoConfig", "MMAudioArchConfig", "MMAudioTransformerConfig", "MMAUDIO_VARIANT_ARCHITECTURES", "MMAUDIO_TRAINING_VARIANTS", "MMAUDIO_44K_TRAINING_VARIANTS", "get_mmaudio_transformer_config", "QwenImage21ArchConfig", "QwenImage21Config" ] diff --git a/fastvideo/configs/models/dits/pan2.py b/fastvideo/configs/models/dits/pan2.py new file mode 100644 index 000000000..8ec5b6d03 --- /dev/null +++ b/fastvideo/configs/models/dits/pan2.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig +from fastvideo.platforms import AttentionBackendEnum + + +def is_transformer_block(n: str, m) -> bool: + return "transformer_blocks" in n and str.isdigit(n.split(".")[-1]) + + +def is_refiner_block(n: str, m) -> bool: + return "context_refiner" in n and str.isdigit(n.split(".")[-1]) + + +@dataclass +class PAN2ArchConfig(DiTArchConfig): + _fsdp_shard_conditions: list = field(default_factory=lambda: [is_transformer_block, is_refiner_block]) + _supported_attention_backends: tuple[AttentionBackendEnum, ...] = ( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ) + + # Diffusers PAN2Transformer3DModel keys -> native keys. The per-stream Q/K/V projections are fused on load. + param_names_mapping: dict = field( + default_factory=lambda: { + r"^x_embedder\.(.*)$": r"x_embedder.proj.\1", + r"^time_embedder\.linear_1\.(.*)$": r"time_embedder.mlp.fc_in.\1", + r"^time_embedder\.linear_2\.(.*)$": r"time_embedder.mlp.fc_out.\1", + r"^context_embedder\.0\.(.*)$": r"context_norm.\1", + r"^context_embedder\.1\.(.*)$": r"context_proj.\1", + r"^context_refiner\.(\d+)\.attn\.to_q\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 0, 3), + r"^context_refiner\.(\d+)\.attn\.to_k\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 1, 3), + r"^context_refiner\.(\d+)\.attn\.to_v\.(.*)$": (r"context_refiner.\1.attn.to_qkv.\2", 2, 3), + r"^context_refiner\.(\d+)\.attn\.to_out\.0\.(.*)$": r"context_refiner.\1.attn.to_out.\2", + r"^context_refiner\.(\d+)\.ff\.net\.0\.proj\.(.*)$": r"context_refiner.\1.ff.fc_in.\2", + r"^context_refiner\.(\d+)\.ff\.net\.2\.(.*)$": r"context_refiner.\1.ff.fc_out.\2", + r"^transformer_blocks\.(\d+)\.attn\.to_q\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 0, 3), + r"^transformer_blocks\.(\d+)\.attn\.to_k\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 1, 3), + r"^transformer_blocks\.(\d+)\.attn\.to_v\.(.*)$": (r"transformer_blocks.\1.attn.to_qkv.\2", 2, 3), + r"^transformer_blocks\.(\d+)\.attn\.add_q_proj\.(.*)$": + (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 0, 3), + r"^transformer_blocks\.(\d+)\.attn\.add_k_proj\.(.*)$": + (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 1, 3), + r"^transformer_blocks\.(\d+)\.attn\.add_v_proj\.(.*)$": + (r"transformer_blocks.\1.attn.add_qkv_proj.\2", 2, 3), + r"^transformer_blocks\.(\d+)\.attn\.to_out\.0\.(.*)$": r"transformer_blocks.\1.attn.to_out.\2", + r"^transformer_blocks\.(\d+)\.(ff|ff_context)\.net\.0\.proj\.(.*)$": r"transformer_blocks.\1.\2.fc_in.\3", + r"^transformer_blocks\.(\d+)\.(ff|ff_context)\.net\.2\.(.*)$": r"transformer_blocks.\1.\2.fc_out.\3", + }) + + reverse_param_names_mapping: dict = field(default_factory=lambda: {}) + + patch_size: tuple[int, ...] = (1, 2, 2) + in_channels: int = 97 + out_channels: int = 48 + num_attention_heads: int = 32 + attention_head_dim: int = 128 + num_layers: int = 44 + num_refiner_layers: int = 4 + mlp_ratio: float = 8.0 + context_mlp_ratio: float = 4.0 + refiner_mlp_ratio: float = 4.0 + text_embed_dim: int = 4096 + rope_axes_dim: tuple[int, ...] = (16, 56, 56) + rope_theta: float = 256.0 + # Frame positions are rescaled to the frame rate the rotary embedding was trained at (16 fps of 24 fps). + rope_temporal_scale: float = 16 / 24 + # PAN2 runs its text stream in the DiT dtype. + cast_prompt_embeds_to_dit_dtype: bool = True + exclude_lora_layers: list[str] = field( + default_factory=lambda: ["x_embedder", "time_embedder", "context_norm", "context_proj", "context_refiner"]) + + def __post_init__(self): + super().__post_init__() + self.patch_size = tuple(self.patch_size) + self.rope_axes_dim = tuple(self.rope_axes_dim) + self.hidden_size: int = self.attention_head_dim * self.num_attention_heads + self.num_channels_latents: int = self.out_channels + + +@dataclass +class PAN2VideoConfig(DiTConfig): + arch_config: DiTArchConfig = field(default_factory=PAN2ArchConfig) + + prefix: str = "PAN2" diff --git a/fastvideo/configs/models/encoders/__init__.py b/fastvideo/configs/models/encoders/__init__.py index d2797f2db..67f442a91 100644 --- a/fastvideo/configs/models/encoders/__init__.py +++ b/fastvideo/configs/models/encoders/__init__.py @@ -16,6 +16,7 @@ from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig from fastvideo.configs.models.encoders.minimax_h3_qwen3_vl import (MiniMaxH3Qwen3VLArchConfig, MiniMaxH3Qwen3VLConfig) from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig from fastvideo.configs.models.encoders.qwen_image21 import QwenImage21Qwen3VLArchConfig, QwenImage21Qwen3VLConfig +from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextArchConfig, Qwen3_5TextConfig from fastvideo.configs.models.encoders.lingbot_video import LingBotVideoQwen3VLTextConfig from fastvideo.configs.models.encoders.stable_audio_conditioner import ( StableAudioConditionerArchConfig, @@ -52,6 +53,8 @@ __all__ = [ "Qwen3TextConfig", "QwenImage21Qwen3VLArchConfig", "QwenImage21Qwen3VLConfig", + "Qwen3_5TextArchConfig", + "Qwen3_5TextConfig", "Mistral3TextConfig", "LingBotWorld2UMT5ArchConfig", "LingBotWorld2UMT5Config", diff --git a/fastvideo/configs/models/encoders/qwen3_5.py b/fastvideo/configs/models/encoders/qwen3_5.py new file mode 100644 index 000000000..f469657ab --- /dev/null +++ b/fastvideo/configs/models/encoders/qwen3_5.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3.5 text encoder configuration (hybrid Gated DeltaNet + gated attention).""" +from dataclasses import dataclass, field +from typing import Any + +from fastvideo.configs.models.encoders.base import ( + TextEncoderArchConfig, + TextEncoderConfig, +) + +_LEGACY_LAYER_TYPES = { + "mamba": "linear_attention", + "conv": "linear_attention", + "attention": "full_attention", +} +_DEFAULT_ROPE_THETA = 10000.0 +_DEFAULT_PARTIAL_ROTARY_FACTOR = 0.25 + + +def _is_transformer_layer(n: str, m: Any) -> bool: + return "layers" in n and str.isdigit(n.split(".")[-1]) + + +def _is_embeddings(n: str, m: Any) -> bool: + return n.endswith("embed_tokens") + + +def _is_final_norm(n: str, m: Any) -> bool: + return n == "norm" + + +@dataclass +class Qwen3_5TextArchConfig(TextEncoderArchConfig): + """Architecture config for the Qwen3.5 text model (``Qwen3_5TextModel``). + + Layers alternate between Gated DeltaNet linear attention and gated full + attention according to ``layer_types``. The defaults describe the 9B text + tower; checkpoint ``config.json`` values override them at load time. + """ + + vocab_size: int = 248320 + hidden_size: int = 4096 + intermediate_size: int = 12288 + num_hidden_layers: int = 32 + num_attention_heads: int = 16 + num_key_value_heads: int = 4 + head_dim: int = 256 + hidden_act: str = "silu" + max_position_embeddings: int = 32768 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-6 + use_cache: bool = True + tie_word_embeddings: bool = False + attention_bias: bool = False + attention_dropout: float = 0.0 + linear_conv_kernel_dim: int = 4 + linear_key_head_dim: int = 128 + linear_value_head_dim: int = 128 + linear_num_key_heads: int = 16 + linear_num_value_heads: int = 32 + layer_types: list[str] | None = None + full_attention_interval: int = 4 + # RoPE settings as written by Transformers 5 (``rope_parameters``) or by + # older configs (top-level ``rope_theta`` / ``partial_rotary_factor`` / + # ``rope_scaling``). Values inside ``rope_parameters`` take precedence. + rope_parameters: dict[str, Any] | None = None + rope_scaling: dict[str, Any] | None = None + rope_theta: float | None = None + partial_rotary_factor: float | None = None + pad_token_id: int | None = None + bos_token_id: int | None = None + eos_token_id: int | list[int] | None = None + + _fsdp_shard_conditions: list = field( + default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm]) + + def __post_init__(self) -> None: + super().__post_init__() + # Validate eagerly; the resolved views are recomputed on demand because + # ``update_model_arch`` re-runs this hook after replacing raw fields. + self.resolved_layer_types() + self.resolved_rope_parameters() + + def resolved_layer_types(self) -> list[str]: + """Return one ``linear_attention`` / ``full_attention`` entry per layer.""" + if self.layer_types is None: + interval = self.full_attention_interval + layer_types = [ + "linear_attention" if (index + 1) % interval else "full_attention" + for index in range(self.num_hidden_layers) + ] + else: + layer_types = [_LEGACY_LAYER_TYPES.get(kind, kind) for kind in self.layer_types] + if len(layer_types) < self.num_hidden_layers: + raise ValueError(f"Qwen3.5 layer_types has {len(layer_types)} entries for " + f"{self.num_hidden_layers} layers") + unknown = set(layer_types) - {"linear_attention", "full_attention"} + if unknown: + raise ValueError(f"Unsupported Qwen3.5 layer types: {sorted(unknown)}") + return layer_types[:self.num_hidden_layers] + + def resolved_rope_parameters(self) -> dict[str, Any]: + """Return the effective text RoPE settings with Transformers' precedence.""" + rope = dict(self.rope_scaling or self.rope_parameters or {}) + rope.setdefault("rope_type", rope.get("type", "default")) + rope.setdefault("rope_theta", self.rope_theta if self.rope_theta is not None else _DEFAULT_ROPE_THETA) + rope.setdefault( + "partial_rotary_factor", + self.partial_rotary_factor if self.partial_rotary_factor is not None else _DEFAULT_PARTIAL_ROTARY_FACTOR, + ) + return rope + + +@dataclass +class Qwen3_5TextConfig(TextEncoderConfig): + """Top-level config for the Qwen3.5 text encoder.""" + + arch_config: TextEncoderArchConfig = field(default_factory=Qwen3_5TextArchConfig) + prefix: str = "qwen3_5" diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index f42e517b8..a925d5123 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -15,6 +15,7 @@ from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.mmaudio import MMAudioV2AConfig from fastvideo.configs.pipelines.qwen_image21 import QwenImage21PipelineConfig +from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.registry import get_pipeline_config_cls_from_name from fastvideo.models.wan.pipeline_config import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig, @@ -27,5 +28,5 @@ __all__ = [ "Kandinsky5T2VConfig", "Kandinsky5I2VConfig", "Kandinsky5DMDConfig", "Kandinsky6TI2VAConfig", "Kandinsky6SRPipelineConfig", "LingBotWorld2CausalFastI2V480PConfig", "LingBotWorldFastI2V480PConfig", "LingBotVideoT2VConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "MMAudioV2AConfig", - "QwenImage21PipelineConfig", "get_pipeline_config_cls_from_name" + "QwenImage21PipelineConfig", "PAN2PipelineConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/pan2.py b/fastvideo/configs/pipelines/pan2.py new file mode 100644 index 000000000..8e58c33b9 --- /dev/null +++ b/fastvideo/configs/pipelines/pan2.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 pipeline configuration: text-to-video, and image-to-video from a first frame.""" + +from dataclasses import dataclass, field + +from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig +from fastvideo.configs.models.dits import PAN2VideoConfig +from fastvideo.configs.models.encoders import Qwen3_5TextConfig +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.models.wan.vae_config import WanVAEConfig + +PAN2_SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions." + + +@dataclass +class PAN2PipelineConfig(PipelineConfig): + """One configuration serves text-to-video and, given a first frame, image-to-video.""" + + dit_config: DiTConfig = field(default_factory=PAN2VideoConfig) + # The Wan2.2 48-channel VAE; its architecture and latent statistics come from vae/config.json at load time. + vae_config: VAEConfig = field(default_factory=WanVAEConfig) + text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Qwen3_5TextConfig(), )) + text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", )) + + dit_precision: str = "bf16" + vae_precision: str = "fp32" + # Unset: the checkpoint's scheduler config decides, which is shift 9 for IFM/PAN2-TEST, the reference sampler's + # fixed value. A fixed default here would override the released shift with whatever a test checkpoint declares. + flow_shift: float | None = None + # Image-to-video: the first latent frame (the conditioning image) takes the positive prediction instead of the + # guided one, as the reference sampler does by default (``--i2v_skip_first_latent_cfg``). + i2v_skip_first_latent_cfg: bool = True + # PAN2 samples on fp32 latents; keep the Euler update out of autocast. + scheduler_step_in_fp32: bool = True + + # Prompt encoding: PAN2 conditions on the tokens after the chat template's user header, taken from the text + # encoder's hidden states with its last two layers skipped (``hidden_states[-3]``). + text_system_prompt: str = PAN2_SYSTEM_PROMPT + text_encoder_skip_layers: int = 2 + text_max_sequence_length: int = 1536 + + def __post_init__(self) -> None: + # The encoder is needed for the image-to-video first frame. + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + self.text_encoder_configs[0].arch_config.output_hidden_states = True diff --git a/fastvideo/models/dits/pan2.py b/fastvideo/models/dits/pan2.py new file mode 100644 index 000000000..40c809d88 --- /dev/null +++ b/fastvideo/models/dits/pan2.py @@ -0,0 +1,398 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 video transformer. + +A mixture-of-transformers DiT: the video and text streams keep separate weights and meet only in one joint +attention per block, over ``[video; text]``. The text stream is modulated by a learned, timestep-free table. +Under sequence parallelism the video tokens are sharded and the text tokens are replicated on every rank. +""" + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from fastvideo.attention import DistributedAttention, LocalAttention +from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig +from fastvideo.distributed.communication_op import (sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) +from fastvideo.distributed.parallel_state import get_sp_world_size +from fastvideo.layers.layernorm import RMSNorm +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import get_1d_rotary_pos_embed +from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum + + +class PAN2SwiGLU(nn.Module): + """SwiGLU MLP of the text refiner; the fused input projection stores ``[value; gate]``.""" + + def __init__(self, dim: int, inner_dim: int, prefix: str = ""): + super().__init__() + self.fc_in = ReplicatedLinear(dim, 2 * inner_dim, bias=False, prefix=f"{prefix}.fc_in") + self.fc_out = ReplicatedLinear(inner_dim, dim, bias=False, prefix=f"{prefix}.fc_out") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, _ = self.fc_in(x) + value, gate = x.chunk(2, dim=-1) + x, _ = self.fc_out(value * F.silu(gate)) + return x + + +class PAN2RefinerAttention(nn.Module): + """Text-only self-attention of the refiner; the text is replicated on every sequence-parallel rank.""" + + def __init__(self, + dim: int, + num_heads: int, + head_dim: int, + eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + self.num_heads = num_heads + self.to_qkv = ReplicatedLinear(dim, 3 * num_heads * head_dim, bias=False, prefix=f"{prefix}.to_qkv") + self.norm_q = RMSNorm(head_dim, eps=eps) + self.norm_k = RMSNorm(head_dim, eps=eps) + self.to_out = ReplicatedLinear(num_heads * head_dim, dim, bias=False, prefix=f"{prefix}.to_out") + self.attn = LocalAttention(num_heads=num_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + qkv, _ = self.to_qkv(hidden_states) + query, key, value = qkv.unflatten(-1, (3, self.num_heads, -1)).unbind(2) + query = self.norm_q(query).to(value.dtype) + key = self.norm_k(key).to(value.dtype) + hidden_states = self.attn(query, key, value) + hidden_states, _ = self.to_out(hidden_states.flatten(2)) + return hidden_states + + +class PAN2TokenRefinerBlock(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + head_dim: int, + mlp_ratio: float, + eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + self.norm1 = RMSNorm(dim, eps=eps) + self.attn = PAN2RefinerAttention(dim, + num_heads, + head_dim, + eps, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") + self.norm2 = RMSNorm(dim, eps=eps) + self.ff = PAN2SwiGLU(dim, int(dim * mlp_ratio), prefix=f"{prefix}.ff") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states)) + hidden_states = hidden_states + self.ff(self.norm2(hidden_states)) + return hidden_states + + +class PAN2AdaLayerNormZero(nn.Module): + """Timestep modulation of the video stream: modulated attention input plus the gate/shift/scale for the rest.""" + + def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""): + super().__init__() + self.linear = ReplicatedLinear(dim, 6 * dim, bias=True, prefix=f"{prefix}.linear") + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + + def forward(self, x: torch.Tensor, + temb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + emb, _ = self.linear(F.silu(temb)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) + x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None] + return x, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +class PAN2AdaLayerNormOut(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-6, prefix: str = ""): + super().__init__() + self.linear = ReplicatedLinear(dim, 2 * dim, bias=True, prefix=f"{prefix}.linear") + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + + def forward(self, x: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + emb, _ = self.linear(F.silu(temb).to(x.dtype)) + scale, shift = emb.chunk(2, dim=-1) + return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + + +class PAN2Attention(nn.Module): + """Joint attention over ``[video; text]`` with per-stream projections and Q/K RMSNorm. + + RoPE applies to the video tokens only; ``DistributedAttention`` applies it after its all-to-all, where each rank + holds the full (unpadded) video sequence for its heads. + """ + + def __init__(self, + dim: int, + num_heads: int, + head_dim: int, + eps: float, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + self.num_heads = num_heads + inner_dim = num_heads * head_dim + + self.to_qkv = ReplicatedLinear(dim, 3 * inner_dim, bias=True, prefix=f"{prefix}.to_qkv") + self.norm_q = RMSNorm(head_dim, eps=eps) + self.norm_k = RMSNorm(head_dim, eps=eps) + self.to_out = ReplicatedLinear(inner_dim, dim, bias=True, prefix=f"{prefix}.to_out") + + self.add_qkv_proj = ReplicatedLinear(dim, 3 * inner_dim, bias=True, prefix=f"{prefix}.add_qkv_proj") + self.norm_added_q = RMSNorm(head_dim, eps=eps) + self.norm_added_k = RMSNorm(head_dim, eps=eps) + self.to_add_out = ReplicatedLinear(inner_dim, dim, bias=True, prefix=f"{prefix}.to_add_out") + + self.attn = DistributedAttention(num_heads=num_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") + + def _project(self, projection: ReplicatedLinear, + hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + qkv, _ = projection(hidden_states) + query, key, value = qkv.unflatten(-1, (3, self.num_heads, -1)).unbind(2) + return query, key, value + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + original_seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + query, key, value = self._project(self.to_qkv, hidden_states) + query = self.norm_q(query).to(value.dtype) + key = self.norm_k(key).to(value.dtype) + + encoder_query, encoder_key, encoder_value = self._project(self.add_qkv_proj, encoder_hidden_states) + encoder_query = self.norm_added_q(encoder_query).to(encoder_value.dtype) + encoder_key = self.norm_added_k(encoder_key).to(encoder_value.dtype) + + hidden_states, encoder_hidden_states = self.attn(query, + key, + value, + original_seq_len, + replicated_q=encoder_query, + replicated_k=encoder_key, + replicated_v=encoder_value, + freqs_cis=freqs_cis) + hidden_states, _ = self.to_out(hidden_states.flatten(2)) + encoder_hidden_states, _ = self.to_add_out(encoder_hidden_states.flatten(2)) + return hidden_states, encoder_hidden_states + + +class PAN2TransformerBlock(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + head_dim: int, + mlp_ratio: float, + context_mlp_ratio: float, + eps: float = 1e-6, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = ""): + super().__init__() + self.norm1 = PAN2AdaLayerNormZero(dim, eps=eps, prefix=f"{prefix}.norm1") + self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim)) + + self.attn = PAN2Attention(dim, + num_heads, + head_dim, + eps, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") + + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff = MLP(dim, int(dim * mlp_ratio), act_type="gelu_pytorch_tanh", prefix=f"{prefix}.ff") + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff_context = MLP(dim, + int(dim * context_mlp_ratio), + act_type="gelu_pytorch_tanh", + prefix=f"{prefix}.ff_context") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + original_seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(hidden_states, temb) + ( + context_shift_msa, + context_scale_msa, + context_gate_msa, + context_shift_mlp, + context_scale_mlp, + context_gate_mlp, + ) = self.scale_shift_table_context.to(temb.dtype).chunk(6, dim=-1) + norm_encoder_hidden_states = (self.norm1_context(encoder_hidden_states) * + (1 + context_scale_msa) + context_shift_msa) + + attn_output, context_attn_output = self.attn(norm_hidden_states, norm_encoder_hidden_states, freqs_cis, + original_seq_len) + hidden_states = hidden_states + attn_output * gate_msa.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + context_attn_output * context_gate_msa + + norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + norm_encoder_hidden_states = (self.norm2_context(encoder_hidden_states) * + (1 + context_scale_mlp) + context_shift_mlp) + + hidden_states = hidden_states + self.ff(norm_hidden_states) * gate_mlp.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + self.ff_context(norm_encoder_hidden_states) * context_gate_mlp + return hidden_states, encoder_hidden_states + + +class PAN2Transformer3DModel(BaseDiT): + """PAN2 DiT. The input stacks the noisy latents, the conditioning latents and a one-channel conditioning mask + along channels, so one model serves text-to-video and first-frame image-to-video.""" + + _fsdp_shard_conditions = PAN2VideoConfig()._fsdp_shard_conditions + _compile_conditions = PAN2VideoConfig()._compile_conditions + _supported_attention_backends = PAN2VideoConfig()._supported_attention_backends + param_names_mapping = PAN2VideoConfig().param_names_mapping + reverse_param_names_mapping = PAN2VideoConfig().reverse_param_names_mapping + lora_param_names_mapping = PAN2VideoConfig().lora_param_names_mapping + + def __init__(self, config: PAN2VideoConfig, hf_config: dict[str, Any], **kwargs) -> None: + super().__init__(config=config, hf_config=hf_config) + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = inner_dim + self.num_attention_heads = config.num_attention_heads + self.num_channels_latents = config.num_channels_latents + self.out_channels = config.out_channels + self.patch_size = tuple(config.patch_size) + self.rope_axes_dim = tuple(config.rope_axes_dim) + if sum(self.rope_axes_dim) != config.attention_head_dim: + raise ValueError(f"rope_axes_dim {self.rope_axes_dim} must sum to the head dim {config.attention_head_dim}") + if config.num_attention_heads % get_sp_world_size() != 0: + raise ValueError(f"The number of attention heads ({config.num_attention_heads}) must be divisible by the " + f"sequence parallel size ({get_sp_world_size()})") + + self.x_embedder = PatchEmbed(self.patch_size, + config.in_channels, + inner_dim, + prefix=f"{config.prefix}.x_embedder") + self.time_embedder = TimestepEmbedder(inner_dim, act_layer="silu", frequency_embedding_size=256) + + self.context_norm = RMSNorm(config.text_embed_dim, eps=1e-5) + self.context_proj = ReplicatedLinear(config.text_embed_dim, + inner_dim, + bias=True, + prefix=f"{config.prefix}.context_proj") + # Layerwise offload hooks the first nn.ModuleList child, which must be the main blocks. + self.transformer_blocks = nn.ModuleList([ + PAN2TransformerBlock(inner_dim, + config.num_attention_heads, + config.attention_head_dim, + config.mlp_ratio, + config.context_mlp_ratio, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.transformer_blocks.{i}") for i in range(config.num_layers) + ]) + + self.context_refiner = nn.ModuleList([ + PAN2TokenRefinerBlock(inner_dim, + config.num_attention_heads, + config.attention_head_dim, + config.refiner_mlp_ratio, + eps=1e-5, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.context_refiner.{i}") + for i in range(config.num_refiner_layers) + ]) + self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim)) + + self.norm_out = PAN2AdaLayerNormOut(inner_dim, eps=1e-6, prefix=f"{config.prefix}.norm_out") + self.proj_out = ReplicatedLinear(inner_dim, + self.patch_size[0] * self.patch_size[1] * self.patch_size[2] * + self.out_channels, + bias=True, + prefix=f"{config.prefix}.proj_out") + + self.gradient_checkpointing = False + self.__post_init__() + + def _rotary_pos_embed(self, num_frames: int, height: int, width: int, + device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + """Half-dim cos/sin of the 3-axis ``(t, h, w)`` RoPE over the full (unsharded) video token grid.""" + axes = [torch.arange(size, device=device, dtype=torch.float32) for size in (num_frames, height, width)] + grid = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=0) + grid[0] = grid[0] * self.config.rope_temporal_scale + cos_sin = [ + get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), theta=self.config.rope_theta, use_real=False) + for i, dim in enumerate(self.rope_axes_dim) + ] + freqs_cos = torch.cat([cos for cos, _ in cos_sin], dim=1).float() + freqs_sin = torch.cat([sin for _, sin in cos_sin], dim=1).float() + return freqs_cos, freqs_sin + + def forward(self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + guidance=None, + **kwargs) -> torch.Tensor: + if not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + + batch_size, _, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + freqs_cis = self._rotary_pos_embed(post_patch_num_frames, post_patch_height, post_patch_width, + hidden_states.device) + + hidden_states = self.x_embedder(hidden_states) + hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) + + if timestep.dim() == 0: + timestep = timestep.expand(batch_size) + temb = self.time_embedder(timestep).to(hidden_states.dtype) + + encoder_hidden_states, _ = self.context_proj(self.context_norm(encoder_hidden_states)) + for refiner_block in self.context_refiner: + encoder_hidden_states = refiner_block(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states + self.context_type_embedding + + for block in self.transformer_blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, temb, freqs_cis, original_seq_len) + else: + hidden_states, encoder_hidden_states = block(hidden_states, encoder_hidden_states, temb, freqs_cis, + original_seq_len) + + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) + hidden_states, _ = self.proj_out(hidden_states) + + # proj_out emits channels before the patch offsets: (C, p_t, p_h, p_w). + hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height, post_patch_width, + -1, p_t, p_h, p_w) + hidden_states = hidden_states.permute(0, 4, 1, 5, 2, 6, 3, 7) + return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +EntryClass = PAN2Transformer3DModel diff --git a/fastvideo/models/encoders/qwen3_5.py b/fastvideo/models/encoders/qwen3_5.py new file mode 100644 index 000000000..c036620fd --- /dev/null +++ b/fastvideo/models/encoders/qwen3_5.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3.5 text encoder: hybrid Gated DeltaNet linear attention + gated full attention. + +The encoder runs replicated (no tensor or sequence parallelism) and +inference-only, so the linear-attention layers use a pure-PyTorch chunked gated +delta rule instead of the fused ``fla`` / ``causal_conv1d`` kernels. Full +attention calls ``scaled_dot_product_attention`` directly with head-major +``[B, H, L, D]`` inputs: that is the layout the reference text model feeds SDPA, +and keeping it makes bf16 hidden states reproducible against that reference. +""" +from collections.abc import Iterable +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from fastvideo.configs.models.encoders import BaseEncoderOutput +from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextArchConfig, Qwen3_5TextConfig +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.quantization import QuantizationConfig +from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding +from fastvideo.logger import init_logger +from fastvideo.models.encoders.base import TextEncoder +from fastvideo.models.loader.weight_utils import default_weight_loader + +logger = init_logger(__name__) + +_GATED_DELTA_CHUNK_SIZE = 64 +# Checkpoints saved from the multimodal wrapper nest the text tower. +_CHECKPOINT_PREFIXES = ("model.language_model.", "language_model.", "model.") +_SKIPPED_CHECKPOINT_PREFIXES = ("lm_head.", "mtp.", "visual.", "model.visual.") + + +class Qwen3_5RMSNorm(nn.Module): + """Zero-centered RMSNorm: the stored weight is an offset from one.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.zeros(hidden_size)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + output = x.float() + output = output * torch.rsqrt(output.pow(2).mean(-1, keepdim=True) + self.eps) + # The scale is applied in float32 before the downcast. + output = output * (1.0 + self.weight.float()) + return output.type_as(x) + + +class Qwen3_5RMSNormGated(nn.Module): + """Per-head RMSNorm followed by a SiLU output gate (plain, not zero-centered, weight).""" + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + hidden_states = self.weight * hidden_states.to(input_dtype) + hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + return hidden_states.to(input_dtype) + + +def _l2norm(x: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + return x * torch.rsqrt((x * x).sum(dim=-1, keepdim=True) + eps) + + +def chunk_gated_delta_rule( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = _GATED_DELTA_CHUNK_SIZE, +) -> torch.Tensor: + """Chunked gated delta rule over a full prefix with a zero initial state. + + Inputs are ``[batch, seq, heads, dim]`` (``g`` / ``beta`` are + ``[batch, seq, heads]``); q/k must already be L2-normalized. The recurrence + runs in float32 and the output is returned in the query dtype. + """ + output_dtype = query.dtype + query, key, value, beta, g = (x.transpose(1, 2).contiguous().to(torch.float32) + for x in (query, key, value, beta, g)) + + batch_size, num_heads, sequence_length, k_head_dim = key.shape + v_head_dim = value.shape[-1] + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_sequence_length = sequence_length + pad_size + query = query * (1 / (query.shape[-1]**0.5)) + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + query, key, value, k_beta, v_beta = (x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) + for x in (query, key, value, k_beta, v_beta)) + g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) + upper_with_diagonal = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), + diagonal=0) + + # Intra-chunk decay and the (I + strictly-lower)^-1 WY transform, solved by + # forward substitution row by row. + g = g.cumsum(dim=-1) + decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(upper_with_diagonal, 0) + for i in range(1, chunk_size): + row = attn[..., i, :i].clone() + sub = attn[..., :i, :i].clone() + attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) + value = attn @ v_beta + k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + + recurrent_state = torch.zeros(batch_size, + num_heads, + k_head_dim, + v_head_dim, + dtype=value.dtype, + device=value.device) + core_attn_out = torch.zeros_like(value) + for i in range(total_sequence_length // chunk_size): + q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] + # decay_mask is already zero above the diagonal, which keeps this causal. + attn = q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i] + v_prime = k_cumdecay[:, :, i] @ recurrent_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ recurrent_state + core_attn_out[:, :, i] = attn_inter + attn @ v_new + recurrent_state = (recurrent_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new) + + core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1]) + core_attn_out = core_attn_out[:, :, :sequence_length] + return core_attn_out.transpose(1, 2).contiguous().to(output_dtype) + + +class Qwen3_5GatedDeltaNet(nn.Module): + """Gated DeltaNet token mixer used by the ``linear_attention`` layers.""" + + def __init__( + self, + config: Qwen3_5TextArchConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + if self.num_v_heads % self.num_k_heads: + raise ValueError(f"Qwen3.5 linear value heads ({self.num_v_heads}) must be a multiple of key heads " + f"({self.num_k_heads})") + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_size = config.linear_conv_kernel_dim + if config.hidden_act != "silu": + raise ValueError(f"Unsupported Qwen3.5 activation: {config.hidden_act}") + + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + bias=False, + kernel_size=self.conv_kernel_size, + groups=self.conv_dim, + padding=self.conv_kernel_size - 1, + ) + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) + self.A_log = nn.Parameter(torch.zeros(self.num_v_heads)) + self.norm = Qwen3_5RMSNormGated(self.head_v_dim, eps=config.rms_norm_eps) + self.out_proj = ReplicatedLinear(self.value_dim, + self.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.out_proj") + self.in_proj_qkv = ReplicatedLinear(self.hidden_size, + self.conv_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_proj_qkv") + self.in_proj_z = ReplicatedLinear(self.hidden_size, + self.value_dim, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_proj_z") + self.in_proj_b = ReplicatedLinear(self.hidden_size, + self.num_v_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_proj_b") + self.in_proj_a = ReplicatedLinear(self.hidden_size, + self.num_v_heads, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.in_proj_a") + + def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor: + if padding_mask is not None: + # Zero padded tokens so they cannot leak into the recurrent state. + hidden_states = (hidden_states * padding_mask[:, :, None]).to(hidden_states.dtype) + batch_size, seq_len, _ = hidden_states.shape + + mixed_qkv, _ = self.in_proj_qkv(hidden_states) + z, _ = self.in_proj_z(hidden_states) + b, _ = self.in_proj_b(hidden_states) + a, _ = self.in_proj_a(hidden_states) + + # Causal depthwise conv: left padding of kernel-1, keep the first seq_len outputs. + mixed_qkv = mixed_qkv.transpose(1, 2) + mixed_qkv = F.silu(self.conv1d(mixed_qkv.to(self.conv1d.weight.dtype))[:, :, :seq_len]).to(mixed_qkv.dtype) + mixed_qkv = mixed_qkv.transpose(1, 2) + query, key, value = torch.split(mixed_qkv, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + query = query.reshape(batch_size, seq_len, -1, self.head_k_dim) + key = key.reshape(batch_size, seq_len, -1, self.head_k_dim) + value = value.reshape(batch_size, seq_len, -1, self.head_v_dim) + + beta = b.sigmoid() + # float32 keeps exp(A_log) finite for half-precision weights. + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + if self.num_v_heads > self.num_k_heads: + repeats = self.num_v_heads // self.num_k_heads + query = query.repeat_interleave(repeats, dim=2) + key = key.repeat_interleave(repeats, dim=2) + + core_attn_out = chunk_gated_delta_rule(_l2norm(query), _l2norm(key), value, g=g, beta=beta) + + core_attn_out = self.norm(core_attn_out.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim)) + core_attn_out = core_attn_out.reshape(batch_size, seq_len, -1) + output, _ = self.out_proj(core_attn_out) + return output + + +class Qwen3_5RotaryEmbedding(nn.Module): + """Text-only partial RoPE. + + Qwen3.5's interleaved M-RoPE collapses to standard RoPE when the temporal, + height and width positions are identical, which holds for text tokens. + """ + + def __init__(self, config: Qwen3_5TextArchConfig) -> None: + super().__init__() + rope = config.resolved_rope_parameters() + if rope["rope_type"] != "default": + raise NotImplementedError(f"Qwen3.5 rope_type {rope['rope_type']!r} is not supported") + rotary_dim = int(config.head_dim * rope["partial_rotary_factor"]) + inv_freq = 1.0 / (rope["rope_theta"]**( + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cpu") / rotary_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, hidden_states: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = self.inv_freq.to(device=hidden_states.device, dtype=torch.float32) + device_type = hidden_states.device.type if hidden_states.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + inv_freq = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1) + freqs = (inv_freq @ position_ids[:, None, :].float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(hidden_states.dtype), sin.to(hidden_states.dtype) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _repeat_kv(x: torch.Tensor, groups: int) -> torch.Tensor: + batch_size, num_kv_heads, seq_len, head_dim = x.shape + x = x[:, :, None].expand(batch_size, num_kv_heads, groups, seq_len, head_dim) + return x.reshape(batch_size, num_kv_heads * groups, seq_len, head_dim) + + +def _apply_partial_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Rotate the leading ``cos.shape[-1]`` channels of ``[B, H, L, D]`` and pass the rest through.""" + rotary_dim = cos.shape[-1] + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + x_rot, x_pass = x[..., :rotary_dim], x[..., rotary_dim:] + x_rot = x_rot * cos + _rotate_half(x_rot) * sin + return torch.cat([x_rot, x_pass], dim=-1) + + +class Qwen3_5Attention(nn.Module): + """Gated full attention with per-head QK-norm, partial RoPE and GQA.""" + + def __init__( + self, + config: Qwen3_5TextConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + arch = config.arch_config + self.num_heads = arch.num_attention_heads + self.num_kv_heads = arch.num_key_value_heads + if self.num_heads % self.num_kv_heads: + raise ValueError(f"Qwen3.5 attention heads ({self.num_heads}) must be a multiple of KV heads " + f"({self.num_kv_heads})") + self.head_dim = arch.head_dim + self.scaling = self.head_dim**-0.5 + # q_proj packs, per head, the query followed by its output gate. + self.q_proj = ReplicatedLinear(arch.hidden_size, + self.num_heads * self.head_dim * 2, + bias=arch.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.q_proj") + self.k_proj = ReplicatedLinear(arch.hidden_size, + self.num_kv_heads * self.head_dim, + bias=arch.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.k_proj") + self.v_proj = ReplicatedLinear(arch.hidden_size, + self.num_kv_heads * self.head_dim, + bias=arch.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.v_proj") + self.o_proj = ReplicatedLinear(self.num_heads * self.head_dim, + arch.hidden_size, + bias=arch.attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj") + self.q_norm = Qwen3_5RMSNorm(self.head_dim, eps=arch.rms_norm_eps) + self.k_norm = Qwen3_5RMSNorm(self.head_dim, eps=arch.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch_size, seq_len, _ = hidden_states.shape + packed, _ = self.q_proj(hidden_states) + query, gate = torch.chunk(packed.view(batch_size, seq_len, -1, self.head_dim * 2), 2, dim=-1) + gate = gate.reshape(batch_size, seq_len, -1) + key, _ = self.k_proj(hidden_states) + value, _ = self.v_proj(hidden_states) + + query = self.q_norm(query).transpose(1, 2) + key = self.k_norm(key.view(batch_size, seq_len, -1, self.head_dim)).transpose(1, 2) + value = value.view(batch_size, seq_len, -1, self.head_dim).transpose(1, 2) + cos, sin = position_embeddings + query = _apply_partial_rotary(query, cos, sin) + key = _apply_partial_rotary(key, cos, sin) + + groups = self.num_heads // self.num_kv_heads + # Native GQA is only used where SDPA keeps a fused kernel (no mask, head_dim <= 256). + native_gqa = padding_mask is None and self.head_dim <= 256 + if groups > 1 and not native_gqa: + key = _repeat_kv(key, groups) + value = _repeat_kv(value, groups) + if padding_mask is None: + attn_output = F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.0, + is_causal=seq_len > 1, + scale=self.scaling, + enable_gqa=groups > 1 and native_gqa, + ) + else: + causal = torch.ones(seq_len, seq_len, device=query.device, dtype=torch.bool).tril() + mask = causal[None, None] & padding_mask.to(device=query.device, dtype=torch.bool)[:, None, None, :] + attn_output = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=mask, + dropout_p=0.0, + is_causal=False, + scale=self.scaling, + ) + + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, -1) + attn_output = attn_output * torch.sigmoid(gate) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen3_5MLP(nn.Module): + + def __init__( + self, + config: Qwen3_5TextArchConfig, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if config.hidden_act != "silu": + raise ValueError(f"Unsupported Qwen3.5 activation: {config.hidden_act}") + self.gate_proj = ReplicatedLinear(config.hidden_size, + config.intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_proj") + self.up_proj = ReplicatedLinear(config.hidden_size, + config.intermediate_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.up_proj") + self.down_proj = ReplicatedLinear(config.intermediate_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj") + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate, _ = self.gate_proj(x) + up, _ = self.up_proj(x) + output, _ = self.down_proj(F.silu(gate) * up) + return output + + +class Qwen3_5DecoderLayer(nn.Module): + + def __init__( + self, + config: Qwen3_5TextConfig, + layer_type: str, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + arch = config.arch_config + self.layer_type = layer_type + if layer_type == "linear_attention": + self.linear_attn = Qwen3_5GatedDeltaNet(arch, quant_config=quant_config, prefix=f"{prefix}.linear_attn") + else: + self.self_attn = Qwen3_5Attention(config, quant_config=quant_config, prefix=f"{prefix}.self_attn") + self.mlp = Qwen3_5MLP(arch, quant_config=quant_config, prefix=f"{prefix}.mlp") + self.input_layernorm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps) + self.post_attention_layernorm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + if self.layer_type == "linear_attention": + hidden_states = self.linear_attn(hidden_states, padding_mask) + else: + hidden_states = self.self_attn(hidden_states, position_embeddings, padding_mask) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class Qwen3_5TextModel(TextEncoder): + """Qwen3.5 text tower used as a diffusion text encoder (e.g. PAN2). + + ``hidden_states`` follows the Transformers contract: entry 0 is the token + embeddings, entry ``i`` is the output of layer ``i - 1``, and the last entry + is the final-normed output (it replaces the raw last-layer output). + """ + + def __init__(self, config: Qwen3_5TextConfig) -> None: + super().__init__(config) + self.config = config + arch = config.arch_config + self.quant_config = getattr(config, "quant_config", None) + self.layer_types = arch.resolved_layer_types() + + self.embed_tokens = VocabParallelEmbedding( + arch.vocab_size, + arch.hidden_size, + org_num_embeddings=arch.vocab_size, + quant_config=self.quant_config, + ) + self.layers = nn.ModuleList([ + Qwen3_5DecoderLayer( + config, + layer_type, + quant_config=self.quant_config, + prefix=f"{config.prefix}.layers.{index}", + ) for index, layer_type in enumerate(self.layer_types) + ]) + self.norm = Qwen3_5RMSNorm(arch.hidden_size, eps=arch.rms_norm_eps) + self.rotary_emb = Qwen3_5RotaryEmbedding(arch) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_hidden_states: bool | None = None, + **kwargs: Any, + ) -> BaseEncoderOutput: + output_hidden_states = (output_hidden_states + if output_hidden_states is not None else self.config.arch_config.output_hidden_states) + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + assert input_ids is not None + if input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + inputs_embeds = self.get_input_embeddings(input_ids) + hidden_states = inputs_embeds + batch_size, seq_len = hidden_states.shape[:2] + + if position_ids is None: + position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0).expand(batch_size, -1) + elif position_ids.ndim == 1: + position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) + elif position_ids.ndim != 2: + raise ValueError("Qwen3.5 text encoder expects text position_ids shaped [batch, seq]") + + # An all-ones mask is equivalent to no mask; only real padding takes the masked path. + padding_mask = None + if attention_mask is not None: + if attention_mask.ndim != 2: + raise ValueError("Qwen3.5 text encoder expects a 2D [batch, seq] padding mask") + if not bool(attention_mask.to(torch.bool).all()): + padding_mask = attention_mask.to(hidden_states.device) + + position_embeddings = self.rotary_emb(hidden_states, position_ids) + all_hidden_states: tuple[torch.Tensor, ...] | None = () if output_hidden_states else None + for layer in self.layers: + if all_hidden_states is not None: + all_hidden_states += (hidden_states, ) + hidden_states = layer(hidden_states, position_embeddings, padding_mask) + + hidden_states = self.norm(hidden_states) + if all_hidden_states is not None: + all_hidden_states += (hidden_states, ) + + return BaseEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attention_mask=attention_mask, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if name.startswith(_SKIPPED_CHECKPOINT_PREFIXES) or "rotary_emb.inv_freq" in name: + continue + for checkpoint_prefix in _CHECKPOINT_PREFIXES: + if name.startswith(checkpoint_prefix): + name = name[len(checkpoint_prefix):] + break + if name not in params_dict: + logger.debug("Skipping unexpected Qwen3.5 checkpoint tensor %s", name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +EntryClass = Qwen3_5TextModel diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 897e3f276..de6464b31 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -44,6 +44,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"), "Kandinsky6Transformer3DModel": ("dits", "kandinsky6", "Kandinsky6Transformer3DModel"), "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"), + "PAN2Transformer3DModel": ("dits", "pan2", "PAN2Transformer3DModel"), } _IMAGE_TO_VIDEO_DIT_MODELS = { @@ -91,6 +92,7 @@ _TEXT_ENCODER_MODELS = { "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"), "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"), "Mistral3ForConditionalGeneration": ("encoders", "mistral3", "Mistral3ForConditionalGeneration"), + "Qwen3_5TextModel": ("encoders", "qwen3_5", "Qwen3_5TextModel"), } _IMAGE_ENCODER_MODELS: dict[str, tuple] = { diff --git a/fastvideo/pipelines/basic/pan2/__init__.py b/fastvideo/pipelines/basic/pan2/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/pipelines/basic/pan2/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/pipelines/basic/pan2/pan2_pipeline.py b/fastvideo/pipelines/basic/pan2/pan2_pipeline.py new file mode 100644 index 000000000..f3f7dc3f5 --- /dev/null +++ b/fastvideo/pipelines/basic/pan2/pan2_pipeline.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 video generation: text-to-video, and image-to-video from a first frame.""" + +import torch + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.pipelines.basic.pan2.stages import ( + PAN2ConditioningStage, + PAN2DenoisingStage, + PAN2InputValidationStage, + PAN2TextEncodingStage, + PAN2TimestepPreparationStage, +) +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.stages import DecodingStage, LatentPreparationStage + + +class PAN2Pipeline(ComposedPipelineBase): + """One stage chain serves both modes: a request with ``image_path`` (or a PIL ``pil_image``) runs image-to-video + from that first frame, and one without runs text-to-video.""" + + _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage(stage_name="input_validation_stage", stage=PAN2InputValidationStage()) + self.add_stage(stage_name="prompt_encoding_stage", + stage=PAN2TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + self.add_stage(stage_name="timestep_preparation_stage", + stage=PAN2TimestepPreparationStage(scheduler=self.get_module("scheduler"))) + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage(scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + latent_dtype=torch.float32)) + self.add_stage(stage_name="conditioning_stage", stage=PAN2ConditioningStage(vae=self.get_module("vae"))) + self.add_stage(stage_name="denoising_stage", + stage=PAN2DenoisingStage(transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + pipeline=self)) + self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self)) + + +class PAN2ModularPipeline(PAN2Pipeline): + """The pipeline class PAN2 checkpoints declare in their ``modular_model_index.json``.""" + + +EntryClass = [PAN2Pipeline, PAN2ModularPipeline] diff --git a/fastvideo/pipelines/basic/pan2/presets.py b/fastvideo/pipelines/basic/pan2/presets.py new file mode 100644 index 000000000..10409a1d5 --- /dev/null +++ b/fastvideo/pipelines/basic/pan2/presets.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference presets for PAN2 text-to-video and image-to-video.""" + +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="PAN2 flow-matching denoising", + allowed_overrides=frozenset({"num_inference_steps", "guidance_scale"}), +) + +# 704x1248, 313 frames (13 s at 24 fps), 50 steps, CFG 3. The preset sends no negative prompt (an empty one is encoded +# as a single space); the reference sampler sends the release's structured negative caption, so pass it as the +# request's ``negative_prompt`` to follow that recipe. +_SHARED_DEFAULTS = { + "height": 704, + "width": 1248, + "num_frames": 313, + "fps": 24, + "num_inference_steps": 50, + "guidance_scale": 3.0, + "negative_prompt": "", +} + +PAN2_T2V = InferencePreset( + name="pan2_t2v", + version=1, + model_family="pan2", + description="PAN2 text-to-video", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +PAN2_I2V = InferencePreset( + name="pan2_i2v", + version=1, + model_family="pan2", + description="PAN2 image-to-video from a first frame", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults=dict(_SHARED_DEFAULTS), +) + +ALL_PRESETS = (PAN2_T2V, PAN2_I2V) + +__all__ = ["ALL_PRESETS", "PAN2_I2V", "PAN2_T2V"] diff --git a/fastvideo/pipelines/basic/pan2/stages.py b/fastvideo/pipelines/basic/pan2/stages.py new file mode 100644 index 000000000..6ef9001b2 --- /dev/null +++ b/fastvideo/pipelines/basic/pan2/stages.py @@ -0,0 +1,210 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2-specific pipeline stages: prompt encoding, first-frame conditioning, sigma schedule and fp32 sampling.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import numpy as np +import PIL.Image +import torch +import torchvision.transforms.functional as TF +from torch.distributed.tensor import DTensor + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.pipelines.stages.input_validation import InputValidationStage +from fastvideo.pipelines.stages.text_encoding import TextEncodingStage +from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult +from fastvideo.utils import PRECISION_TO_TYPE + +_USER_HEADER = "<|im_start|>user\n" + + +class PAN2InputValidationStage(InputValidationStage): + """Rejects geometry the PAN2 VAE and patch embedding cannot represent.""" + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + batch = super().forward(batch, fastvideo_args) + pipeline_config = fastvideo_args.pipeline_config + vae_arch = pipeline_config.vae_config.arch_config + spatial_multiple = vae_arch.spatial_compression_ratio * pipeline_config.dit_config.arch_config.patch_size[1] + if batch.height % spatial_multiple != 0 or batch.width % spatial_multiple != 0: + raise ValueError(f"PAN2 height and width must be divisible by {spatial_multiple}, " + f"got {batch.height} and {batch.width}.") + if (batch.num_frames - 1) % vae_arch.temporal_compression_ratio != 0: + raise ValueError(f"PAN2 num_frames - 1 must be divisible by {vae_arch.temporal_compression_ratio}, " + f"got {batch.num_frames}.") + if batch.pil_image is not None and not isinstance(batch.pil_image, PIL.Image.Image): + raise ValueError("PAN2 image-to-video expects the first frame as a PIL image or an image_path.") + return batch + + +class PAN2TextEncodingStage(TextEncodingStage): + """Encodes a prompt with PAN2's chat template. + + PAN2 conditions on the tokens after the ``<|im_start|>user\\n`` header (the prompt and the assistant header), + read from the text encoder's hidden states with its last ``text_encoder_skip_layers`` layers skipped. + """ + + @staticmethod + def tokenize(tokenizer: Any, prompt: str, system_prompt: str, max_length: int) -> tuple[torch.Tensor, int]: + """Token ids truncated to ``max_length`` tokens after the user header, and where that header ends.""" + # An empty prompt is encoded as a single space, as in training. + messages = [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt or " "}] + input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, + return_dict=True)["input_ids"] + if isinstance(input_ids, torch.Tensor): + input_ids = input_ids.tolist() + if input_ids and isinstance(input_ids[0], list): + input_ids = input_ids[0] + header_ids = tokenizer(_USER_HEADER, add_special_tokens=False)["input_ids"] + crop_start = next((i + len(header_ids) + for i in range(len(input_ids)) if input_ids[i:i + len(header_ids)] == header_ids), None) + if crop_start is None: + raise ValueError(f"The PAN2 tokenizer's chat template has no {_USER_HEADER!r} header.") + return torch.tensor([input_ids[:crop_start + max_length]], dtype=torch.long), crop_start + + @torch.no_grad() + def encode_text( + self, + text: str | list[str], + fastvideo_args: FastVideoArgs, + encoder_index: int | list[int] | None = None, + return_attention_mask: bool = False, + return_type: str = "list", + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + max_length: int | Sequence[int | None] | None = None, + truncation: bool | None = None, + padding: bool | str | None = None, + ): + pipeline_config = fastvideo_args.pipeline_config + texts = [text] if isinstance(text, str) else list(text) + if len(texts) != 1: + raise ValueError(f"PAN2 encodes one prompt per request, got {len(texts)}.") + if isinstance(max_length, Sequence): + max_length = max_length[0] + max_length = max_length or pipeline_config.text_max_sequence_length + + tokenizer = self.tokenizers[0] + text_encoder = self.text_encoders[0] + target_device = torch.device(device) if device is not None else get_local_torch_device() + # A module parked on CPU by text_encoder_cpu_offload computes on the target device and goes back afterwards; + # FSDP-wrapped (DTensor) encoders move their own shards. + first_param = next(text_encoder.parameters(), None) + input_device = first_param.device if first_param is not None else target_device + moved_for_forward = False + if (first_param is not None and not isinstance(first_param, DTensor) + and input_device.type != target_device.type): + text_encoder = text_encoder.to(target_device) + input_device = target_device + moved_for_forward = True + + input_ids, crop_start = self.tokenize(tokenizer, texts[0], pipeline_config.text_system_prompt, max_length) + with set_forward_context(current_timestep=0, attn_metadata=None): + outputs = text_encoder(input_ids=input_ids.to(input_device), output_hidden_states=True) + prompt_embeds = outputs.hidden_states[-(pipeline_config.text_encoder_skip_layers + 1)][:, crop_start:] + # fp32 is exact for the bf16 encoder output; the denoiser casts the embeddings to the DiT dtype. + prompt_embeds = prompt_embeds.to(device=target_device, dtype=dtype or torch.float32, copy=True) + attention_mask = torch.ones(prompt_embeds.shape[:2], dtype=torch.long, device=target_device) + if moved_for_forward and fastvideo_args.text_encoder_cpu_offload: + text_encoder.to("cpu") + return self.return_embeds([prompt_embeds], [attention_mask], return_type, return_attention_mask, [0]) + + +class PAN2TimestepPreparationStage(TimestepPreparationStage): + """PAN2 samples on uniformly spaced sigmas from 1 down to (excluding) 0; the scheduler applies the shift.""" + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + if batch.timesteps is None and batch.sigmas is None: + batch.sigmas = np.linspace(1.0, 0.0, batch.num_inference_steps + 1)[:-1].tolist() + return super().forward(batch, fastvideo_args) + + +class PAN2ConditioningStage(PipelineStage): + """Builds the conditioning latents and mask that PAN2 stacks after the noisy latents along channels. + + Text-to-video conditions on zeros. Image-to-video fills the first latent frame with the normalized VAE latents + of the first frame, scaled to cover the output and center cropped, and sets that frame's mask channel to one. + """ + + def __init__(self, vae) -> None: + super().__init__() + self.vae = vae + + @staticmethod + def preprocess_image(image: PIL.Image.Image, height: int, width: int) -> torch.Tensor: + """``[1, 3, 1, H, W]`` pixels in ``[-1, 1]``: scale to cover the canvas, round the size, center crop.""" + image = image.convert("RGB") + scale = max(width / image.width, height / image.height) + image = image.resize((round(image.width * scale), round(image.height * scale)), resample=PIL.Image.LANCZOS) + left, top = round((image.width - width) / 2), round((image.height - height) / 2) + image = image.crop((left, top, left + width, top + height)) + return (TF.to_tensor(image) * 2.0 - 1.0)[None, :, None] + + def _encode_first_frame(self, image: PIL.Image.Image, batch: ForwardBatch, + fastvideo_args: FastVideoArgs) -> torch.Tensor: + device = get_local_torch_device() + vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + self.vae = self.vae.to(device) + pixels = self.preprocess_image(image, batch.height, batch.width).to(device=device, dtype=vae_dtype) + latents = self.vae.encode(pixels).mode().float() + shift = self.vae.shift_factor.to(latents.device, latents.dtype) + scale = self.vae.scaling_factor.to(latents.device, latents.dtype) + if fastvideo_args.vae_cpu_offload: + self.vae.to("cpu") + return (latents - shift) * scale + + @torch.no_grad() + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + latents = batch.latents + batch_size, num_channels, num_frames, height, width = latents.shape + condition = torch.zeros(batch_size, + num_channels + 1, + num_frames, + height, + width, + device=latents.device, + dtype=torch.float32) + if batch.pil_image is not None: + image_latents = self._encode_first_frame(batch.pil_image, batch, fastvideo_args) + condition[:, :num_channels, :1] = image_latents.to(condition.device) + condition[:, num_channels:, :1] = 1.0 + batch.image_latent = condition + return batch + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + result.add_check("height", batch.height, V.positive_int) + result.add_check("width", batch.width, V.positive_int) + return result + + def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)]) + return result + + +class PAN2DenoisingStage(DenoisingStage): + """Keeps the latents in fp32 across the loop; they are cast to the DiT dtype for each forward. + + Image-to-video leaves the conditioning image's latent frame out of classifier-free guidance: that frame takes the + positive prediction, as in the reference sampler. + """ + + def finish_guidance(self, noise_pred, noise_pred_text, batch, fastvideo_args) -> torch.Tensor: + if batch.pil_image is not None and fastvideo_args.pipeline_config.i2v_skip_first_latent_cfg: + noise_pred[:, :, :1] = noise_pred_text[:, :, :1] + return noise_pred + + def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor: + return noise_pred.float() diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index acdc478ae..7e20544a8 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -472,6 +472,9 @@ class DenoisingStage(PipelineStage): noise_pred_text, guidance_rescale=batch.guidance_rescale, ) + + noise_pred = self.finish_guidance(noise_pred, noise_pred_text, batch, fastvideo_args) + noise_pred = self.prepare_scheduler_model_output(noise_pred, state) if scheduler_fp32: # Diffusers-style: fp32 Euler update outside autocast avoids BF16 drift. latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0] @@ -605,6 +608,14 @@ class DenoisingStage(PipelineStage): def prepare_timestep(self, timestep, latent_model_input, batch, state) -> torch.Tensor: return timestep.repeat(latent_model_input.shape[0]) + def finish_guidance(self, noise_pred, noise_pred_text, batch, fastvideo_args) -> torch.Tensor: + """The guided prediction handed on to the scheduler; a model overrides it to keep part of the output out of CFG.""" + return noise_pred + + def prepare_scheduler_model_output(self, noise_pred, state) -> torch.Tensor: + """The model output handed to the scheduler step, which returns the next latents in its dtype.""" + return noise_pred + def finish_step(self, latents, state) -> torch.Tensor: return latents diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index e0116af9a..ef0e94df8 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -30,11 +30,16 @@ class LatentPreparationStage(PipelineStage): denoised during the diffusion process. """ - def __init__(self, scheduler, transformer, use_btchw_layout: bool = False) -> None: + def __init__(self, + scheduler, + transformer, + use_btchw_layout: bool = False, + latent_dtype: torch.dtype | None = None) -> None: super().__init__() self.scheduler = scheduler self.transformer = transformer self.use_btchw_layout = use_btchw_layout + self.latent_dtype = latent_dtype def forward( self, @@ -88,7 +93,7 @@ class LatentPreparationStage(PipelineStage): batch.prompt_embeds = [dummy_prompt] batch.negative_prompt_embeds = [] batch.do_classifier_free_guidance = False - dtype = batch.prompt_embeds[0].dtype + dtype = self.latent_dtype or batch.prompt_embeds[0].dtype device = get_local_torch_device() generator = batch.generator latents = batch.latents diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 0a0d18361..5efe6b10e 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -10,6 +10,7 @@ from __future__ import annotations import dataclasses import os +import re from collections.abc import Callable from functools import lru_cache from typing import TYPE_CHECKING, Any @@ -44,6 +45,7 @@ from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.minimax_h3 import MiniMaxH3PipelineConfig from fastvideo.configs.pipelines.pika_h3_flex import PikaH3FlexPipelineConfig from fastvideo.configs.pipelines.open_h3_regenerate_2k import OpenH3Regenerate2KPipelineConfig +from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig from fastvideo.configs.pipelines.mmaudio import ( MMAudioLarge44kV2PipelineConfig, MMAudioLarge44kV2AConfig, @@ -1227,6 +1229,22 @@ def _register_configs() -> None: pipeline_cls_name="MiniMaxH3ModularPipeline", ) + # PAN2: one pipeline runs text-to-video, and image-to-video when the request carries a first frame. + def _is_pan2(name: str) -> bool: + name = name.lower() + return "pan2modularpipeline" in name or re.search(r"(^|[/_-])pan-?2([/_.-]|$)", name) is not None + + register_configs( + sampling_param_cls=None, + pipeline_config_cls=PAN2PipelineConfig, + workload_types=(WorkloadType.T2V, WorkloadType.I2V), + hf_model_paths=["IFM/PAN2"], + model_detectors=[_is_pan2], + model_family="pan2", + default_preset="pan2_t2v", + pipeline_cls_name="PAN2ModularPipeline", + ) + # SD3.5 register_configs( sampling_param_cls=None, @@ -1404,6 +1422,8 @@ def _register_presets() -> None: ALL_PRESETS as OPEN_H3_REGENERATE_2K_PRESETS, ) from fastvideo.pipelines.basic.mmaudio.presets import ( ALL_PRESETS as MMAUDIO_PRESETS, ) + from fastvideo.pipelines.basic.pan2.presets import ( + ALL_PRESETS as PAN2_PRESETS, ) from fastvideo.pipelines.basic.sd35.presets import ( ALL_PRESETS as SD35_PRESETS, ) from fastvideo.pipelines.basic.stable_audio.presets import ( @@ -1438,6 +1458,7 @@ def _register_presets() -> None: PIKA_H3_FLEX_PRESETS, OPEN_H3_REGENERATE_2K_PRESETS, MMAUDIO_PRESETS, + PAN2_PRESETS, SD35_PRESETS, STABLE_AUDIO_PRESETS, WAN_PRESETS, diff --git a/fastvideo/tests/entrypoints/test_serving_family_baselines.py b/fastvideo/tests/entrypoints/test_serving_family_baselines.py index bb7cf8b24..be67249f3 100644 --- a/fastvideo/tests/entrypoints/test_serving_family_baselines.py +++ b/fastvideo/tests/entrypoints/test_serving_family_baselines.py @@ -25,6 +25,7 @@ from fastvideo.entrypoints.openai.request_adapter import ( prepare_reference_media, ) from fastvideo.pipelines.basic.hunyuan15.presets import HUNYUAN15_I2V_480P_DISTILLED +from fastvideo.pipelines.basic.pan2.presets import PAN2_T2V ROOT = Path(__file__).resolve().parents[3] @@ -37,6 +38,7 @@ BASELINES = { "openai_kandinsky6_ti2va.yaml": "t2v", "openai_kandinsky6_ti2va_distilled.yaml": "t2v", "openai_ltx2_distilled.yaml": "t2v", + "openai_pan2.yaml": "t2v", "openai_flux1_dev.yaml": "t2i", "openai_flux2_dev.yaml": "t2i", "openai_flux2_klein.yaml": "t2i", @@ -112,6 +114,32 @@ def test_hunyuan_i2v_480p_keeps_the_registered_checkpoint_semantics(tmp_path: Pa assert adapted.output.output_path.endswith(".mp4") +def test_pan2_serving_matches_its_preset_and_runs_image_to_video_from_a_reference(tmp_path: Path) -> None: + config = _load("openai_pan2.yaml") + sampling = config.default_request.sampling + assert config.generator.model_path == "IFM/PAN2" + assert config.generator.engine.offload.dit is False + assert config.generator.engine.offload.dit_layerwise is False + for field in ("height", "width", "num_frames", "fps", "num_inference_steps", "guidance_scale"): + assert getattr(sampling, field) == PAN2_T2V.defaults[field], field + + image = tmp_path / "first-frame.png" + image.write_bytes(b"not-decoded-at-admission") + for reference in (None, {"image_url": str(image)}): + request = VideoGenerationRequest(prompt="a red fox", image_reference=reference) + adapted = build_generation_request( + "pan2_test", + request, + _args(config, "t2v"), + served_model_name=config.server.served_model_name or "", + output_dir=str(tmp_path), + default_request=config.default_request, + ) + assert adapted.sampling.num_frames == 313 + assert adapted.inputs.image_path == (str(image) if reference else None) + assert adapted.output.output_path.endswith(".mp4") + + def test_stable_audio_seconds_override_produces_wav(tmp_path: Path) -> None: config = _load("openai_stable_audio.yaml") request = VideoGenerationRequest(prompt="rain on a metal roof", seconds=6) diff --git a/fastvideo/tests/stages/test_pan2_first_latent_cfg.py b/fastvideo/tests/stages/test_pan2_first_latent_cfg.py new file mode 100644 index 000000000..e74314091 --- /dev/null +++ b/fastvideo/tests/stages/test_pan2_first_latent_cfg.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 image-to-video keeps the first latent frame out of classifier-free guidance; no weights or CUDA needed.""" + +import pytest +import torch + +from fastvideo.tests.stages._denoising_fixtures import ( + NullProgressBar, TinyDenoiser, TinyScheduler, _patch_denoising_module, _tiny_args, _tiny_batch, +) + + +def _one_step(monkeypatch, env_overrides, *, image_to_video: bool, skip_first_latent: bool, guidance_rescale=0.0, + frame_values=(0.0, 0.0, 0.0)): + _patch_denoising_module(monkeypatch, env_overrides, None) + from fastvideo.pipelines.basic.pan2.stages import PAN2DenoisingStage + + stage = PAN2DenoisingStage(TinyDenoiser(), TinyScheduler()) + stage.progress_bar = lambda iterable=None, total=None: NullProgressBar() + batch = _tiny_batch() + batch.latents = torch.tensor(frame_values).reshape(1, 1, 3, 1, 1) + batch.raw_latent_shape = (1, 1, 3, 1, 1) + batch.timesteps = torch.tensor([0.0]) + batch.num_inference_steps = 1 + batch.pil_image = object() if image_to_video else None + batch.guidance_rescale = guidance_rescale + args = _tiny_args() + args.pipeline_config.i2v_skip_first_latent_cfg = skip_first_latent + args.pipeline_config.scheduler_step_in_fp32 = True + return stage.forward(batch, args).latents.flatten().tolist() + + +def test_image_to_video_takes_the_positive_prediction_for_the_first_latent_frame(monkeypatch, env_overrides): + # With a zero latent and timestep the cond prediction is 0.5 and the uncond -0.5, so CFG 2 guides to 1.5. + latents = _one_step(monkeypatch, env_overrides, image_to_video=True, skip_first_latent=True) + + assert latents == pytest.approx([-0.005, -0.015, -0.015]) + + +@pytest.mark.parametrize("image_to_video,skip_first_latent", [(True, False), (False, True)]) +def test_every_frame_is_guided_without_a_conditioning_image_or_with_the_skip_off(monkeypatch, env_overrides, + image_to_video, skip_first_latent): + latents = _one_step(monkeypatch, env_overrides, image_to_video=image_to_video, skip_first_latent=skip_first_latent) + + assert latents == pytest.approx([-0.015, -0.015, -0.015]) + + +def test_first_latent_frame_keeps_the_positive_prediction_after_guidance_rescale(monkeypatch, env_overrides): + # Frames 0, 1, 2 predict 0.5, 0.7, 0.9 (cond) and 1.5, 1.7, 1.9 (CFG 2), so the rescale factor is 1 and only the + # first frame's replacement by 0.5 remains. Replacing before the rescale would scale it by std(cond)/std(mixed). + latents = _one_step(monkeypatch, + env_overrides, + image_to_video=True, + skip_first_latent=True, + guidance_rescale=1.0, + frame_values=(0.0, 1.0, 2.0)) + + assert latents == pytest.approx([-0.005, 0.983, 1.981]) diff --git a/fastvideo/tests/stages/test_pan2_latent_preparation.py b/fastvideo/tests/stages/test_pan2_latent_preparation.py new file mode 100644 index 000000000..80c09dd4d --- /dev/null +++ b/fastvideo/tests/stages/test_pan2_latent_preparation.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 draws its initial noise in fp32 whatever dtype the prompt embeddings arrive in; no weights or CUDA needed.""" + +import types + +import pytest +import torch + +from fastvideo.pipelines.basic.pan2 import pan2_pipeline +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages import LatentPreparationStage, latent_preparation + + +@pytest.fixture() +def cpu_device(monkeypatch): + monkeypatch.setattr(latent_preparation, "get_local_torch_device", lambda: torch.device("cpu")) + + +def _args() -> types.SimpleNamespace: + arch = types.SimpleNamespace(temporal_compression_ratio=4, spatial_compression_ratio=16) + return types.SimpleNamespace(pipeline_config=types.SimpleNamespace(vae_config=types.SimpleNamespace(arch_config=arch))) + + +def _noise(stage, embeds_dtype: torch.dtype) -> torch.Tensor: + batch = ForwardBatch(data_type="video", + prompt="a cat", + prompt_embeds=[torch.zeros(1, 4, 8, dtype=embeds_dtype)], + generator=torch.Generator("cpu").manual_seed(7), + height=32, + width=32, + num_frames=5) + return stage.forward(batch, _args()).latents + + +def test_noise_is_fp32_and_independent_of_the_prompt_embedding_dtype(cpu_device): + transformer = types.SimpleNamespace(num_channels_latents=3) + stage = LatentPreparationStage(scheduler=types.SimpleNamespace(), + transformer=transformer, + latent_dtype=torch.float32) + + from_bf16, from_fp32 = _noise(stage, torch.bfloat16), _noise(stage, torch.float32) + + assert from_bf16.dtype == torch.float32 + assert torch.equal(from_bf16, from_fp32) + assert _noise(LatentPreparationStage(types.SimpleNamespace(), transformer), torch.bfloat16).dtype == torch.bfloat16 + + +def test_pan2_pipeline_requests_fp32_latents(monkeypatch): + seen = {} + + class _Stage: + + def __init__(self, *args, **kwargs): + seen.setdefault("kwargs", []).append(kwargs) + + for name in ("PAN2InputValidationStage", "PAN2TextEncodingStage", "PAN2TimestepPreparationStage", + "PAN2ConditioningStage", "PAN2DenoisingStage", "DecodingStage"): + monkeypatch.setattr(pan2_pipeline, name, _Stage) + monkeypatch.setattr(pan2_pipeline, "LatentPreparationStage", _Stage) + pipeline = pan2_pipeline.PAN2Pipeline.__new__(pan2_pipeline.PAN2Pipeline) + stages = {} + pipeline.add_stage = lambda stage_name, stage: stages.__setitem__(stage_name, stage) + pipeline.get_module = lambda name: name + + pipeline.create_pipeline_stages(None) + + assert any(kwargs.get("latent_dtype") is torch.float32 for kwargs in seen["kwargs"]) diff --git a/tests/local_tests/encoders/test_pan2_qwen3_5_parity.py b/tests/local_tests/encoders/test_pan2_qwen3_5_parity.py new file mode 100644 index 000000000..35a835d82 --- /dev/null +++ b/tests/local_tests/encoders/test_pan2_qwen3_5_parity.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3.5 text-encoder parity for PAN2 against Transformers' ``Qwen3_5TextModel``. + +Coverage scope: both. The checkpoint test loads the public tiny PAN2 text +encoder through FastVideo's production ``TextEncoderLoader`` and feeds both +models the exact PAN2 chat-templated token ids. The random-config tests use a +larger hybrid stack (GQA, partial RoPE, more value than key heads in the Gated +DeltaNet, sequences that cross several 64-token delta-rule chunks) and also +cover bf16 and padded batches. + +Everything runs on CPU. Set ``PAN2_TINY_MODEL_DIR`` to a local copy of +``wuqing157/tiny-pan2-modular-pipe``; otherwise the HF cache is used and the +checkpoint test skips when it is absent. +""" + +from __future__ import annotations + +import copy +import os +import socket +from pathlib import Path +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from torch.testing import assert_close + +from fastvideo.configs.models.encoders.qwen3_5 import Qwen3_5TextConfig +from fastvideo.forward_context import set_forward_context +from fastvideo.models.loader.utils import set_default_torch_dtype + +PARITY_SCOPE = "both" +TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe" +PAN2_SYSTEM_MESSAGE = "You are a helpful assistant. Generate the video with the following instructions." +PAN2_TEXT_ENCODER_LAYERS_TO_SKIP = 2 +PAN2_PROMPTS = ( + "A red fox crosses fresh snow at sunrise.", + "A slow dolly shot glides through a rain-soaked neon market at night. Steam rises from food stalls while " + "vendors in bright raincoats shout over the hum of generators. Reflections of pink, teal and amber signs ripple " + "across puddles as a cyclist weaves between crowds, her umbrella catching droplets that sparkle in the light. " + "In the background, a tram rattles past an old clock tower, its windows glowing warm yellow. The camera rises to " + "reveal rooftops crowded with antennas and laundry lines, then tilts toward a thunderstorm rolling in over the " + "harbor, lightning briefly silhouetting cranes and container ships. Cinematic lighting, shallow depth of field, " + "35mm film grain, realistic motion blur.", +) +FP32_TOLERANCE = {"atol": 1e-5, "rtol": 1e-5} +BF16_TOLERANCE = {"atol": 2e-2, "rtol": 2e-2} + + +def _hf_qwen3_5() -> tuple[Any, Any]: + module = pytest.importorskip("transformers.models.qwen3_5", reason="transformers without Qwen3.5 support") + return module.Qwen3_5TextConfig, module.Qwen3_5TextModel + + +def _tiny_model_dir() -> Path | None: + local = os.getenv("PAN2_TINY_MODEL_DIR") + if local: + return Path(local) + try: + from huggingface_hub import snapshot_download + + return Path(snapshot_download(TINY_REPO_ID, local_files_only=True)) + except Exception: + return None + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +@pytest.fixture(scope="module") +def cpu_single_rank(): + """Single-rank CPU runtime; the loader would otherwise target ``mps`` on CPU-only Linux hosts.""" + import fastvideo.distributed.parallel_state as parallel_state + import fastvideo.models.loader.component_loader as component_loader + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + + with pytest.MonkeyPatch.context() as patch: + patch.setattr(parallel_state, "get_local_torch_device", lambda: torch.device("cpu")) + patch.setattr(component_loader, "get_local_torch_device", lambda: torch.device("cpu")) + for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"), + ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")): + patch.setenv(name, value) + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + cleanup_dist_env_and_memory() + + +def _loader_args(config: Qwen3_5TextConfig, precision: str = "fp32") -> SimpleNamespace: + return SimpleNamespace( + pipeline_config=SimpleNamespace(text_encoder_configs=(config, ), text_encoder_precisions=(precision, )), + text_encoder_cpu_offload=False, + override_text_encoder_quant=None, + override_text_encoder_safetensors=None, + pin_cpu_memory=False, + disable_offload_on_unified_memory=lambda device_id, offload_flag=None: False, + ) + + +def _pan2_input_ids(tokenizer: Any, prompt: str) -> tuple[torch.Tensor, int]: + """Tokenize exactly like PAN2's ``encode_prompt`` and return ids plus the user-header crop.""" + messages = [{"role": "system", "content": PAN2_SYSTEM_MESSAGE}, {"role": "user", "content": prompt or " "}] + input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, + return_dict=True).input_ids + header_ids = tokenizer("<|im_start|>user\n", add_special_tokens=False).input_ids + crop_start = next( + (i + len(header_ids) for i in range(len(input_ids)) if input_ids[i:i + len(header_ids)] == header_ids), + None, + ) + assert crop_start is not None, "PAN2 chat template has no user header" + return torch.tensor([input_ids]), crop_start + + +def _run_both(reference: torch.nn.Module, native: torch.nn.Module, **inputs: Any) -> tuple[Any, Any]: + with torch.inference_mode(): + expected = reference(**inputs, output_hidden_states=True) + with set_forward_context(current_timestep=0, attn_metadata=None): + actual = native(**inputs, output_hidden_states=True) + return expected, actual + + +def _assert_hidden_states_close( + expected: tuple[torch.Tensor, ...], + actual: tuple[torch.Tensor, ...] | None, + label: str, + tolerance: dict[str, float], + valid: torch.Tensor | None = None, +) -> float: + assert actual is not None, f"{label}: native encoder returned no hidden_states" + assert len(actual) == len(expected), f"{label}: {len(actual)} hidden states, reference has {len(expected)}" + worst = 0.0 + for index, (reference_state, native_state) in enumerate(zip(expected, actual, strict=True)): + assert native_state.shape == reference_state.shape, f"{label} hidden_states[{index}] shape" + assert torch.isfinite(native_state).all(), f"{label} hidden_states[{index}] has non-finite values" + if valid is not None: + reference_state = reference_state[valid] + native_state = native_state[valid] + drift = (native_state.float() - reference_state.float()).abs().max().item() + worst = max(worst, drift) + print(f"{label} hidden_states[{index}]: max_abs={drift:.3e}", flush=True) + assert_close(native_state, reference_state, **tolerance, msg=lambda m: f"{label} hidden_states[{index}]: {m}") + return worst + + +def test_pan2_qwen3_5_arch_config_matches_transformers_rope_and_layer_types() -> None: + """Config.json fields resolve to the same RoPE settings and layer schedule as Transformers.""" + hf_config_cls, _ = _hf_qwen3_5() + cases = ( + {}, + { + "rope_theta": 1_000_000.0 + }, + { + "partial_rotary_factor": 0.5, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 5_000_000.0 + } + }, + { + "partial_rotary_factor": 0.5, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 5_000_000.0, + "partial_rotary_factor": 0.25 + } + }, + { + "num_hidden_layers": 6, + "full_attention_interval": 3 + }, + { + "num_hidden_layers": 3, + "layer_types": ["mamba", "attention", "linear_attention"] + }, + ) + for case in cases: + reference = hf_config_cls.from_dict(dict(case)) + config = Qwen3_5TextConfig() + config.update_model_arch(dict(case)) + rope = config.arch_config.resolved_rope_parameters() + for key in ("rope_type", "rope_theta", "partial_rotary_factor"): + assert rope[key] == reference.rope_parameters[key], f"{case}: {key}" + assert config.arch_config.resolved_layer_types() == list(reference.layer_types), case + + +def test_pan2_qwen3_5_tiny_checkpoint_loader_parity(cpu_single_rank) -> None: + """Production-loader encoder matches every HF hidden state on real PAN2 prompts (fp32 CPU).""" + _, hf_model_cls = _hf_qwen3_5() + model_dir = _tiny_model_dir() + if model_dir is None or not (model_dir / "text_encoder" / "config.json").is_file(): + pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR") + if not (model_dir / "tokenizer").is_dir(): + pytest.skip(f"{TINY_REPO_ID} tokenizer is not available locally") + from transformers import AutoTokenizer + + from fastvideo.models.encoders.qwen3_5 import Qwen3_5TextModel + from fastvideo.models.loader.component_loader import TextEncoderLoader + + text_encoder_dir = model_dir / "text_encoder" + tokenizer = AutoTokenizer.from_pretrained(model_dir / "tokenizer", local_files_only=True) + reference, loading_info = hf_model_cls.from_pretrained( + text_encoder_dir, + local_files_only=True, + dtype=torch.float32, + output_loading_info=True, + ) + load_errors = { + name: loading_info.get(name) + for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") if loading_info.get(name) + } + assert not load_errors, f"Reference Qwen3.5 checkpoint did not load strictly: {load_errors}" + reference = reference.eval() + + native = TextEncoderLoader().load(str(text_encoder_dir), _loader_args(Qwen3_5TextConfig())) + assert isinstance(native, Qwen3_5TextModel) + assert native.layer_types == list(reference.config.layer_types) + assert next(native.parameters()).dtype == torch.float32 + assert_close(native.rotary_emb.inv_freq, reference.rotary_emb.inv_freq.float(), atol=0.0, rtol=0.0) + + for prompt_index, prompt in enumerate(PAN2_PROMPTS): + input_ids, crop_start = _pan2_input_ids(tokenizer, prompt) + label = f"tiny prompt{prompt_index} (L={input_ids.shape[1]})" + expected, actual = _run_both(reference, native, input_ids=input_ids) + worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, label, FP32_TOLERANCE) + assert_close(actual.last_hidden_state, expected.last_hidden_state, **FP32_TOLERANCE) + + skip = PAN2_TEXT_ENCODER_LAYERS_TO_SKIP + 1 + reference_embeds = expected.hidden_states[-skip][:, crop_start:] + native_embeds = actual.hidden_states[-skip][:, crop_start:] + assert_close(native_embeds, reference_embeds, **FP32_TOLERANCE) + print( + f"{label} PAN2 prompt_embeds hidden_states[-{skip}][:, {crop_start}:] " + f"max_abs={(native_embeds - reference_embeds).abs().max().item():.3e} all_layers_max_abs={worst:.3e}", + flush=True, + ) + + +@pytest.fixture(scope="module") +def random_pair(cpu_single_rank): + """A randomly initialized 8-layer hybrid reference and the native model loaded from its state dict.""" + hf_config_cls, hf_model_cls = _hf_qwen3_5() + from fastvideo.models.encoders.qwen3_5 import Qwen3_5TextModel + + torch.manual_seed(0) + hf_config = hf_config_cls( + vocab_size=1024, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=8, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=32, + linear_conv_kernel_dim=4, + linear_key_head_dim=16, + linear_value_head_dim=24, + linear_num_key_heads=2, + linear_num_value_heads=4, + layer_types=[ + "linear_attention", + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + "linear_attention", + "linear_attention", + "full_attention", + ], + rope_parameters={ + "rope_type": "default", + "rope_theta": 1_000_000.0, + "partial_rotary_factor": 0.25 + }, + initializer_range=0.08, + ) + reference = hf_model_cls(hf_config).eval() + with torch.no_grad(): + # Move the zero-centered norms, gates and decay parameters off their + # init values so every term contributes. + for name, parameter in reference.named_parameters(): + if name.endswith("norm.weight") or name.endswith(("A_log", "dt_bias")): + parameter.add_(0.2 * torch.randn_like(parameter)) + + config = Qwen3_5TextConfig() + config.update_model_arch(hf_config.to_dict()) + with set_default_torch_dtype(torch.float32): + native = Qwen3_5TextModel(config).eval() + loaded = native.load_weights(reference.state_dict().items()) + expected_names = {name for name, _ in native.named_parameters()} + assert loaded == expected_names, f"unloaded native parameters: {sorted(expected_names - loaded)}" + assert set(reference.state_dict()) == expected_names + return reference, native + + +def test_pan2_qwen3_5_random_config_fp32_parity(random_pair) -> None: + """Multi-chunk delta rule (L=301), GQA and partial RoPE match HF in fp32.""" + reference, native = random_pair + generator = torch.Generator().manual_seed(1) + input_ids = torch.randint(0, reference.config.vocab_size, (2, 301), generator=generator) + expected, actual = _run_both(reference, native, input_ids=input_ids) + worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, "random fp32", FP32_TOLERANCE) + print(f"random fp32 all_layers_max_abs={worst:.3e}", flush=True) + # Each layer must move the residual stream, or the comparison is vacuous. + for index in range(1, len(expected.hidden_states) - 1): + change = (expected.hidden_states[index] - expected.hidden_states[index - 1]).abs().max().item() + assert change > 0.1, f"layer {index - 1} barely changes the hidden state ({change:.3e})" + + +def test_pan2_qwen3_5_random_config_padding_parity(random_pair) -> None: + """Right- and left-padded rows match HF on every non-padded position.""" + reference, native = random_pair + generator = torch.Generator().manual_seed(2) + input_ids = torch.randint(0, reference.config.vocab_size, (2, 200), generator=generator) + attention_mask = torch.ones_like(input_ids) + attention_mask[0, 150:] = 0 + attention_mask[1, :37] = 0 + expected, actual = _run_both(reference, native, input_ids=input_ids, attention_mask=attention_mask) + worst = _assert_hidden_states_close( + expected.hidden_states, + actual.hidden_states, + "random padded", + FP32_TOLERANCE, + valid=attention_mask.bool(), + ) + print(f"random padded all_layers_max_abs={worst:.3e}", flush=True) + + +def test_pan2_qwen3_5_random_config_bf16_parity(random_pair) -> None: + """Half-precision dtype handling (fp32 norms, gates and delta rule) follows HF.""" + reference, native = (copy.deepcopy(model).to(torch.bfloat16) for model in random_pair) + generator = torch.Generator().manual_seed(3) + input_ids = torch.randint(0, reference.config.vocab_size, (1, 173), generator=generator) + expected, actual = _run_both(reference, native, input_ids=input_ids) + assert actual.hidden_states is not None and actual.hidden_states[-1].dtype == torch.bfloat16 + worst = _assert_hidden_states_close(expected.hidden_states, actual.hidden_states, "random bf16", BF16_TOLERANCE) + print(f"random bf16 all_layers_max_abs={worst:.3e}", flush=True) diff --git a/tests/local_tests/pan2/PORT_STATUS.md b/tests/local_tests/pan2/PORT_STATUS.md new file mode 100644 index 000000000..7c5b53a3b --- /dev/null +++ b/tests/local_tests/pan2/PORT_STATUS.md @@ -0,0 +1,71 @@ +# PAN2 Port Status + +## Summary +- model_family: pan2 +- workload_types: T2V, I2V (one pipeline; I2V when the request carries a first frame) +- official_ref: vllm-project/vllm-omni#8602 at `e0c4c73158a5e4f44eda8eaa8879f7fd9cedbd74` (the diffusers integration is private) +- official_ref_dir: none; rebuilt in `tests/local_tests/pan2/reference.py` +- hf_weights_path: `IFM/PAN2` (not public); tests use `wuqing157/tiny-pan2-modular-pipe@3ff3130b` +- local_weights_dir: HF cache, or `PAN2_TINY_MODEL_DIR` +- source_layout: Modular Diffusers (`modular_model_index.json`), diffusers component subfolders +- local_tests_readme: `tests/local_tests/pan2/README.md` + +## Current Phase +- phase: 11 (post-parity review and handoff) +- status: in_progress (GPU and real-weight validation pending, see I001 and I002) +- owner: orchestrator +- last_updated: 2026-10-20 + +## Component Matrix +| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +|---|---|---|---|---|---|---|---|---|---| +| transformer | DiT | port | vLLM-Omni `pan2_transformer.py` (diffusers `PAN2Transformer3DModel`) | `transformer/config.json` | `fastvideo/models/dits/pan2.py`, `fastvideo/configs/models/dits/pan2.py` | done | load-time key mapping (Q/K/V fused) | non_skip_pass | I002 | +| text_encoder | encoder | port | transformers `Qwen3_5TextModel` | `text_encoder/config.json` | `fastvideo/models/encoders/qwen3_5.py`, `fastvideo/configs/models/encoders/qwen3_5.py` | done | passthrough | non_skip_pass | I002 | +| tokenizer | tokenizer | reuse | `Qwen2Tokenizer` | `tokenizer/` | AutoTokenizer | n/a | passthrough | covered by pipeline parity | — | +| vae | VAE | reuse | diffusers `AutoencoderKLWan` (Wan2.2, z_dim 48) | `vae/config.json` | `fastvideo/models/wan/vae.py` | done | passthrough | non_skip_pass | — | +| scheduler | scheduler | reuse | diffusers `FlowMatchEulerDiscreteScheduler` | `scheduler/scheduler_config.json` (shift 7) | `fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py` | done | passthrough | non_skip_pass | — | + +## Conversion State +- conversion_script: none; the transformer maps diffusers keys at load time (`PAN2ArchConfig.param_names_mapping`) +- converted_weights_dir: n/a +- source_layout: diffusers +- strict_load_status: tiny checkpoint loads strictly (no missing or unexpected keys) +- passthrough_components: text_encoder, tokenizer, vae, scheduler +- retry_history: none + +## Parity Commands +| Scope | Command | Last Result | Notes | +|---|---|---|---| +| DiT | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py` | 13 passed (CPU) | SP=2/4 vs one rank over gloo, T2V/I2V, one/two pad tokens; max abs 7.7e-6 | +| Text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | 5 passed (CPU) | exact | +| VAE | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | 9 passed (CPU) | exact | +| Scheduler | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | 9 passed (CPU) | exact | +| Pipeline | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | 3 passed (CPU) | T2V and I2V latents exact | + +## Open Questions +| ID | Question | Owner | Needed By Phase | Status | Resolution | +|---|---|---|---|---|---| +| Q001 | Which reference, with the diffusers integration private? | orchestrator | 2 | resolved | vLLM-Omni #8602, rebuilt from diffusers/transformers blocks in `reference.py` | +| Q002 | Port the `pan2-guardrail` prompt/video guardrails? | orchestrator | 7 | resolved | Not ported: FastVideo has no guardrail stage. | +| Q003 | Port Cache-DiT `quality` and CFG-parallel? | orchestrator | 7 | resolved | Not ported: vLLM-Omni serving features. Ulysses SP is supported through `DistributedAttention`. | + +## Issues And Blockers +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +|---|---|---|---|---|---|---|---|---| +| I001 | 9-10 | all | medium | No real-weight parity or SSIM regression | `IFM/PAN2` is not public | orchestrator | open | Add `fastvideo/tests/ssim/test_pan2_similarity.py` once weights are public. | +| I002 | 9 | transformer, text_encoder | medium | No GPU parity run yet (bf16 + FlashAttention, NCCL SP, 9B encoder speed) | Local parity tests ran on a CPU login node | orchestrator | open | Run the same tests on a GPU node. | +| I003 | 7 | pipeline | low | The shared denoiser casts latents to the model-output dtype in the scheduler step; PAN2 samples on fp32 latents | pipeline parity fails (8e-2) without the fix | orchestrator | resolved | `DenoisingStage.prepare_scheduler_model_output` hook; `PAN2DenoisingStage` returns `noise_pred.float()`. | + +## Escape Hatches +| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | +|---|---|---|---|---|---|---| + +## Decisions +| Date | Decision | Rationale | Impact | +|---|---|---|---| +| 2026-10-15 | Native Qwen3.5 encoder with a PyTorch chunked gated delta rule | No `fla`/`causal_conv1d` dependency; encoder runs once per prompt | Slower than fused kernels for long prompts | +| 2026-10-15 | Fuse the per-stream Q/K/V projections at load time | One GEMM per stream instead of three | Mapping in `PAN2ArchConfig.param_names_mapping` | +| 2026-10-15 | Text embeddings kept in fp32 after encoding | The shared latent stage draws noise in the embedding dtype; PAN2 uses fp32 noise and latents | Denoiser casts embeddings to the DiT dtype (`cast_prompt_embeds_to_dit_dtype`) | + +## Handoff Notes +- The remaining work is I001 and I002. Both need hardware or weights, not code changes. diff --git a/tests/local_tests/pan2/README.md b/tests/local_tests/pan2/README.md new file mode 100644 index 000000000..2fac8e256 --- /dev/null +++ b/tests/local_tests/pan2/README.md @@ -0,0 +1,67 @@ +# PAN2 Local Tests + +Runbook for validating the PAN2 port (text-to-video and first-frame image-to-video). These tests run manually and are +not part of the default CI suite. + +## Reference + +PAN2's diffusers integration (`PAN2Transformer3DModel`, `PAN2ModularPipeline`) is not public yet, and the full +`IFM/PAN2` weights are not released. The reference is the vLLM-Omni port, +[vllm-project/vllm-omni#8602](https://github.com/vllm-project/vllm-omni/pull/8602) (head `e0c4c731`), which reports +parity with the diffusers implementation. `reference.py` rebuilds it from diffusers and transformers building blocks: +the DiT from `AdaLayerNormZero`, `AdaLayerNormContinuous`, `FeedForward` (gelu-approximate and swiglu), `Timesteps`, +`TimestepEmbedding` and the diffusers RoPE helpers, using the checkpoint's module names so the checkpoint loads +strictly. The pipeline is `reference_generate`, a step-by-step copy of the vLLM-Omni `PAN2Pipeline.forward` +using `Qwen3_5TextModel`, `AutoencoderKLWan` and `FlowMatchEulerDiscreteScheduler`. + +The checkpoint is the public random-weight [`wuqing157/tiny-pan2-modular-pipe`](https://huggingface.co/wuqing157/tiny-pan2-modular-pipe) +(revision `3ff3130b`), the same one vLLM-Omni tests against. It has two DiT blocks, a four-layer Qwen3.5 text model +(three Gated DeltaNet layers and one full-attention layer) and a small Wan2.2-layout VAE. + +## Setup + +```bash +huggingface-cli download wuqing157/tiny-pan2-modular-pipe --revision 3ff3130b4527f61aaab7bde01465c9884dc84cdd +# or point the tests at a local copy: +export PAN2_TINY_MODEL_DIR=/path/to/tiny-pan2-modular-pipe +``` + +The tests need `diffusers` with `AutoencoderKLWan` and `transformers` with `Qwen3_5TextModel` (5.x). They run on CPU +when no GPU is present. On a GPU-less host, `import fastvideo` also needs the `fastvideo_kernel` stub described in +`fastvideo/tests/AGENTS.md`. + +## Tests + +| Area | Command | Result (CPU, fp32 unless noted) | +|---|---|---| +| DiT: tiny checkpoint, strict load + forward | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k tiny_checkpoint` | max abs 6.0e-7 | +| DiT: random weights, batch 1 and 2 | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k random_weights` | max abs 0 and 1.9e-5 (outputs ~16) | +| DiT: Ulysses SP=2/4 vs one rank, T2V/I2V with 15/18 tokens | `pytest tests/local_tests/transformers/test_pan2_transformer_parity.py -k sequence_parallel` | max abs 4.8e-6 / 7.7e-6 (gloo on CPU; NCCL on 2/4 GPUs) | +| Qwen3.5 text encoder | `pytest tests/local_tests/encoders/test_pan2_qwen3_5_parity.py` | all hidden states exact (fp32, bf16, padded batch) | +| Wan2.2 VAE encode/decode + normalization | `pytest tests/local_tests/vaes/test_pan2_vae_parity.py` | exact; normalization 2.4e-7 | +| Flow-match scheduler, PAN2 sigmas + shift | `pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py` | exact | +| Pipeline, T2V and I2V (bf16 DiT) | `pytest tests/local_tests/pipelines/test_pan2_pipeline_parity.py` | final latents and video exact | + +The pipeline test builds the FastVideo pipeline through the registry and component loaders and compares the final +latents (relative L2 < 2e-2, cosine > 0.9995) and the decoded video (PSNR > 40 dB) with `reference_generate`. To +confirm the threshold catches real deviations, three injected changes were each run against it: conditioning on +`hidden_states[-2]`, letting the latents fall to bf16 between steps, and shifting the first-frame crop by one pixel. +Each moved the latents by about 8e-2 relative L2. + +Upstream revision `e0c4c731` always masks video padding, fixing Ulysses=4 drift when padding was left unmasked. +FastVideo already excludes this padding by passing the original video length to `DistributedAttention`, which trims +it before RoPE and joint attention, then pads the output for the reverse all-to-all. The distributed test covers one +pad token (15 tokens at SP=2/4), two pad tokens (18 at SP=4, matching 704x1248x313 modulo 4), and an unpadded +control (18 at SP=2), with both zero T2V and first-frame I2V conditioning. No masking toggle is required. +The pipeline math and checkpoint layout are unchanged between upstream revisions `59dd8c9a` and `e0c4c731`. + +On 2026-10-20, all 39 PAN2 local tests passed on CPU (Torch 2.12.0, Transformers 5.15.0, Diffusers 0.39.0). +The broader API/contract suite had 499 passes and two failures also reproduced on clean `main` at `cdda1b753`: +the schema inventory lacks MiniMax-H3's `audio_scheduler_shift`/`video_scheduler_shift`, and +`test_load_kandinsky5.py` reads `FASTVIDEO_ATTENTION_BACKEND` outside the environment registry. +Importing the PAN2 tests and running their fixtures preserves both an unset and a pre-existing attention backend. + +## Not covered + +- Real-weight parity and an SSIM regression: `IFM/PAN2` is not public. +- GPU runs: everything above ran on CPU. The bf16 GPU path, FlashAttention, and SP over NCCL still need a GPU node. diff --git a/tests/local_tests/pan2/reference.py b/tests/local_tests/pan2/reference.py new file mode 100644 index 000000000..41a207580 --- /dev/null +++ b/tests/local_tests/pan2/reference.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 reference implementation for local parity tests. + +The diffusers ``PAN2Transformer3DModel`` and ``PAN2ModularPipeline`` are not public yet. The reference here is +assembled from diffusers' and transformers' own building blocks (``AdaLayerNormZero``, ``AdaLayerNormContinuous``, +``FeedForward`` with ``gelu-approximate``/``swiglu``, ``Timesteps``/``TimestepEmbedding``, +``get_1d_rotary_pos_embed``/``apply_rotary_emb``, ``Qwen3_5TextModel``, ``AutoencoderKLWan``, +``FlowMatchEulerDiscreteScheduler``) and follows the vLLM-Omni PAN2 port (vllm-project/vllm-omni#8602), which was +validated against the diffusers implementation (upstream revision ``e0c4c731``). Module names are the checkpoint's, +so checkpoints load strictly. Its unconditional SP padding mask is equivalent to unpadded reference attention. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np +import PIL.Image +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.attention import FeedForward +from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed +from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, RMSNorm + +TINY_REPO_ID = "wuqing157/tiny-pan2-modular-pipe" +TINY_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd" +SYSTEM_PROMPT = "You are a helpful assistant. Generate the video with the following instructions." +TEXT_ENCODER_LAYERS_TO_SKIP = 2 +MAX_SEQUENCE_LENGTH = 1536 + + +def tiny_model_dir() -> Path | None: + """``PAN2_TINY_MODEL_DIR``, else the pinned tiny checkpoint in the local HF cache.""" + local = os.getenv("PAN2_TINY_MODEL_DIR") + if local: + return Path(local) + try: + from huggingface_hub import snapshot_download + + return Path(snapshot_download(TINY_REPO_ID, revision=TINY_REVISION, local_files_only=True)) + except Exception: + return None + + +def _sdpa(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + out = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)) + return out.transpose(1, 2).flatten(2) + + +class _RefRefinerAttention(nn.Module): + + def __init__(self, dim: int, heads: int, dim_head: int, eps: float): + super().__init__() + self.heads = heads + self.to_q = nn.Linear(dim, heads * dim_head, bias=False) + self.to_k = nn.Linear(dim, heads * dim_head, bias=False) + self.to_v = nn.Linear(dim, heads * dim_head, bias=False) + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim, bias=False)]) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1))) + key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1))) + value = self.to_v(x).unflatten(-1, (self.heads, -1)) + return self.to_out[0](_sdpa(query, key, value)) + + +class _RefRefinerBlock(nn.Module): + + def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, eps: float): + super().__init__() + self.norm1 = RMSNorm(dim, eps=eps) + self.attn = _RefRefinerAttention(dim, heads, dim_head, eps) + self.norm2 = RMSNorm(dim, eps=eps) + self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="swiglu", bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.norm1(x)) + return x + self.ff(self.norm2(x)) + + +class _RefJointAttention(nn.Module): + + def __init__(self, dim: int, heads: int, dim_head: int, eps: float): + super().__init__() + self.heads = heads + self.to_q = nn.Linear(dim, heads * dim_head) + self.to_k = nn.Linear(dim, heads * dim_head) + self.to_v = nn.Linear(dim, heads * dim_head) + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + self.to_out = nn.ModuleList([nn.Linear(heads * dim_head, dim)]) + self.add_q_proj = nn.Linear(dim, heads * dim_head) + self.add_k_proj = nn.Linear(dim, heads * dim_head) + self.add_v_proj = nn.Linear(dim, heads * dim_head) + self.norm_added_q = RMSNorm(dim_head, eps=eps) + self.norm_added_k = RMSNorm(dim_head, eps=eps) + self.to_add_out = nn.Linear(heads * dim_head, dim) + + def forward(self, x: torch.Tensor, context: torch.Tensor, + rope: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + query = self.norm_q(self.to_q(x).unflatten(-1, (self.heads, -1))) + key = self.norm_k(self.to_k(x).unflatten(-1, (self.heads, -1))) + value = self.to_v(x).unflatten(-1, (self.heads, -1)) + # Interleaved (GPT-J) rotary embedding on the video tokens only. + query = apply_rotary_emb(query, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1) + key = apply_rotary_emb(key, rope, use_real=True, use_real_unbind_dim=-1, sequence_dim=1) + + context_query = self.norm_added_q(self.add_q_proj(context).unflatten(-1, (self.heads, -1))) + context_key = self.norm_added_k(self.add_k_proj(context).unflatten(-1, (self.heads, -1))) + context_value = self.add_v_proj(context).unflatten(-1, (self.heads, -1)) + + out = _sdpa(torch.cat([query, context_query], dim=1), torch.cat([key, context_key], dim=1), + torch.cat([value, context_value], dim=1)) + out, context_out = out.split([x.shape[1], context.shape[1]], dim=1) + return self.to_out[0](out), self.to_add_out(context_out) + + +class _RefBlock(nn.Module): + + def __init__(self, dim: int, heads: int, dim_head: int, mlp_ratio: float, context_mlp_ratio: float): + super().__init__() + eps = 1e-6 + self.norm1 = AdaLayerNormZero(dim) + self.norm1_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.scale_shift_table_context = nn.Parameter(torch.zeros(6 * dim)) + self.attn = _RefJointAttention(dim, heads, dim_head, eps) + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff = FeedForward(dim, inner_dim=int(dim * mlp_ratio), activation_fn="gelu-approximate") + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=eps) + self.ff_context = FeedForward(dim, inner_dim=int(dim * context_mlp_ratio), activation_fn="gelu-approximate") + + def forward(self, x, context, temb, rope): + norm_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(x, emb=temb) + c_shift_msa, c_scale_msa, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = ( + self.scale_shift_table_context.chunk(6, dim=-1)) + norm_context = self.norm1_context(context) * (1 + c_scale_msa) + c_shift_msa + attn_out, context_attn_out = self.attn(norm_x, norm_context, rope) + x = x + attn_out * gate_msa.unsqueeze(1) + context = context + context_attn_out * c_gate_msa + norm_x = self.norm2(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + norm_context = self.norm2_context(context) * (1 + c_scale_mlp) + c_shift_mlp + x = x + self.ff(norm_x) * gate_mlp.unsqueeze(1) + context = context + self.ff_context(norm_context) * c_gate_mlp + return x, context + + +class ReferencePAN2Transformer(nn.Module): + + def __init__(self, patch_size=(1, 2, 2), in_channels=97, out_channels=48, num_attention_heads=32, + attention_head_dim=128, num_layers=44, num_refiner_layers=4, mlp_ratio=8.0, context_mlp_ratio=4.0, + refiner_mlp_ratio=4.0, text_embed_dim=4096, rope_axes_dim=(16, 56, 56), rope_theta=256.0, + rope_temporal_scale=16 / 24, **_): + super().__init__() + self.patch_size = tuple(patch_size) + self.rope_axes_dim = tuple(rope_axes_dim) + self.rope_theta = rope_theta + self.rope_temporal_scale = rope_temporal_scale + inner_dim = num_attention_heads * attention_head_dim + self.x_embedder = nn.Conv3d(in_channels, inner_dim, kernel_size=self.patch_size, stride=self.patch_size) + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) + self.time_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=inner_dim) + self.context_embedder = nn.Sequential(RMSNorm(text_embed_dim, eps=1e-5), nn.Linear(text_embed_dim, inner_dim)) + self.context_refiner = nn.ModuleList([ + _RefRefinerBlock(inner_dim, num_attention_heads, attention_head_dim, refiner_mlp_ratio, eps=1e-5) + for _ in range(num_refiner_layers) + ]) + self.context_type_embedding = nn.Parameter(torch.zeros(inner_dim)) + self.transformer_blocks = nn.ModuleList([ + _RefBlock(inner_dim, num_attention_heads, attention_head_dim, mlp_ratio, context_mlp_ratio) + for _ in range(num_layers) + ]) + self.norm_out = AdaLayerNormContinuous(inner_dim, inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out = nn.Linear(inner_dim, self.patch_size[0] * self.patch_size[1] * self.patch_size[2] * out_channels) + + def _rope(self, num_frames: int, height: int, width: int, device) -> tuple[torch.Tensor, torch.Tensor]: + axes = [torch.arange(size, device=device, dtype=torch.float32) for size in (num_frames, height, width)] + grid = torch.stack(torch.meshgrid(*axes, indexing="ij"), dim=0) + grid[0] = grid[0] * self.rope_temporal_scale + cos, sin = zip(*[ + get_1d_rotary_pos_embed(dim, grid[i].reshape(-1), self.rope_theta, use_real=True, + repeat_interleave_real=True) for i, dim in enumerate(self.rope_axes_dim) + ]) + return torch.cat(cos, dim=1).float(), torch.cat(sin, dim=1).float() + + def forward(self, hidden_states, timestep, encoder_hidden_states): + batch_size, _, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + t, h, w = num_frames // p_t, height // p_h, width // p_w + rope = self._rope(t, h, w, hidden_states.device) + x = self.x_embedder(hidden_states).flatten(2).transpose(1, 2) + temb = self.time_embedder(self.time_proj(timestep).to(x.dtype)) + context = self.context_embedder(encoder_hidden_states) + for block in self.context_refiner: + context = block(context) + context = context + self.context_type_embedding + for block in self.transformer_blocks: + x, context = block(x, context, temb, rope) + x = self.proj_out(self.norm_out(x, temb)) + x = x.reshape(batch_size, t, h, w, -1, p_t, p_h, p_w).permute(0, 4, 1, 5, 2, 6, 3, 7) + return x.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +def reference_encode_prompt(tokenizer, text_encoder, prompt: str, device) -> torch.Tensor: + messages = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": prompt or " "}] + input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True, + return_dict=True).input_ids + header_ids = tokenizer("<|im_start|>user\n", add_special_tokens=False).input_ids + crop_start = next(i + len(header_ids) for i in range(len(input_ids)) + if input_ids[i:i + len(header_ids)] == header_ids) + input_ids = torch.tensor([input_ids[:crop_start + MAX_SEQUENCE_LENGTH]], device=device) + hidden_states = text_encoder(input_ids=input_ids, output_hidden_states=True).hidden_states + return hidden_states[-(TEXT_ENCODER_LAYERS_TO_SKIP + 1)][:, crop_start:] + + +def reference_condition_latents(vae, latents: torch.Tensor, image: PIL.Image.Image | None, height: int, + width: int) -> torch.Tensor: + batch_size, num_channels, num_frames, latent_height, latent_width = latents.shape + condition = latents.new_zeros(batch_size, num_channels + 1, num_frames, latent_height, latent_width) + if image is None: + return condition + from diffusers.video_processor import VideoProcessor + + image = image.convert("RGB") + scale = max(width / image.width, height / image.height) + image = image.resize((round(image.width * scale), round(image.height * scale)), resample=PIL.Image.LANCZOS) + left, top = round((image.width - width) / 2), round((image.height - height) / 2) + image = image.crop((left, top, left + width, top + height)) + pixels = VideoProcessor(vae_scale_factor=16).preprocess(image, height=height, width=width) + pixels = pixels.unsqueeze(2).to(device=latents.device, dtype=vae.dtype) + image_latents = vae.encode(pixels).latent_dist.mode() + mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(image_latents) + std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(image_latents) + condition[:, :num_channels, :1] = (image_latents - mean) / std + condition[:, num_channels:, :1] = 1.0 + return condition + + +@torch.no_grad() +def reference_generate(components: dict, *, prompt: str, negative_prompt: str, image: PIL.Image.Image | None, + height: int, width: int, num_frames: int, num_inference_steps: int, guidance_scale: float, + seed: int, dit_dtype: torch.dtype, device="cpu", + skip_first_latent_cfg: bool = True) -> tuple[torch.Tensor, torch.Tensor]: + """Final latents and the decoded video in ``[0, 1]``, following the vLLM-Omni pipeline step by step. + + ``skip_first_latent_cfg`` is the reference sampler's ``--i2v_skip_first_latent_cfg``: with an image, the first + latent frame takes the positive prediction instead of the guided one. + """ + from diffusers.utils.torch_utils import randn_tensor + + tokenizer, text_encoder = components["tokenizer"], components["text_encoder"] + transformer, vae, scheduler = components["transformer"], components["vae"], components["scheduler"] + do_cfg = guidance_scale > 1.0 + prompt_embeds = reference_encode_prompt(tokenizer, text_encoder, prompt, device).to(dit_dtype) + negative_embeds = (reference_encode_prompt(tokenizer, text_encoder, negative_prompt, device).to(dit_dtype) + if do_cfg else None) + + shape = (1, vae.config.z_dim, (num_frames - 1) // vae.config.scale_factor_temporal + 1, + height // vae.config.scale_factor_spatial, width // vae.config.scale_factor_spatial) + generator = [torch.Generator("cpu").manual_seed(seed)] + latents = randn_tensor(shape, generator=generator, device=torch.device(device), dtype=torch.float32) + condition = reference_condition_latents(vae, latents, image, height, width).to(dit_dtype) + + scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, num_inference_steps + 1)[:-1], device=device) + for t in scheduler.timesteps: + model_input = torch.cat([latents.to(dit_dtype), condition], dim=1) + timestep = t.expand(model_input.shape[0]) + noise_pred = transformer(model_input, timestep, prompt_embeds) + if do_cfg: + noise_uncond = transformer(model_input, timestep, negative_embeds) + guided = noise_uncond + guidance_scale * (noise_pred - noise_uncond) + if image is not None and skip_first_latent_cfg: + guided[:, :, :1] = noise_pred[:, :, :1] + noise_pred = guided + latents = scheduler.step(noise_pred.float(), t, latents, return_dict=False)[0] + + mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents) + std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents) + video = vae.decode((latents * std + mean).to(vae.dtype), return_dict=False)[0] + return latents, (video.float() / 2 + 0.5).clamp(0, 1) + + +def load_reference_components(model_dir: Path, dit_dtype: torch.dtype) -> dict: + import json + + from diffusers import AutoencoderKLWan, FlowMatchEulerDiscreteScheduler + from safetensors.torch import load_file + from transformers import AutoTokenizer, Qwen3_5TextModel + + transformer_dir = model_dir / "transformer" + transformer = ReferencePAN2Transformer(**json.loads((transformer_dir / "config.json").read_text())) + transformer.load_state_dict(load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors")), strict=True) + return { + "tokenizer": AutoTokenizer.from_pretrained(model_dir / "tokenizer"), + "text_encoder": Qwen3_5TextModel.from_pretrained(model_dir / "text_encoder", torch_dtype=torch.float32).eval(), + "transformer": transformer.to(dit_dtype).eval(), + "vae": AutoencoderKLWan.from_pretrained(model_dir / "vae", torch_dtype=torch.float32).eval(), + "scheduler": FlowMatchEulerDiscreteScheduler.from_pretrained(model_dir / "scheduler"), + } diff --git a/tests/local_tests/pipelines/test_pan2_pipeline_parity.py b/tests/local_tests/pipelines/test_pan2_pipeline_parity.py new file mode 100644 index 000000000..daf608fb2 --- /dev/null +++ b/tests/local_tests/pipelines/test_pan2_pipeline_parity.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 end-to-end pipeline parity on the public tiny checkpoint. + +The FastVideo pipeline (built through the real registry and component loaders) is compared against +``tests/local_tests/pan2/reference.py::reference_generate``, which follows the vLLM-Omni ``PAN2Pipeline.forward`` +step by step with transformers/diffusers components: prompt template and crop, fp32 noise, first-frame +conditioning, the uniform sigma schedule with shift 7, sequential CFG, fp32 Euler steps and the VAE decode. + +Both sides run the DiT in bf16 (FastVideo's denoiser always feeds it bf16), so the comparison allows bf16 rounding. +Runs on CPU when no GPU is present. Set ``PAN2_TINY_MODEL_DIR`` to a local copy of the tiny checkpoint. +""" + +from __future__ import annotations + +import socket +from contextlib import ExitStack + +import PIL.Image +import pytest +import torch + +from fastvideo import envs +from tests.local_tests.pan2.reference import load_reference_components, reference_generate, tiny_model_dir + +HEIGHT, WIDTH, NUM_FRAMES = 64, 96, 9 +NUM_INFERENCE_STEPS = 3 +# A large scale makes the guided steps visible through the tiny random-weight VAE. +GUIDANCE_SCALE = 30.0 +SEED = 42 +PROMPT = "a cat walks on the grass" +NEGATIVE_PROMPT = "blurry" +# Not the output aspect ratio, so the cover-resize and center crop are exercised. +CONDITION_IMAGE = PIL.Image.radial_gradient("L").resize((120, 64)).convert("RGB") + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +@pytest.fixture(scope="module") +def pan2_runtime(): + model_dir = tiny_model_dir() + if model_dir is None or not (model_dir / "modular_model_index.json").is_file(): + pytest.skip("wuqing157/tiny-pan2-modular-pipe is not available locally; set PAN2_TINY_MODEL_DIR") + from fastvideo.configs.pipelines.pan2 import PAN2PipelineConfig + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + import fastvideo.distributed.parallel_state as parallel_state + from fastvideo.fastvideo_args import FastVideoArgs + from fastvideo.pipelines import build_pipeline + + with ExitStack() as stack: + stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + monkeypatch = stack.enter_context(pytest.MonkeyPatch.context()) + if not torch.cuda.is_available(): + # Reach modules that imported the device helper by name on a GPU-less Linux host. + def _cpu_device() -> torch.device: + return torch.device("cpu") + + monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__) + for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"), + ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")): + stack.enter_context(envs.override_external(name, value)) + stack.callback(cleanup_dist_env_and_memory) + maybe_init_distributed_environment_and_model_parallel(1, 1) + + pipeline_config = PAN2PipelineConfig(dit_precision="bf16", + vae_precision="fp32", + text_encoder_precisions=("fp32", )) + args = FastVideoArgs(model_path=str(model_dir), + num_gpus=1, + pipeline_config=pipeline_config, + dit_cpu_offload=False, + dit_layerwise_offload=False, + text_encoder_cpu_offload=False, + vae_cpu_offload=False, + use_fsdp_inference=False) + pipeline = build_pipeline(args) + device = parallel_state.get_local_torch_device() + yield pipeline, args, model_dir, device + + +def _fastvideo_generate(pipeline, args, image, guidance_scale: float = GUIDANCE_SCALE): + from fastvideo.pipelines import ForwardBatch + + batch = ForwardBatch(data_type="video", + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + seed=SEED, + height=HEIGHT, + width=WIDTH, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=guidance_scale, + pil_image=image) + with torch.inference_mode(): + batch = pipeline.forward(batch, args) + return batch.latents.float().cpu(), batch.output.float().cpu() + + +def _psnr(a: torch.Tensor, b: torch.Tensor) -> float: + mse = (a - b).pow(2).mean().clamp_min(1e-12) + return float(10 * torch.log10(1.0 / mse)) + + +@pytest.mark.parametrize( + ("image", "skip_first_latent_cfg"), + [(None, True), (CONDITION_IMAGE, True), (CONDITION_IMAGE, False)], + ids=["t2v", "i2v", "i2v-guided-first-frame"], +) +def test_pan2_pipeline_matches_reference(pan2_runtime, monkeypatch, image, skip_first_latent_cfg): + pipeline, args, model_dir, device = pan2_runtime + monkeypatch.setattr(args.pipeline_config, "i2v_skip_first_latent_cfg", skip_first_latent_cfg) + latents, video = _fastvideo_generate(pipeline, args, image) + + components = load_reference_components(model_dir, torch.bfloat16) + components = { + name: module.to(device) if isinstance(module, torch.nn.Module) else module + for name, module in components.items() + } + ref_latents, ref_video = reference_generate(components, + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + image=image, + height=HEIGHT, + width=WIDTH, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=GUIDANCE_SCALE, + seed=SEED, + dit_dtype=torch.bfloat16, + device=device, + skip_first_latent_cfg=skip_first_latent_cfg) + ref_latents, ref_video = ref_latents.float().cpu(), ref_video.float().cpu() + + assert latents.dtype == torch.float32 + assert latents.shape == ref_latents.shape == (1, 48, 3, HEIGHT // 16, WIDTH // 16) + assert video.shape == ref_video.shape == (1, 3, NUM_FRAMES, HEIGHT, WIDTH) + latent_rel = float((latents - ref_latents).norm() / ref_latents.norm()) + cosine = float(torch.nn.functional.cosine_similarity(latents.flatten(), ref_latents.flatten(), dim=0)) + psnr = _psnr(video, ref_video) + print(f"latents rel_l2={latent_rel:.3e} cosine={cosine:.6f}; video psnr={psnr:.1f} dB") + assert latent_rel < 2e-2 + assert cosine > 0.9995 + assert psnr > 40.0 + + +def test_pan2_image_conditioning_and_guidance_change_every_frame(pan2_runtime): + pipeline, args, _, _ = pan2_runtime + _, t2v = _fastvideo_generate(pipeline, args, None) + _, i2v = _fastvideo_generate(pipeline, args, CONDITION_IMAGE) + _, unguided = _fastvideo_generate(pipeline, args, None, guidance_scale=1.0) + + def per_frame_max_diff(a, b): + return (a - b).abs().amax(dim=(0, 1, 3, 4)) + + assert per_frame_max_diff(t2v, i2v).max() > 1e-3 + assert per_frame_max_diff(t2v, unguided).min() > 1e-3 diff --git a/tests/local_tests/schedulers/test_pan2_scheduler_parity.py b/tests/local_tests/schedulers/test_pan2_scheduler_parity.py new file mode 100644 index 000000000..3fcdb7497 --- /dev/null +++ b/tests/local_tests/schedulers/test_pan2_scheduler_parity.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 flow-matching Euler schedule on FastVideo's FlowMatchEulerDiscreteScheduler. + +Coverage scope: both. PAN2 loads Diffusers' ``FlowMatchEulerDiscreteScheduler`` +from ``/scheduler`` (shift 7.0, no dynamic shifting), optionally +overrides the shift with ``flow_shift``, calls +``set_timesteps(sigmas=np.linspace(1.0, 0.0, N + 1)[:-1])`` and steps float32 +latents. The FastVideo side is loaded by the production ``SchedulerLoader`` and +its schedule is built by the stock ``TimestepPreparationStage`` from +``batch.sigmas``. Timesteps, sigmas, and every Euler update must be bit-exact. + +The default asset is the tiny PAN2 checkpoint (config only), so this runs on a +CPU-only host:: + + pytest tests/local_tests/schedulers/test_pan2_scheduler_parity.py -v -s + +Set ``PAN2_MODEL_PATH`` to a local PAN2 root containing ``scheduler/`` to use +another checkpoint. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from torch.testing import assert_close + +try: + from diffusers import FlowMatchEulerDiscreteScheduler as DiffusersFlowMatchEulerDiscreteScheduler +except ImportError: # pragma: no cover - environment dependent + pytest.skip("diffusers with FlowMatchEulerDiscreteScheduler is required", allow_module_level=True) + +from fastvideo.models.loader.component_loader import SchedulerLoader +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler, ) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage + +PARITY_SCOPE = "both" +TINY_PAN2_REPO = "wuqing157/tiny-pan2-modular-pipe" +TINY_PAN2_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd" +MODEL_PATH_ENV = "PAN2_MODEL_PATH" +LATENT_SHAPE = (1, 48, 3, 4, 6) + + +def _pan2_root(subfolder: str) -> Path: + """Resolve a PAN2 root that contains ``subfolder``; skip when unavailable.""" + override = os.getenv(MODEL_PATH_ENV) + if override: + root = Path(override) + else: + try: + from huggingface_hub import snapshot_download + except ImportError: + pytest.skip("huggingface_hub is required to resolve the tiny PAN2 checkpoint") + root = None + last_error: Exception | None = None + for local_files_only in (True, False): + try: + root = Path( + snapshot_download( + TINY_PAN2_REPO, + revision=TINY_PAN2_REVISION, + allow_patterns=[f"{subfolder}/*"], + local_files_only=local_files_only, + )) + break + except Exception as exc: # noqa: BLE001 - offline/missing cache must skip. + last_error = exc + if root is None: + pytest.skip(f"Tiny PAN2 checkpoint unavailable: {last_error}") + if not (root / subfolder).is_dir(): + pytest.skip(f"PAN2 {subfolder}/ missing under {root}") + return root + + +@pytest.fixture(scope="module") +def pan2_root() -> Path: + root = _pan2_root("scheduler") + if not (root / "scheduler" / "scheduler_config.json").is_file(): + pytest.skip(f"PAN2 scheduler config missing under {root}") + return root + + +@pytest.fixture +def cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None: + """Pin FastVideo's device helper to CPU; it reports ``mps`` on CPU-only Linux.""" + monkeypatch.setattr("fastvideo.pipelines.stages.timestep_preparation.get_local_torch_device", + lambda: torch.device("cpu")) + + +def _pan2_sigmas(num_steps: int) -> np.ndarray: + """PAN2's pre-shift grid: uniform from 1 down to, but excluding, 0.""" + return np.linspace(1.0, 0.0, num_steps + 1)[:-1] + + +def _reference(root: Path, flow_shift: float | None = None) -> DiffusersFlowMatchEulerDiscreteScheduler: + """Mirror the upstream PAN2 pipeline's scheduler construction.""" + scheduler = DiffusersFlowMatchEulerDiscreteScheduler.from_pretrained(str(root), + subfolder="scheduler", + local_files_only=True) + if flow_shift is not None: + scheduler._shift = flow_shift + return scheduler + + +def _fastvideo(root: Path, flow_shift: float | None = None) -> FastVideoFlowMatchEulerDiscreteScheduler: + """Load through the production SchedulerLoader, which applies ``flow_shift``.""" + fastvideo_args = SimpleNamespace(pipeline_config=SimpleNamespace(flow_shift=flow_shift)) + scheduler = SchedulerLoader().load(str(root / "scheduler"), fastvideo_args) + assert type(scheduler) is FastVideoFlowMatchEulerDiscreteScheduler + return scheduler + + +def _prepare_fastvideo_timesteps(scheduler: FastVideoFlowMatchEulerDiscreteScheduler, + num_steps: int) -> ForwardBatch: + """Drive the stock TimestepPreparationStage with PAN2's custom sigmas.""" + batch = ForwardBatch(data_type="video", num_inference_steps=num_steps, sigmas=_pan2_sigmas(num_steps).tolist()) + stage = TimestepPreparationStage(scheduler) + assert stage.verify_input(batch, SimpleNamespace()).is_valid() + batch = stage.forward(batch, SimpleNamespace()) + assert batch.timesteps is scheduler.timesteps + return batch + + +def test_pan2_scheduler_loader_keeps_checkpoint_config(pan2_root: Path) -> None: + """Without ``flow_shift`` the loader keeps every released config value, including shift 7.0.""" + raw = json.loads((pan2_root / "scheduler" / "scheduler_config.json").read_text()) + assert raw["_class_name"] == "FlowMatchEulerDiscreteScheduler" + scheduler = _fastvideo(pan2_root) + for key, value in raw.items(): + if not key.startswith("_"): + assert scheduler.config[key] == value, key + assert scheduler.shift == raw["shift"] == 7.0 + assert not scheduler.config.use_dynamic_shifting + + +@pytest.mark.parametrize("num_steps", [1, 4, 8, 30, 50]) +def test_pan2_scheduler_schedule_matches_diffusers(num_steps: int, pan2_root: Path, cpu_runtime: None) -> None: + """Timesteps and sigmas from ``batch.sigmas`` are bit-exact for short and long schedules.""" + reference = _reference(pan2_root) + reference.set_timesteps(sigmas=_pan2_sigmas(num_steps), device="cpu") + fastvideo = _fastvideo(pan2_root) + batch = _prepare_fastvideo_timesteps(fastvideo, num_steps) + + print(f"[pan2-scheduler] N={num_steps} timesteps[:4]={reference.timesteps[:4].tolist()} " + f"sigmas[-3:]={reference.sigmas[-3:].tolist()}") + assert len(batch.timesteps) == num_steps + assert_close(batch.timesteps, reference.timesteps, atol=0, rtol=0) + assert_close(fastvideo.sigmas, reference.sigmas, atol=0, rtol=0) + assert fastvideo.sigmas[0].item() == 1.0 and fastvideo.sigmas[-1].item() == 0.0 + + +@pytest.mark.parametrize("flow_shift", [None, 5.0], ids=["checkpoint_shift", "flow_shift_override"]) +def test_pan2_scheduler_steps_match_diffusers(flow_shift: float | None, pan2_root: Path, cpu_runtime: None) -> None: + """Five float32 Euler updates are bit-exact, with and without a shift override.""" + num_steps = 5 + reference = _reference(pan2_root, flow_shift) + reference.set_timesteps(sigmas=_pan2_sigmas(num_steps), device="cpu") + fastvideo = _fastvideo(pan2_root, flow_shift) + _prepare_fastvideo_timesteps(fastvideo, num_steps) + assert fastvideo.shift == reference.shift == (flow_shift or 7.0) + assert_close(fastvideo.timesteps, reference.timesteps, atol=0, rtol=0) + + generator = torch.Generator().manual_seed(2026) + reference_latents = torch.randn(LATENT_SHAPE, generator=generator) + fastvideo_latents = reference_latents.clone() + max_diff = 0.0 + for index in range(num_steps): + noise_pred = torch.randn(LATENT_SHAPE, generator=generator) + reference_latents = reference.step(noise_pred, reference.timesteps[index], reference_latents).prev_sample + fastvideo_latents = fastvideo.step(noise_pred, fastvideo.timesteps[index], fastvideo_latents).prev_sample + assert fastvideo_latents.dtype == torch.float32 + max_diff = max(max_diff, (fastvideo_latents - reference_latents).abs().max().item()) + assert_close(fastvideo_latents, reference_latents, atol=0, rtol=0) + print(f"[pan2-scheduler] flow_shift={flow_shift} steps={num_steps} max_abs={max_diff:.3e}") + assert fastvideo.step_index == reference.step_index == num_steps + + +def test_pan2_scheduler_step_returns_model_output_dtype(pan2_root: Path, cpu_runtime: None) -> None: + """Both implementations cast ``prev_sample`` to the model output dtype. + + PAN2 keeps float32 latents by stepping with ``noise_pred.float()``; a bf16 + noise prediction would silently drop the latents to bf16 on either side. + """ + reference = _reference(pan2_root) + reference.set_timesteps(sigmas=_pan2_sigmas(2), device="cpu") + fastvideo = _fastvideo(pan2_root) + _prepare_fastvideo_timesteps(fastvideo, 2) + latents = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(5)) + noise_pred = torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(6)).to(torch.bfloat16) + expected = reference.step(noise_pred, reference.timesteps[0], latents).prev_sample + actual = fastvideo.step(noise_pred, fastvideo.timesteps[0], latents).prev_sample + assert expected.dtype == actual.dtype == torch.bfloat16 + assert_close(actual, expected, atol=0, rtol=0) diff --git a/tests/local_tests/test_cookbook_serving.py b/tests/local_tests/test_cookbook_serving.py index f0a2e11e8..2f4e37822 100644 --- a/tests/local_tests/test_cookbook_serving.py +++ b/tests/local_tests/test_cookbook_serving.py @@ -264,3 +264,23 @@ def test_kandinsky6_cookbook_recipes_match_maintained_examples(): assert all("INPUT_VIDEO" in by_id[recipe_id]["command"] for recipe_id in ("kandinsky6-vsr", "kandinsky6-vsr-distilled")) assert "VSR-distilled2steps" in by_id["kandinsky6-vsr-distilled"]["command"] validate_cookbook() + + +def test_pan2_recipe_serves_text_and_image_from_one_deployment(): + recipes = json.loads(COOKBOOK_DATA.read_text())["recipes"] + recipe = next(item for item in recipes if item["family"] == "pan2") + assert recipe["id"] == "pan2-ti2v" + assert recipe["model"] == "IFM/PAN2" + assert recipe["evidence"] == "Source-backed" + assert recipe["hardware"] == {"platform": "cuda", "gpu_count": 1, "evidence": "source-configured"} + profile = cookbook_serving_profile(recipe) + assert profile["command"] == "fastvideo serve --config examples/serving/openai_pan2.yaml --server.host 127.0.0.1" + assert profile["model"] == "pan2" + assert profile["playground_url"] is None + assert profile["sampling"]["num_frames"] == 313 + assert any("image_reference" in item for item in recipe["serving"]["limitations"]) + manifest = yaml.safe_load((ROOT / "docs/cookbook/recipes/pan2.yaml").read_text()) + deployment = manifest["deployments"]["cuda-rest"] + assert deployment["workload"] == "t2v" + assert deployment["defaults"] == recipe["serving"]["source"] + validate_cookbook() diff --git a/tests/local_tests/transformers/test_pan2_transformer_parity.py b/tests/local_tests/transformers/test_pan2_transformer_parity.py new file mode 100644 index 000000000..4df532f06 --- /dev/null +++ b/tests/local_tests/transformers/test_pan2_transformer_parity.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 transformer weight-schema and forward parity. + +The reference is ``tests/local_tests/pan2/reference.py``: diffusers building blocks arranged as the (not yet public) +diffusers ``PAN2Transformer3DModel``, following the vLLM-Omni port. Its module names are the checkpoint's, so the +public tiny checkpoint loads into it strictly. + +Set ``PAN2_TINY_MODEL_DIR`` to a local copy of ``wuqing157/tiny-pan2-modular-pipe``; otherwise the HF cache is +used and the checkpoint test skips when it is absent. +""" + +from __future__ import annotations + +import argparse +import json +import os +import socket +import subprocess +import sys +from contextlib import ExitStack +from dataclasses import fields +from pathlib import Path + +import pytest +import torch +import torch.nn as nn +from torch.testing import assert_close + +from fastvideo import envs +from fastvideo.forward_context import set_forward_context +from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict +from fastvideo.models.loader.utils import get_param_names_mapping +from tests.local_tests.pan2.reference import TINY_REPO_ID, ReferencePAN2Transformer, tiny_model_dir + +REPO_ROOT = Path(__file__).resolve().parents[3] +SP_WORLD_SIZES = (2, 4) + + +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _small_kwargs() -> dict: + """Larger than the tiny checkpoint: several heads, asymmetric RoPE axes, more blocks, distinct MLP ratios.""" + return { + "patch_size": (1, 2, 2), + "in_channels": 9, + "out_channels": 4, + "num_attention_heads": 4, + "attention_head_dim": 32, + "num_layers": 3, + "num_refiner_layers": 2, + "mlp_ratio": 3.0, + "context_mlp_ratio": 2.0, + "refiner_mlp_ratio": 1.5, + "text_embed_dim": 48, + "rope_axes_dim": (8, 12, 12), + "rope_theta": 256.0, + "rope_temporal_scale": 16 / 24, + } + + +@pytest.fixture +def single_rank_dist(monkeypatch): + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + import fastvideo.distributed.parallel_state as parallel_state + + if not torch.cuda.is_available(): + def _cpu_device() -> torch.device: + return torch.device("cpu") + + monkeypatch.setattr(parallel_state.get_local_torch_device, "__code__", _cpu_device.__code__) + with ExitStack() as stack: + stack.enter_context(envs.FASTVIDEO_ATTENTION_BACKEND.override("TORCH_SDPA")) + for name, value in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", str(_free_port())), ("RANK", "0"), + ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")): + stack.enter_context(envs.override_external(name, value)) + try: + maybe_init_distributed_environment_and_model_parallel(1, 1) + yield + finally: + cleanup_dist_env_and_memory() + + +def _build_native(arch_kwargs: dict): + from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig + from fastvideo.models.dits.pan2 import PAN2Transformer3DModel + + config = PAN2VideoConfig() + config.update_model_arch(arch_kwargs) + return PAN2Transformer3DModel(config=config, hf_config=arch_kwargs).float().eval() + + +def _load_into_native(native, weights) -> None: + incompatible = load_model_from_full_model_state_dict( + native, + iter(weights), + device=torch.device("cpu"), + param_dtype=torch.float32, + strict=True, + param_names_mapping=get_param_names_mapping(native.param_names_mapping), + training_mode=False, + ) + assert incompatible.missing_keys == [] + assert incompatible.unexpected_keys == [] + + +def _randomize(module: nn.Module, seed: int) -> None: + generator = torch.Generator().manual_seed(seed) + with torch.no_grad(): + for param in module.parameters(): + param.copy_(torch.randn(param.shape, generator=generator) * 0.2) + + +def _make_inputs(in_channels: int, text_dim: int, batch_size: int = 2, seed: int = 0) -> dict[str, torch.Tensor]: + generator = torch.Generator().manual_seed(seed) + return { + # 9 frames at 64x96 in latent space: 3 latent frames of 4x6. + "hidden_states": torch.randn(batch_size, in_channels, 3, 4, 6, generator=generator), + "timestep": torch.tensor([700.0, 125.5][:batch_size]), + "encoder_hidden_states": torch.randn(batch_size, 13, text_dim, generator=generator), + } + + +def _compare(native, reference, inputs: dict[str, torch.Tensor], scope: str) -> None: + with torch.inference_mode(): + expected = reference(**inputs) + with set_forward_context(current_timestep=0, attn_metadata=None): + actual = native(**inputs) + assert actual.shape == expected.shape + diff = (actual - expected).abs() + print(f"{scope}: shape={tuple(actual.shape)} ref_abs_mean={expected.abs().mean().item():.3e} " + f"diff_max={diff.max().item():.3e} diff_mean={diff.mean().item():.3e}") + assert torch.isfinite(expected).all() and expected.abs().mean() > 1e-3 + assert_close(actual, expected, atol=2e-5, rtol=2e-5) + + +def test_pan2_transformer_registry_resolves_native_class(): + from fastvideo.models.dits.pan2 import PAN2Transformer3DModel + from fastvideo.models.registry import ModelRegistry + + model_cls, architecture = ModelRegistry.resolve_model_cls("PAN2Transformer3DModel") + assert model_cls is PAN2Transformer3DModel + assert architecture == "PAN2Transformer3DModel" + + +@pytest.mark.parametrize(("module_name", "expected"), [ + ("transformer_blocks.0", True), + ("transformer_blocks.43", True), + ("context_refiner.3", True), + ("transformer_blocks.0.attn.to_qkv", False), + ("transformer_blocks.0.ff.fc_in", False), +]) +def test_pan2_fsdp_shard_conditions_match_only_blocks(module_name: str, expected: bool): + from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig + + conditions = PAN2ArchConfig()._fsdp_shard_conditions + assert any(condition(module_name, None) for condition in conditions) is expected + + +def test_pan2_default_arch_config_is_the_29b_model(): + from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig + + config = PAN2ArchConfig() + assert config.hidden_size == 4096 + assert config.num_channels_latents == 48 + assert config.in_channels == 2 * 48 + 1 + assert sum(config.rope_axes_dim) == config.attention_head_dim + + +def test_pan2_tiny_checkpoint_config_fields_are_arch_fields(): + model_dir = tiny_model_dir() + if model_dir is None or not (model_dir / "transformer" / "config.json").is_file(): + pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR") + from fastvideo.configs.models.dits.pan2 import PAN2ArchConfig + + checkpoint_config = json.loads((model_dir / "transformer" / "config.json").read_text()) + assert checkpoint_config.pop("_class_name") == "PAN2Transformer3DModel" + checkpoint_config.pop("_diffusers_version", None) + assert set(checkpoint_config) <= {field.name for field in fields(PAN2ArchConfig)} + + +def test_pan2_tiny_checkpoint_strict_load_and_forward_parity(single_rank_dist): + model_dir = tiny_model_dir() + if model_dir is None or not (model_dir / "transformer").is_dir(): + pytest.skip(f"{TINY_REPO_ID} is not available locally; set PAN2_TINY_MODEL_DIR") + from safetensors.torch import load_file + + transformer_dir = model_dir / "transformer" + arch_kwargs = json.loads((transformer_dir / "config.json").read_text()) + state_dict = load_file(str(transformer_dir / "diffusion_pytorch_model.safetensors")) + + reference = ReferencePAN2Transformer(**arch_kwargs).float().eval() + reference.load_state_dict(state_dict, strict=True) + native = _build_native(arch_kwargs) + _load_into_native(native, state_dict.items()) + + inputs = _make_inputs(arch_kwargs["in_channels"], arch_kwargs["text_embed_dim"]) + _compare(native, reference, inputs, "tiny checkpoint") + + +def test_pan2_layerwise_offload_hooks_the_main_blocks(single_rank_dist): + native = _build_native(_small_kwargs()) + + first_module_list = next(child for child in native.children() if isinstance(child, torch.nn.ModuleList)) + + assert first_module_list is native.transformer_blocks + + +@pytest.mark.parametrize("batch_size", [1, 2]) +def test_pan2_random_weights_forward_parity(single_rank_dist, batch_size: int): + kwargs = _small_kwargs() + reference = ReferencePAN2Transformer(**kwargs).float().eval() + # Zero-initialized tables (text modulation, type embedding) must be exercised too. + _randomize(reference, seed=1234) + native = _build_native(kwargs) + _load_into_native(native, reference.state_dict().items()) + + inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"], batch_size=batch_size, seed=batch_size) + _compare(native, reference, inputs, f"random weights, batch {batch_size}") + + +def _run_sp_worker(sp_size: int, output_path: Path) -> None: + from fastvideo.configs.models.dits.pan2 import PAN2VideoConfig + from fastvideo.distributed import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel + from fastvideo.models.dits.pan2 import PAN2Transformer3DModel + import torch.distributed as dist + + if torch.cuda.is_available(): + device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}") + torch.cuda.set_device(device) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + else: + # GPU-less host: gloo ranks on CPU. fastvideo resolves the local device to mps there; swapping the code + # object reaches every module that imported get_local_torch_device by name. + import fastvideo.distributed.parallel_state as parallel_state + + def _cpu_device() -> torch.device: + return torch.device("cpu") + + parallel_state.get_local_torch_device.__code__ = _cpu_device.__code__ + device = torch.device("cpu") + try: + maybe_init_distributed_environment_and_model_parallel(1, sp_size) + kwargs = _small_kwargs() + reference = ReferencePAN2Transformer(**kwargs).float() + _randomize(reference, seed=1234) + config = PAN2VideoConfig() + config.update_model_arch(kwargs) + native = PAN2Transformer3DModel(config=config, hf_config=kwargs).float().eval() + _load_into_native(native, reference.state_dict().items()) + native = native.to(device) + outputs = {} + # 15 tokens need one pad at SP=2/4; 18 need two at SP=4, as at 704x1248x313. + for latent_frames in (5, 6): + for image_conditioned in (False, True): + inputs = _make_inputs(kwargs["in_channels"], kwargs["text_embed_dim"]) + video = torch.randn(2, kwargs["in_channels"], latent_frames, 2, 6, + generator=torch.Generator().manual_seed(5)) + video[:, kwargs["out_channels"]:] = 0 + if image_conditioned: + video[:, kwargs["out_channels"]:-1, :1] = torch.randn( + 2, kwargs["out_channels"], 1, 2, 6, generator=torch.Generator().manual_seed(6)) + video[:, -1:, :1] = 1 + inputs["hidden_states"] = video + inputs = {name: value.to(device) for name, value in inputs.items()} + with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None): + output = native(**inputs) + outputs[f"{latent_frames}_{image_conditioned}"] = output.detach().cpu() + if int(os.environ.get("RANK", "0")) == 0: + torch.save(outputs, output_path) + dist.barrier() + finally: + cleanup_dist_env_and_memory() + + +def _run_torchrun(nproc_per_node: int, output_path: Path) -> None: + command = [ + sys.executable, "-m", "torch.distributed.run", "--nnodes", "1", "--nproc_per_node", + str(nproc_per_node), "--master_port", + str(_free_port()), + str(Path(__file__).resolve()), "--pan2-sp-worker", "--sp-size", str(nproc_per_node), "--output", + str(output_path) + ] + environment = os.environ.copy() + environment["PYTHONPATH"] = os.pathsep.join(filter(None, [str(REPO_ROOT), environment.get("PYTHONPATH")])) + environment["FASTVIDEO_ATTENTION_BACKEND"] = "TORCH_SDPA" + process = subprocess.run(command, capture_output=True, text=True, env=environment, timeout=180) + if process.returncode != 0: + raise RuntimeError(f"SP={nproc_per_node} worker failed with code {process.returncode}\n" + f"STDOUT:\n{process.stdout}\nSTDERR:\n{process.stderr}") + + +@pytest.fixture(scope="module") +def single_rank_sp_outputs(tmp_path_factory): + path = tmp_path_factory.mktemp("pan2-single") / "single.pt" + _run_torchrun(1, path) + return torch.load(path, map_location="cpu", weights_only=True) + + +@pytest.mark.parametrize("sp_size", SP_WORLD_SIZES) +def test_pan2_sequence_parallel_matches_single_rank(tmp_path: Path, single_rank_sp_outputs, sp_size: int): + """T2V and I2V joint attention excludes padding at Ulysses SP=2/4. Uses gloo on CPU or NCCL on 2/4 GPUs.""" + if torch.cuda.is_available() and torch.cuda.device_count() < sp_size: + pytest.skip(f"PAN2 SP parity requires {sp_size} CUDA devices.") + sp_path = tmp_path / f"sp{sp_size}.pt" + _run_torchrun(sp_size, sp_path) + sp_outputs = torch.load(sp_path, map_location="cpu", weights_only=True) + assert sp_outputs.keys() == single_rank_sp_outputs.keys() + for case, single in single_rank_sp_outputs.items(): + sp = sp_outputs[case] + diff = (sp - single).abs() + print(f"sp{sp_size}-vs-single {case} diff_max={diff.max().item():.3e}") + assert_close(sp, single, atol=1e-5, rtol=1e-5) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--pan2-sp-worker", action="store_true") + parser.add_argument("--sp-size", type=int, choices=[1, *SP_WORLD_SIZES], required=True) + parser.add_argument("--output", type=str, required=True) + args = parser.parse_args() + if not args.pan2_sp_worker: + raise SystemExit("This module is intended to be run by pytest.") + _run_sp_worker(args.sp_size, Path(args.output)) diff --git a/tests/local_tests/vaes/test_pan2_vae_parity.py b/tests/local_tests/vaes/test_pan2_vae_parity.py new file mode 100644 index 000000000..3f56210f3 --- /dev/null +++ b/tests/local_tests/vaes/test_pan2_vae_parity.py @@ -0,0 +1,387 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PAN2 reuse of FastVideo's Wan2.2 (48-channel) VAE. + +Coverage scope: both. PAN2 ships a Diffusers ``AutoencoderKLWan`` in the +Wan2.2 layout (z_dim 48, residual blocks, 2x2 patchify, 16x spatial and 4x +temporal compression). The FastVideo side loads it through ``VAELoader``, the +path a pipeline uses, and is compared on CPU fp32 against +``diffusers.AutoencoderKLWan.from_pretrained``: + +* posterior mode/logvar for a single first frame (PAN2's I2V conditioning), +* posterior mode/logvar for a 9-frame clip (4x temporal downsampling + cache), +* decode of 3 latent frames (first-chunk handling + temporal upsampling cache). + +The normalization tests use a copy of the checkpoint with non-trivial +``latents_mean``/``latents_std`` (the tiny checkpoint ships 0/1, which would +hide a missing or doubled normalization) and pin where it happens: the raw +``encode``/``decode`` never normalize, ``PAN2ConditioningStage`` applies +PAN2's ``(z - mean) / std`` once, and ``DecodingStage`` applies +``z * std + mean`` once. + +The default asset is the random-weight tiny PAN2 checkpoint, so this runs on a +CPU-only host:: + + pytest tests/local_tests/vaes/test_pan2_vae_parity.py -v -s + +Set ``PAN2_MODEL_PATH`` to a local PAN2 root containing ``vae/`` to use other +weights. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import load_file +from torch.testing import assert_close + +try: + from diffusers import AutoencoderKLWan as DiffusersAutoencoderKLWan +except ImportError: # pragma: no cover - environment dependent + pytest.skip("diffusers with AutoencoderKLWan is required", allow_module_level=True) + +from fastvideo.models.loader.component_loader import VAELoader +from fastvideo.models.wan.vae import AutoencoderKLWan as FastVideoAutoencoderKLWan +from fastvideo.models.wan.vae_config import WanVAEConfig + +PARITY_SCOPE = "both" +TINY_PAN2_REPO = "wuqing157/tiny-pan2-modular-pipe" +TINY_PAN2_REVISION = "3ff3130b4527f61aaab7bde01465c9884dc84cdd" +MODEL_PATH_ENV = "PAN2_MODEL_PATH" + +# H and W must be multiples of 32 (16x spatial compression times the 2x2 +# patchify inside the VAE). +HEIGHT, WIDTH = 64, 96 +VIDEO_FRAMES = 9 +LATENT_FRAMES = (VIDEO_FRAMES - 1) // 4 + 1 +ATOL = RTOL = 1e-5 + +# Structural fields of Wan-AI/Wan2.2-TI2V-5B-Diffusers vae/config.json. The +# statistics do not affect parameter shapes, so placeholders keep this inline. +WAN22_TI2V_VAE_CONFIG = { + "attn_scales": [], + "base_dim": 160, + "clip_output": False, + "decoder_base_dim": 256, + "dim_mult": [1, 2, 4, 4], + "dropout": 0.0, + "in_channels": 12, + "is_residual": True, + "latents_mean": [0.0] * 48, + "latents_std": [1.0] * 48, + "num_res_blocks": 2, + "out_channels": 12, + "patch_size": 2, + "scale_factor_spatial": 16, + "scale_factor_temporal": 4, + "temperal_downsample": [False, True, True], + "z_dim": 48, +} + + +def _pan2_root(subfolder: str) -> Path: + """Resolve a PAN2 root that contains ``subfolder``; skip when unavailable.""" + override = os.getenv(MODEL_PATH_ENV) + if override: + root = Path(override) + else: + try: + from huggingface_hub import snapshot_download + except ImportError: + pytest.skip("huggingface_hub is required to resolve the tiny PAN2 checkpoint") + root = None + last_error: Exception | None = None + for local_files_only in (True, False): + try: + root = Path( + snapshot_download( + TINY_PAN2_REPO, + revision=TINY_PAN2_REVISION, + allow_patterns=[f"{subfolder}/*"], + local_files_only=local_files_only, + )) + break + except Exception as exc: # noqa: BLE001 - offline/missing cache must skip. + last_error = exc + if root is None: + pytest.skip(f"Tiny PAN2 checkpoint unavailable: {last_error}") + if not (root / subfolder).is_dir(): + pytest.skip(f"PAN2 {subfolder}/ missing under {root}") + return root + + +@pytest.fixture(scope="module") +def vae_dir() -> Path: + path = _pan2_root("vae") / "vae" + for name in ("config.json", "diffusion_pytorch_model.safetensors"): + if not (path / name).is_file(): + pytest.skip(f"PAN2 VAE asset missing: {path / name}") + return path + + +@pytest.fixture +def cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None: + """Pin FastVideo's device helper to CPU; it reports ``mps`` on CPU-only Linux.""" + for module in ( + "fastvideo.models.loader.component_loader", + "fastvideo.pipelines.stages.decoding", + "fastvideo.pipelines.basic.pan2.stages", + ): + monkeypatch.setattr(f"{module}.get_local_torch_device", lambda: torch.device("cpu")) + + +def _pan2_vae_config() -> WanVAEConfig: + """The VAE config a PAN2 pipeline config would carry (arch filled by the loader).""" + config = WanVAEConfig() + config.load_encoder = True + config.load_decoder = True + return config + + +def _load_fastvideo(vae_dir: Path) -> FastVideoAutoencoderKLWan: + """Load through the production VAELoader, exactly as a pipeline does.""" + fastvideo_args = SimpleNamespace( + model_paths={}, + pipeline_config=SimpleNamespace(vae_config=_pan2_vae_config(), vae_precision="fp32"), + vae_cpu_offload=False, + ) + vae = VAELoader().load(str(vae_dir), fastvideo_args) + assert isinstance(vae, FastVideoAutoencoderKLWan) + assert fastvideo_args.model_paths["vae"] == str(vae_dir) + assert all(param.dtype == torch.float32 and param.device.type == "cpu" for param in vae.parameters()) + return vae + + +def _load_diffusers(vae_dir: Path) -> DiffusersAutoencoderKLWan: + model, info = DiffusersAutoencoderKLWan.from_pretrained( + str(vae_dir), + torch_dtype=torch.float32, + local_files_only=True, + output_loading_info=True, + ) + problems = {key: value for key, value in info.items() if value} + assert not problems, f"Diffusers reference did not load strictly: {problems}" + return model.eval() + + +def _pixels(num_frames: int, seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.rand((1, 3, num_frames, HEIGHT, WIDTH), generator=generator) * 2.0 - 1.0 + + +def _latents(seed: int) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randn((1, 48, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16), generator=generator) + + +def _report(name: str, actual: torch.Tensor, expected: torch.Tensor, clamp_bound: float | None = None) -> None: + assert actual.shape == expected.shape, f"{name}: {tuple(actual.shape)} != {tuple(expected.shape)}" + diff = (actual.float() - expected.float()).abs() + message = (f"[pan2-vae] {name}: shape={tuple(actual.shape)} max_abs={diff.max().item():.3e} " + f"mean_abs={diff.mean().item():.3e} ref_absmax={expected.abs().max().item():.3e}") + if clamp_bound is not None: + # Clamped pixels compare trivially; report how many there are. + message += f" clamped={(expected.abs() >= clamp_bound).float().mean().item():.2%}" + print(message) + + +def _nontrivial_stats() -> tuple[list[float], list[float]]: + mean = [round(0.05 * (index - 24), 4) for index in range(48)] + std = [round(0.4 + 0.03 * index, 4) for index in range(48)] + return mean, std + + +@pytest.fixture +def stats_vae_dir(vae_dir: Path, tmp_path: Path) -> Path: + """Checkpoint copy with non-trivial latent statistics (weights symlinked).""" + config = json.loads((vae_dir / "config.json").read_text()) + config["latents_mean"], config["latents_std"] = _nontrivial_stats() + target = tmp_path / "vae" + target.mkdir() + (target / "config.json").write_text(json.dumps(config)) + (target / "diffusion_pytorch_model.safetensors").symlink_to( + (vae_dir / "diffusion_pytorch_model.safetensors").resolve()) + return target + + +def _diffusers_stats(model: DiffusersAutoencoderKLWan, like: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mean = torch.tensor(model.config.latents_mean).view(1, -1, 1, 1, 1).to(like) + std = torch.tensor(model.config.latents_std).view(1, -1, 1, 1, 1).to(like) + return mean, std + + +def test_pan2_vae_production_loader_matches_checkpoint(vae_dir: Path, cpu_runtime: None) -> None: + """The loader fills the tiny Wan2.2 arch from config.json and covers every weight. + + ``VAELoader`` loads ``AutoencoderKLWan`` with ``strict=False``, so key + coverage is checked here instead of trusting the loader. + """ + raw = json.loads((vae_dir / "config.json").read_text()) + vae = _load_fastvideo(vae_dir) + arch = vae.config.arch_config + for key in ("base_dim", "decoder_base_dim", "z_dim", "num_res_blocks", "in_channels", "out_channels", + "is_residual", "patch_size", "scale_factor_spatial", "scale_factor_temporal", "dropout"): + assert getattr(arch, key) == raw[key], key + for key in ("dim_mult", "attn_scales", "temperal_downsample", "latents_mean", "latents_std"): + assert list(getattr(arch, key)) == list(raw[key]), key + assert vae.config.use_feature_cache and not vae.config.use_tiling + assert arch.spatial_compression_ratio == 16 and arch.temporal_compression_ratio == 4 + assert tuple(vae.scaling_factor.shape) == tuple(vae.shift_factor.shape) == (1, 48, 1, 1, 1) + assert_close(vae.scaling_factor.flatten(), 1.0 / torch.tensor(raw["latents_std"])) + assert_close(vae.shift_factor.flatten(), torch.tensor(raw["latents_mean"])) + + checkpoint = load_file(str(vae_dir / "diffusion_pytorch_model.safetensors")) + state = vae.state_dict() + assert set(state) == set(checkpoint), (sorted(set(checkpoint) - set(state)), sorted(set(state) - set(checkpoint))) + for key, tensor in checkpoint.items(): + assert torch.equal(state[key], tensor.float()), key + + +@pytest.mark.parametrize("variant", ["wan22_ti2v_5b", "tiny_pan2"]) +def test_pan2_vae_config_shapes_match_diffusers(variant: str, request: pytest.FixtureRequest) -> None: + """``WanVAEConfig`` expresses both the real Wan2.2 and the tiny PAN2 VAE configs. + + Both models are built on the meta device, so the full-size config costs no + memory; parameter names and shapes must match one-to-one. + """ + if variant == "tiny_pan2": + raw = json.loads((request.getfixturevalue("vae_dir") / "config.json").read_text()) + config = {key: value for key, value in raw.items() if not key.startswith("_")} + else: + config = WAN22_TI2V_VAE_CONFIG + fastvideo_config = _pan2_vae_config() + fastvideo_config.update_model_arch(config) + with torch.device("meta"): + fastvideo = FastVideoAutoencoderKLWan(fastvideo_config) + reference = DiffusersAutoencoderKLWan.from_config(config) + fastvideo_shapes = {key: tuple(value.shape) for key, value in fastvideo.state_dict().items()} + reference_shapes = {key: tuple(value.shape) for key, value in reference.state_dict().items()} + assert fastvideo_shapes == reference_shapes + + +def test_pan2_vae_wan22_topology_random_weight_parity() -> None: + """Numerically cover the Wan2.2 paths the tiny config skips. + + The tiny checkpoint uses ``dim_mult`` [1, 1, 1, 1] and no separate decoder + width, so channel-changing shortcuts and resamplers never run. This narrow + random-weight model keeps the real topology (``dim_mult`` [1, 2, 4, 4], two + res blocks, ``decoder_base_dim`` > ``base_dim``). + """ + config = dict(WAN22_TI2V_VAE_CONFIG, base_dim=8, decoder_base_dim=16) + torch.manual_seed(0) + reference = DiffusersAutoencoderKLWan.from_config(config).eval() + fastvideo_config = _pan2_vae_config() + fastvideo_config.update_model_arch(config) + fastvideo = FastVideoAutoencoderKLWan(fastvideo_config).eval() + fastvideo.load_state_dict(reference.state_dict(), strict=True) + + video = _pixels(VIDEO_FRAMES, seed=11) + latents = _latents(seed=12) + with torch.inference_mode(): + expected_mode = reference.encode(video).latent_dist.mode() + actual_mode = fastvideo.encode(video).mode() + expected_video = reference.decode(latents).sample + actual_video = fastvideo.decode(latents) + _report("wan22-topology encode mode", actual_mode, expected_mode) + _report("wan22-topology decode", actual_video, expected_video, clamp_bound=1.0) + assert_close(actual_mode, expected_mode, atol=ATOL, rtol=RTOL) + assert_close(actual_video, expected_video, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize("num_frames", [1, VIDEO_FRAMES], ids=["first_frame", "video_9f"]) +def test_pan2_vae_encode_matches_diffusers(num_frames: int, vae_dir: Path, cpu_runtime: None) -> None: + """Posterior mode (PAN2 uses ``latent_dist.mode()``) and logvar match Diffusers. + + FastVideo's ``encode`` returns the ``DiagonalGaussianDistribution`` itself, + so the port calls ``vae.encode(x).mode()`` instead of ``.latent_dist.mode()``. + """ + reference = _load_diffusers(vae_dir) + fastvideo = _load_fastvideo(vae_dir) + pixels = _pixels(num_frames, seed=100 + num_frames) + with torch.inference_mode(): + expected = reference.encode(pixels).latent_dist + actual = fastvideo.encode(pixels) + assert tuple(actual.mode().shape) == (1, 48, (num_frames - 1) // 4 + 1, HEIGHT // 16, WIDTH // 16) + _report(f"encode mode ({num_frames}f)", actual.mode(), expected.mode()) + _report(f"encode logvar ({num_frames}f)", actual.logvar, expected.logvar) + assert_close(actual.mode(), expected.mode(), atol=ATOL, rtol=RTOL) + assert_close(actual.logvar, expected.logvar, atol=ATOL, rtol=RTOL) + + +def test_pan2_vae_decode_matches_diffusers(vae_dir: Path, cpu_runtime: None) -> None: + """Decode three latent frames so first-chunk and temporal cache paths run.""" + reference = _load_diffusers(vae_dir) + fastvideo = _load_fastvideo(vae_dir) + latents = _latents(seed=7) + with torch.inference_mode(): + expected = reference.decode(latents).sample + actual = fastvideo.decode(latents) + assert tuple(actual.shape) == (1, 3, VIDEO_FRAMES, HEIGHT, WIDTH) + _report("decode", actual, expected, clamp_bound=1.0) + assert_close(actual, expected, atol=ATOL, rtol=RTOL) + + +def test_pan2_conditioning_stage_matches_reference(stats_vae_dir: Path, cpu_runtime: None) -> None: + """``PAN2ConditioningStage`` builds the reference I2V condition: cover-resize and center crop the first frame, + encode it, apply ``(z - mean) / std`` once, and set the mask channel on the first latent frame only. + + The raw ``encode`` is unnormalized, so the stage is the only place the statistics apply. + """ + import PIL.Image + + from fastvideo.pipelines.basic.pan2.stages import PAN2ConditioningStage + from tests.local_tests.pan2.reference import reference_condition_latents + + reference = _load_diffusers(stats_vae_dir) + fastvideo = _load_fastvideo(stats_vae_dir) + # Not the output aspect ratio, so the resize and crop are exercised. + image = PIL.Image.radial_gradient("L").resize((120, 64)).convert("RGB") + latents = torch.zeros(1, 48, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16) + with torch.inference_mode(): + pixels = PAN2ConditioningStage.preprocess_image(image, HEIGHT, WIDTH) + raw = reference.encode(pixels).latent_dist.mode() + assert_close(fastvideo.encode(pixels).mode(), raw, atol=ATOL, rtol=RTOL) + expected = reference_condition_latents(reference, latents, image, HEIGHT, WIDTH) + + batch = SimpleNamespace(latents=latents, pil_image=image, height=HEIGHT, width=WIDTH, image_latent=None) + stage_args = SimpleNamespace(pipeline_config=SimpleNamespace(vae_precision="fp32"), vae_cpu_offload=False) + actual = PAN2ConditioningStage(fastvideo).forward(batch, stage_args).image_latent + _report("PAN2 I2V condition", actual, expected) + assert actual.shape == (1, 49, LATENT_FRAMES, HEIGHT // 16, WIDTH // 16) + assert_close(actual, expected, atol=ATOL, rtol=RTOL) + assert torch.equal(actual[:, 48, 0], torch.ones_like(actual[:, 48, 0])) + assert not actual[:, :, 1:].any() + assert not torch.allclose(raw, expected[:, :48, :1], atol=1e-2), "statistics too close to identity" + + +def test_pan2_decoding_stage_denormalizes_once(stats_vae_dir: Path, cpu_runtime: None) -> None: + """``DecodingStage`` performs PAN2's ``z * std + mean`` itself. + + The PAN2 denoising loop must hand normalized latents to ``DecodingStage``. + Denormalizing beforehand (as the upstream pipeline does inline) would apply + the statistics twice, which the negative check below detects. + """ + from fastvideo.pipelines.stages.decoding import DecodingStage + + reference = _load_diffusers(stats_vae_dir) + fastvideo = _load_fastvideo(stats_vae_dir) + latents = _latents(seed=31) + # vae_tiling=True is the PipelineConfig default; the feature-cache decode ignores tiling, so it must not matter. + decode_args = SimpleNamespace( + pipeline_config=SimpleNamespace(vae_precision="fp32", vae_decode_precision="fp32", vae_tiling=True), + disable_autocast=False, + ) + stage = DecodingStage(fastvideo) + with torch.inference_mode(): + mean, std = _diffusers_stats(reference, latents) + expected = reference.decode(latents * std + mean).sample + expected = (expected * 0.5 + 0.5).clamp(0, 1) + actual = stage.decode(latents, decode_args) + doubled = stage.decode(latents * std + mean, decode_args) + _report("DecodingStage output", actual, expected) + assert_close(actual, expected, atol=ATOL, rtol=RTOL) + assert (doubled - expected).abs().max().item() > 1e-3, "double denormalization went undetected"