78 lines
2.5 KiB
Python
78 lines
2.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.common.weights.loading.weight_loader import WeightLoader
|
|
from mflux.models.flux.variants.txt2img.flux import Flux1
|
|
from mflux.utils.version_util import VersionUtil
|
|
|
|
PATH = "tests/4bit/"
|
|
|
|
|
|
class TestModelSaving:
|
|
@pytest.mark.slow
|
|
def test_save_and_load_4bit_model(self):
|
|
# Clean up any existing temporary directories from previous test runs
|
|
TestModelSaving.delete_folder_if_exists(PATH)
|
|
|
|
try:
|
|
# given a saved quantized model (and an image from that model)
|
|
fluxA = Flux1(
|
|
model_config=ModelConfig.dev(),
|
|
quantize=4,
|
|
)
|
|
image1 = fluxA.generate_image(
|
|
seed=42,
|
|
prompt="Luxury food photograph",
|
|
num_inference_steps=15,
|
|
height=341,
|
|
width=768,
|
|
)
|
|
fluxA.save_model(PATH)
|
|
del fluxA
|
|
|
|
# Verify that the mflux version is correctly saved in the model's metadata
|
|
_, quantization_level, mflux_version = WeightLoader._try_load_mflux_format(Path(PATH) / "vae")
|
|
assert mflux_version == VersionUtil.get_mflux_version(), "mflux version not correctly saved in metadata" # fmt: off
|
|
assert quantization_level == 4, "quantization level not correctly saved in metadata" # fmt: off
|
|
|
|
# when loading the quantized model (also without specifying bits)
|
|
fluxB = Flux1(
|
|
model_config=ModelConfig.dev(),
|
|
model_path=PATH,
|
|
)
|
|
|
|
# then we can load the model and generate the identical image
|
|
image2 = fluxB.generate_image(
|
|
seed=42,
|
|
prompt="Luxury food photograph",
|
|
num_inference_steps=15,
|
|
height=341,
|
|
width=768,
|
|
)
|
|
np.testing.assert_array_equal(
|
|
np.array(image1.image),
|
|
np.array(image2.image),
|
|
err_msg="image2 doesn't match image1.",
|
|
)
|
|
|
|
finally:
|
|
# cleanup
|
|
TestModelSaving.delete_folder(PATH)
|
|
|
|
@staticmethod
|
|
def delete_folder(path: str) -> None:
|
|
return shutil.rmtree(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}")
|