|
| 1 | +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +""" |
| 16 | +Functional tests for FLUX HF <-> Megatron checkpoint conversion. |
| 17 | +
|
| 18 | +Uses a tiny randomly initialized diffusers FluxTransformer2DModel saved under a local |
| 19 | +``transformer/`` layout (same as FLUX.1-dev) and drives |
| 20 | +``examples/diffusion/recipes/flux/conversion/convert_checkpoints.py``. |
| 21 | +
|
| 22 | +""" |
| 23 | + |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +import subprocess |
| 27 | +from pathlib import Path |
| 28 | + |
| 29 | +import pytest |
| 30 | +import torch |
| 31 | + |
| 32 | + |
| 33 | +diffusers = pytest.importorskip("diffusers") |
| 34 | +FluxTransformer2DModel = diffusers.FluxTransformer2DModel |
| 35 | + |
| 36 | +# Repo root: tests/functional_tests/diffusion/flux -> five parents |
| 37 | +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent.parent |
| 38 | +CONVERT_SCRIPT = "examples/diffusion/recipes/flux/conversion/convert_checkpoints.py" |
| 39 | + |
| 40 | + |
| 41 | +def _build_toy_flux_hf_root(base: Path) -> Path: |
| 42 | + """ |
| 43 | + Create ``base`` as a FLUX-style HF root with weights under ``base/transformer/``. |
| 44 | +
|
| 45 | + Dimensions are chosen so ``sum(axes_dims_rope) == attention_head_dim`` (128), matching |
| 46 | + real FLUX layouts, while keeping layers minimal for speed. |
| 47 | + """ |
| 48 | + hf_root = base / "flux_toy_hf" |
| 49 | + transformer_dir = hf_root / "transformer" |
| 50 | + transformer_dir.mkdir(parents=True, exist_ok=True) |
| 51 | + |
| 52 | + model = FluxTransformer2DModel( |
| 53 | + patch_size=1, |
| 54 | + in_channels=64, |
| 55 | + num_layers=1, |
| 56 | + num_single_layers=2, |
| 57 | + attention_head_dim=128, |
| 58 | + num_attention_heads=8, |
| 59 | + joint_attention_dim=4096, |
| 60 | + pooled_projection_dim=768, |
| 61 | + guidance_embeds=False, |
| 62 | + axes_dims_rope=(16, 56, 56), |
| 63 | + ) |
| 64 | + model = model.to(dtype=torch.bfloat16) |
| 65 | + model.save_pretrained(transformer_dir, safe_serialization=True) |
| 66 | + return hf_root |
| 67 | + |
| 68 | + |
| 69 | +class TestFluxCheckpointConversion: |
| 70 | + """HF FLUX transformer (diffusers) <-> Megatron via convert_checkpoints.py.""" |
| 71 | + |
| 72 | + @pytest.fixture(scope="class") |
| 73 | + def flux_toy_hf_root(self, tmp_path_factory: pytest.TempPathFactory) -> Path: |
| 74 | + base = tmp_path_factory.mktemp("flux_conv_toy") |
| 75 | + return _build_toy_flux_hf_root(Path(base)) |
| 76 | + |
| 77 | + @pytest.fixture(scope="class") |
| 78 | + def flux_megatron_ckpt_dir(self, flux_toy_hf_root: Path, tmp_path_factory: pytest.TempPathFactory) -> Path: |
| 79 | + """ |
| 80 | + Run ``convert_checkpoints.py import`` once for the class. |
| 81 | +
|
| 82 | + Export test no longer repeats this (~10–30s+ per run depending on GPU). |
| 83 | + """ |
| 84 | + megatron_out = Path(tmp_path_factory.mktemp("flux_megatron_shared")) |
| 85 | + megatron_out.mkdir(parents=True, exist_ok=True) |
| 86 | + cmd = [ |
| 87 | + "python", |
| 88 | + "-m", |
| 89 | + "coverage", |
| 90 | + "run", |
| 91 | + "--data-file=/opt/Megatron-Bridge/.coverage", |
| 92 | + "--source=/opt/Megatron-Bridge/", |
| 93 | + "--parallel-mode", |
| 94 | + CONVERT_SCRIPT, |
| 95 | + "import", |
| 96 | + "--hf-model", |
| 97 | + str(flux_toy_hf_root), |
| 98 | + "--megatron-path", |
| 99 | + str(megatron_out), |
| 100 | + "--torch-dtype", |
| 101 | + "bfloat16", |
| 102 | + ] |
| 103 | + result = subprocess.run(cmd, capture_output=True, text=True, cwd=REPO_ROOT) |
| 104 | + if result.returncode != 0: |
| 105 | + print("STDOUT:", result.stdout) |
| 106 | + print("STDERR:", result.stderr) |
| 107 | + pytest.fail(f"FLUX import (class fixture) failed with exit code {result.returncode}") |
| 108 | + assert "Successfully imported model to:" in result.stdout |
| 109 | + assert any(megatron_out.iterdir()), f"Megatron output empty: {megatron_out}" |
| 110 | + return megatron_out |
| 111 | + |
| 112 | + def test_toy_flux_layout(self, flux_toy_hf_root: Path) -> None: |
| 113 | + tdir = flux_toy_hf_root / "transformer" |
| 114 | + assert tdir.is_dir() |
| 115 | + assert (tdir / "config.json").is_file() |
| 116 | + assert (tdir / "diffusion_pytorch_model.safetensors").is_file() or ( |
| 117 | + tdir / "diffusion_pytorch_model.bin" |
| 118 | + ).is_file() |
| 119 | + |
| 120 | + def test_import_hf_to_megatron(self, flux_megatron_ckpt_dir: Path) -> None: |
| 121 | + """Asserts the shared class import produced a non-empty Megatron checkpoint.""" |
| 122 | + assert flux_megatron_ckpt_dir.is_dir() |
| 123 | + assert any(flux_megatron_ckpt_dir.iterdir()) |
| 124 | + |
| 125 | + def test_export_megatron_to_hf_roundtrip( |
| 126 | + self, flux_toy_hf_root: Path, flux_megatron_ckpt_dir: Path, tmp_path: Path |
| 127 | + ) -> None: |
| 128 | + hf_export = tmp_path / "hf_export" |
| 129 | + hf_export.mkdir(parents=True, exist_ok=True) |
| 130 | + |
| 131 | + export_cmd = [ |
| 132 | + "python", |
| 133 | + "-m", |
| 134 | + "coverage", |
| 135 | + "run", |
| 136 | + "--data-file=/opt/Megatron-Bridge/.coverage", |
| 137 | + "--source=/opt/Megatron-Bridge/", |
| 138 | + "--parallel-mode", |
| 139 | + CONVERT_SCRIPT, |
| 140 | + "export", |
| 141 | + "--hf-model", |
| 142 | + str(flux_toy_hf_root), |
| 143 | + "--megatron-path", |
| 144 | + str(flux_megatron_ckpt_dir), |
| 145 | + "--hf-path", |
| 146 | + str(hf_export), |
| 147 | + "--no-progress", |
| 148 | + ] |
| 149 | + er = subprocess.run(export_cmd, capture_output=True, text=True, cwd=REPO_ROOT) |
| 150 | + if er.returncode != 0: |
| 151 | + print("EXPORT STDOUT:", er.stdout) |
| 152 | + print("EXPORT STDERR:", er.stderr) |
| 153 | + pytest.fail(f"FLUX export failed with exit code {er.returncode}") |
| 154 | + |
| 155 | + assert "Successfully exported model to:" in er.stdout |
| 156 | + transformer_out = hf_export / "transformer" |
| 157 | + assert transformer_out.is_dir(), f"Missing {transformer_out}" |
| 158 | + assert (transformer_out / "config.json").is_file() |
| 159 | + st = transformer_out / "diffusion_pytorch_model.safetensors" |
| 160 | + pt = transformer_out / "diffusion_pytorch_model.bin" |
| 161 | + assert st.is_file() or pt.is_file(), f"No weights under {transformer_out}" |
0 commit comments