monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Refactor: Remove ConfigControlnet

Browse Source
This commit is contained in:
filipstrand 2025-02-08 13:29:32 +01:00
parent 22f98fc5a4
commit e4af40e4ee
10 changed files with 33 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/mflux/__init__.py
View File

@ -1,4 +1,4 @@
from mflux.config.config import Config, ConfigControlnet
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig, ModelLookup
from mflux.controlnet.flux_controlnet import Flux1Controlnet
from mflux.error.exceptions import StopImageGenerationException
@ -9,7 +9,6 @@ __all__ = [
"Flux1",
"Flux1Controlnet",
"Config",
"ConfigControlnet",
"ModelConfig",
"ModelLookup",
"ImageUtil",

13
src/mflux/config/config.py
View File

@ -17,6 +17,7 @@ class Config:
guidance: float = 4.0,
init_image_path: Path | None = None,
init_image_strength: float | None = None,
controlnet_strength: float | None = None,
):
if width % 16 != 0 or height % 16 != 0:
log.warning("Width and height should be multiples of 16. Rounding down.")
@ -26,16 +27,4 @@ class Config:
self.guidance = guidance
self.init_image_path = init_image_path
self.init_image_strength = init_image_strength
class ConfigControlnet(Config):
def __init__(
self,
num_inference_steps: int = 4,
width: int = 1024,
height: int = 1024,
guidance: float = 4.0,
controlnet_strength: float = 1.0,
):
super().__init__(num_inference_steps=num_inference_steps, width=width, height=height, guidance=guidance)
self.controlnet_strength = controlnet_strength

19
src/mflux/config/runtime_config.py
View File

@ -3,7 +3,7 @@ import logging
import mlx.core as mx
import numpy as np
from mflux.config.config import Config, ConfigControlnet
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
logger = logging.getLogger(__name__)
@ -12,7 +12,7 @@ logger = logging.getLogger(__name__)
class RuntimeConfig:
def __init__(
self,
config: Config | ConfigControlnet,
config: Config,
model_config: ModelConfig,
):
self.config = config
@ -54,23 +54,20 @@ class RuntimeConfig:
@property
def init_time_step(self) -> int:
if self.config.init_image_path is None:
# text to image, always begin at time step 0
# For text-to-image, always begin at time step 0.
return 0
else:
# we skip to the time step as informed by the init_image_strength
# the higher the strength number, the more time steps we skip
# For image-to-image: higher strength means we skip more steps.
strength = max(0.0, min(1.0, self.config.init_image_strength))
# if the strength is too small to even influence the image
# help the user round up so the init_image has influence at step 1
t = max(1, int(self.num_inference_steps * strength))
return t
@property
def controlnet_strength(self) -> float:
if isinstance(self.config, ConfigControlnet):
def controlnet_strength(self) -> float | None:
if self.config.controlnet_strength is not None:
return self.config.controlnet_strength
else:
raise NotImplementedError("Controlnet conditioning scale is only available for ConfigControlnet")
return None
@staticmethod
def _create_sigmas(config, model) -> mx.array:

10
src/mflux/controlnet/flux_controlnet.py
View File

@ -4,7 +4,7 @@ import mlx.core as mx
from mlx import nn
from tqdm import tqdm
from mflux.config.config import ConfigControlnet
from mflux.config.config import Config
from mflux.config.model_config import ModelConfig
from mflux.config.runtime_config import RuntimeConfig
from mflux.controlnet.controlnet_util import ControlnetUtil
@ -56,10 +56,10 @@ class Flux1Controlnet(nn.Module):
output: str,
controlnet_image_path: str,
controlnet_save_canny: bool = False,
config: ConfigControlnet = ConfigControlnet(),
config: Config = Config(),
stepwise_output_dir: Path = None,
) -> GeneratedImage:
# Create a new runtime config based on the model type and input parameters
# Convert the user config to a runtime config with derived parameters.
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.num_inference_steps))
stepwise_handler = StepwiseHandler(
@ -135,14 +135,14 @@ class Flux1Controlnet(nn.Module):
decoded = self.vae.decode(latents)
return ImageUtil.to_image(
decoded_latents=decoded,
config=config,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
config=config,
controlnet_image_path=controlnet_image_path,
generation_time=time_steps.format_dict["elapsed"],
)
def save_model(self, base_path: str) -> None:

6
src/mflux/flux/flux.py
View File

@ -53,7 +53,7 @@ class Flux1(nn.Module):
config: Config = Config(),
stepwise_output_dir: Path = None,
) -> GeneratedImage:
# Create a new runtime config based on the model type and input parameters
# Convert the user config to a runtime config with derived parameters.
config = RuntimeConfig(config, self.model_config)
time_steps = tqdm(range(config.init_time_step, config.num_inference_steps))
stepwise_handler = StepwiseHandler(
@ -108,15 +108,15 @@ class Flux1(nn.Module):
decoded = self.vae.decode(latents)
return ImageUtil.to_image(
decoded_latents=decoded,
config=config,
seed=seed,
prompt=prompt,
quantization=self.bits,
generation_time=time_steps.format_dict["elapsed"],
lora_paths=self.lora_paths,
lora_scales=self.lora_scales,
init_image_path=config.init_image_path,
init_image_strength=config.init_image_strength,
config=config,
generation_time=time_steps.format_dict["elapsed"],
)
@staticmethod

6
src/mflux/generate.py
View File

@ -14,7 +14,7 @@ def main():
parser.add_output_arguments()
args = parser.parse_args()
# Load the model
# 1. Load the model
flux = Flux1(
model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model),
quantize=args.quantize,
@ -25,7 +25,7 @@ def main():
try:
for seed_value in args.seed:
# Generate an image for each seed value
# 2. Generate an image for each seed value
image = flux.generate_image(
seed=seed_value,
prompt=args.prompt,
@ -39,7 +39,7 @@ def main():
init_image_strength=args.init_image_strength,
),
)
# Save the image
# 3. Save the image
image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata)
except StopImageGenerationException as stop_exc:
print(stop_exc)

