153 lines
5.5 KiB
Python
153 lines
5.5 KiB
Python
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from mflux.models.common.config import ModelConfig
|
|
from mflux.models.flux.variants.txt2img.flux import Flux1
|
|
|
|
PATH = "tests/4bit/"
|
|
SIZE_TOLERANCE_RATIO = 0.05 # allow small metadata/header differences
|
|
|
|
|
|
class TestModelSavingLora:
|
|
LORA_FILES = [
|
|
"FLUX-dev-lora-MiaoKa-Yarn-World.safetensors",
|
|
"Flux_-_Renaissance_art_style.safetensors",
|
|
]
|
|
LORA_SCALES = [0.6, 0.4]
|
|
|
|
@pytest.mark.slow
|
|
def test_save_and_load_4bit_model_with_lora(self):
|
|
# Clean up any existing temporary directories from previous test runs
|
|
TestModelSavingLora._delete_folder_if_exists(PATH)
|
|
|
|
try:
|
|
lora_paths = TestModelSavingLora._get_lora_paths()
|
|
assert len(lora_paths) == len(TestModelSavingLora.LORA_SCALES)
|
|
|
|
# given a saved quantized model on disk (without LoRA)...
|
|
fluxA = Flux1(
|
|
model_config=ModelConfig.schnell(),
|
|
quantize=4,
|
|
)
|
|
fluxA.save_model(PATH)
|
|
del fluxA
|
|
|
|
# ...and given an 'on-the-fly' quantized model which we generate an image from
|
|
fluxB = Flux1(
|
|
model_config=ModelConfig.schnell(),
|
|
quantize=4,
|
|
lora_paths=lora_paths,
|
|
lora_scales=TestModelSavingLora.LORA_SCALES,
|
|
)
|
|
image1 = fluxB.generate_image(
|
|
seed=44,
|
|
prompt="mkym this is made of wool, pizza",
|
|
num_inference_steps=2,
|
|
height=341,
|
|
width=768,
|
|
)
|
|
del fluxB
|
|
|
|
# when loading the quantized model from a local path (also without specifying bits) with a LoRA...
|
|
fluxC = Flux1(
|
|
model_config=ModelConfig.schnell(),
|
|
model_path=PATH,
|
|
lora_paths=lora_paths,
|
|
lora_scales=TestModelSavingLora.LORA_SCALES,
|
|
)
|
|
|
|
# ...and generating the identical image
|
|
image2 = fluxC.generate_image(
|
|
seed=44,
|
|
prompt="mkym this is made of wool, pizza",
|
|
num_inference_steps=2,
|
|
height=341,
|
|
width=768,
|
|
)
|
|
|
|
# then we confirm that we get the exact *identical* image in both cases
|
|
np.testing.assert_array_equal(
|
|
np.array(image1.image),
|
|
np.array(image2.image),
|
|
err_msg="image2 doesn't match image1.",
|
|
)
|
|
|
|
finally:
|
|
# cleanup
|
|
TestModelSavingLora._delete_folder_if_exists(PATH)
|
|
|
|
@pytest.mark.slow
|
|
def test_save_with_lora_has_same_shard_count_as_base(self):
|
|
base_path = "tests/4bit_base/"
|
|
lora_path = "tests/4bit_with_lora/"
|
|
|
|
TestModelSavingLora._delete_folder_if_exists(base_path)
|
|
TestModelSavingLora._delete_folder_if_exists(lora_path)
|
|
|
|
try:
|
|
flux_base = Flux1(
|
|
model_config=ModelConfig.schnell(),
|
|
quantize=4,
|
|
)
|
|
flux_base.save_model(base_path)
|
|
del flux_base
|
|
|
|
flux_lora = Flux1(
|
|
model_config=ModelConfig.schnell(),
|
|
quantize=4,
|
|
lora_paths=TestModelSavingLora._get_lora_paths(),
|
|
lora_scales=TestModelSavingLora.LORA_SCALES,
|
|
)
|
|
flux_lora.save_model(lora_path)
|
|
del flux_lora
|
|
|
|
base_shards = list((Path(base_path) / "transformer").glob("*.safetensors"))
|
|
lora_shards = list((Path(lora_path) / "transformer").glob("*.safetensors"))
|
|
|
|
assert len(base_shards) == len(lora_shards), "LoRA save should not inflate transformer shard count"
|
|
|
|
# Also assert the total saved size stays within a tight bound to catch shard bloat
|
|
base_size = TestModelSavingLora._dir_size_bytes(base_path)
|
|
lora_size = TestModelSavingLora._dir_size_bytes(lora_path)
|
|
|
|
# LoRA baking should leave sizes effectively unchanged; allow a small tolerance for metadata
|
|
max_allowed = base_size * (1 + SIZE_TOLERANCE_RATIO)
|
|
assert lora_size <= max_allowed, f"LoRA save size grew unexpectedly: base={base_size}B vs lora={lora_size}B"
|
|
finally:
|
|
TestModelSavingLora._delete_folder_if_exists(base_path)
|
|
TestModelSavingLora._delete_folder_if_exists(lora_path)
|
|
|
|
@staticmethod
|
|
def _delete_folder_if_exists(path: str) -> None:
|
|
if os.path.exists(path):
|
|
shutil.rmtree(path)
|
|
print(f"Deleted folder: {path}")
|
|
else:
|
|
print(f"Folder does not exist: {path}")
|
|
|
|
@staticmethod
|
|
def _resolve_path(path) -> Path | None:
|
|
if path is None:
|
|
return None
|
|
return Path(__file__).parent.parent / "resources" / path
|
|
|
|
@staticmethod
|
|
def _get_lora_paths() -> list[str]:
|
|
resolved_paths = [TestModelSavingLora._resolve_path(path) for path in TestModelSavingLora.LORA_FILES]
|
|
missing = [p for p in resolved_paths if p is None or not p.exists()]
|
|
if missing:
|
|
missing_names = ", ".join(sorted(p.name if p else "unknown" for p in missing))
|
|
pytest.skip(f"Missing local LoRA test asset(s): {missing_names}")
|
|
return [str(path) for path in resolved_paths if path is not None]
|
|
|
|
@staticmethod
|
|
def _dir_size_bytes(path: str | Path) -> int:
|
|
p = Path(path)
|
|
if not p.exists():
|
|
return 0
|
|
return sum(f.stat().st_size for f in p.rglob("*") if f.is_file())
|