81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
import os
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
|
|
from mflux.config.config import Config
|
|
from mflux.config.model_config import ModelConfig
|
|
from mflux.flux.flux import Flux1
|
|
from mflux.utils.version_util import VersionUtil
|
|
from mflux.weights.weight_handler import WeightHandler
|
|
|
|
PATH = "tests/4bit/"
|
|
|
|
|
|
class TestModelSaving:
|
|
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",
|
|
config=Config(
|
|
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 = WeightHandler._load_vae(root_path=Path(PATH))
|
|
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(),
|
|
local_path=PATH,
|
|
)
|
|
|
|
# then we can load the model and generate the identical image
|
|
image2 = fluxB.generate_image(
|
|
seed=42,
|
|
prompt="Luxury food photograph",
|
|
config=Config(
|
|
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}")
|