10
src/mflux/generate_controlnet.py
View File

@ -1,6 +1,6 @@
from pathlib import Path
from mflux import ConfigControlnet, Flux1Controlnet, ModelLookup, StopImageGenerationException
from mflux import Config, Flux1Controlnet, ModelLookup, StopImageGenerationException
from mflux.ui.cli.parsers import CommandLineParser
@ -13,7 +13,7 @@ def main():
parser.add_output_arguments()
args = parser.parse_args()
# Load the model
# 1. Load the model
flux = Flux1Controlnet(
model_config=ModelLookup.from_name(model_name=args.model, base_model=args.base_model),
quantize=args.quantize,
@ -24,7 +24,7 @@ def main():
try:
for seed_value in args.seed:
# Generate an image for each seed value
# 2. Generate an image for each seed value
image = flux.generate_image(
seed=seed_value,
prompt=args.prompt,
@ -32,7 +32,7 @@ def main():
controlnet_image_path=args.controlnet_image_path,
controlnet_save_canny=args.controlnet_save_canny,
stepwise_output_dir=Path(args.stepwise_image_output_dir) if args.stepwise_image_output_dir else None,
config=ConfigControlnet(
config=Config(
num_inference_steps=args.steps,
height=args.height,
width=args.width,
@ -41,7 +41,7 @@ def main():
),
)
# Save the image
# 3. Save the image
image.save(path=args.output.format(seed=seed_value), export_json_metadata=args.metadata)
except StopImageGenerationException as stop_exc:
print(stop_exc)

8
src/mflux/post_processing/image_util.py
View File

@ -8,25 +8,23 @@ import numpy as np
import piexif
import PIL.Image
from mflux.config.config import ConfigControlnet
from mflux.config.runtime_config import RuntimeConfig
from mflux.post_processing.generated_image import GeneratedImage
log = logging.getLogger(__name__)
RuntimeConfig: t.TypeAlias = "mflux.config.runtime_config.RuntimeConfig" # noqa: F821
class ImageUtil:
@staticmethod
def to_image(
decoded_latents: mx.array,
config: RuntimeConfig,
seed: int,
prompt: str,
quantization: int,
generation_time: float,
lora_paths: list[str],
lora_scales: list[float],
config: RuntimeConfig,
controlnet_image_path: str | None = None,
init_image_path: str | None = None,
init_image_strength: float | None = None,
@ -49,7 +47,7 @@ class ImageUtil:
init_image_path=init_image_path,
init_image_strength=init_image_strength,
controlnet_image_path=controlnet_image_path,
controlnet_strength=config.controlnet_strength if isinstance(config.config, ConfigControlnet) else None,
controlnet_strength=config.controlnet_strength,
)
@staticmethod

4
src/mflux/post_processing/stepwise_handler.py
View File

@ -40,13 +40,13 @@ class StepwiseHandler:
stepwise_decoded = self.flux.vae.decode(unpack_latents)
stepwise_img = ImageUtil.to_image(
decoded_latents=stepwise_decoded,
config=self.config,
seed=self.seed,
prompt=self.prompt,
quantization=self.flux.bits,
generation_time=self.time_steps.format_dict["elapsed"],
lora_paths=self.flux.lora_paths,
lora_scales=self.flux.lora_scales,
config=self.config,
generation_time=self.time_steps.format_dict["elapsed"],
)
self.step_wise_images.append(stepwise_img)

4
tests/image_generation/helpers/image_generation_controlnet_test_helper.py
View File

@ -3,7 +3,7 @@ import os
import numpy as np
from PIL import Image
from mflux import ConfigControlnet, Flux1Controlnet, ModelConfig
from mflux import Config, Flux1Controlnet, ModelConfig
from tests.image_generation.helpers.image_generation_test_helper import ImageGeneratorTestHelper
@ -43,7 +43,7 @@ class ImageGeneratorControlnetTestHelper:
output=str(output_image_path),
controlnet_image_path=controlnet_image_path,
controlnet_save_canny=False,
config=ConfigControlnet(
config=Config(
num_inference_steps=steps,
height=768,
width=493,