847404748f
Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: filipstrand <filipstrand@users.noreply.github.com>
92 lines
3.2 KiB
Python
92 lines
3.2 KiB
Python
from pathlib import Path
|
|
|
|
from mflux import Config, ModelConfig, StopImageGenerationException
|
|
from mflux.callbacks.callback_manager import CallbackManager
|
|
from mflux.error.exceptions import PromptFileReadError
|
|
from mflux.flux_tools.redux.flux_redux import Flux1Redux
|
|
from mflux.ui import defaults as ui_defaults
|
|
from mflux.ui.cli.parsers import CommandLineParser
|
|
from mflux.ui.prompt_utils import get_effective_prompt
|
|
|
|
|
|
def main():
|
|
# 0. Parse command line arguments
|
|
parser = CommandLineParser(description="Generate an image using the fill tool to complete masked areas.")
|
|
parser.add_general_arguments()
|
|
parser.add_model_arguments(require_model_arg=False)
|
|
parser.add_lora_arguments()
|
|
parser.add_image_generator_arguments(supports_metadata_config=False)
|
|
parser.add_redux_arguments()
|
|
parser.add_output_arguments()
|
|
args = parser.parse_args()
|
|
|
|
# 0. Set default guidance value if not provided by user
|
|
if args.guidance is None:
|
|
args.guidance = ui_defaults.GUIDANCE_SCALE
|
|
|
|
# Validate and normalize redux image strengths
|
|
redux_image_strengths = _validate_redux_image_strengths(
|
|
redux_image_paths=args.redux_image_paths,
|
|
redux_image_strengths=args.redux_image_strengths,
|
|
)
|
|
|
|
# 1. Load the model
|
|
flux = Flux1Redux(
|
|
model_config=ModelConfig.dev_redux(),
|
|
quantize=args.quantize,
|
|
local_path=args.path,
|
|
lora_paths=args.lora_paths,
|
|
lora_scales=args.lora_scales,
|
|
)
|
|
|
|
# 2. Register callbacks
|
|
memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
|
|
|
|
try:
|
|
for seed in args.seed:
|
|
# 3. Generate an image for each seed value
|
|
image = flux.generate_image(
|
|
seed=seed,
|
|
prompt=get_effective_prompt(args),
|
|
config=Config(
|
|
num_inference_steps=args.steps,
|
|
height=args.height,
|
|
width=args.width,
|
|
guidance=args.guidance,
|
|
redux_image_paths=args.redux_image_paths,
|
|
redux_image_strengths=redux_image_strengths,
|
|
),
|
|
)
|
|
|
|
# 4. Save the image
|
|
image.save(path=args.output.format(seed=seed), export_json_metadata=args.metadata)
|
|
except (StopImageGenerationException, PromptFileReadError) as exc:
|
|
print(exc)
|
|
finally:
|
|
if memory_saver:
|
|
print(memory_saver.memory_stats())
|
|
|
|
|
|
def _validate_redux_image_strengths(
|
|
redux_image_paths: list[Path],
|
|
redux_image_strengths: list[float] | None,
|
|
) -> list[float] | None:
|
|
if not redux_image_strengths or len(redux_image_strengths) == 0:
|
|
return redux_image_strengths
|
|
|
|
# If strengths provided but not enough for all images, fill with default (1.0)
|
|
if len(redux_image_strengths) < len(redux_image_paths):
|
|
redux_image_strengths.extend([1.0] * (len(redux_image_paths) - len(redux_image_strengths)))
|
|
|
|
# If too many strengths provided, raise error
|
|
elif len(redux_image_strengths) > len(redux_image_paths):
|
|
raise ValueError(
|
|
f"Too many strengths provided ({len(redux_image_strengths)}), expted at most {len(redux_image_paths)}."
|
|
)
|
|
|
|
return redux_image_strengths
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|