107 lines
3.9 KiB
Python
107 lines
3.9 KiB
Python
import sys
|
|
|
|
import PIL.Image
|
|
|
|
from mflux.callbacks.callback_manager import CallbackManager
|
|
from mflux.config.config import Config
|
|
from mflux.config.model_config import ModelConfig
|
|
from mflux.controlnet.flux_controlnet import Flux1Controlnet
|
|
from mflux.error.exceptions import PromptFileReadError, StopImageGenerationException
|
|
from mflux.ui import defaults as ui_defaults
|
|
from mflux.ui.cli.parsers import CommandLineParser
|
|
from mflux.ui.prompt_utils import get_effective_prompt
|
|
from mflux.ui.scale_factor import ScaleFactor
|
|
|
|
|
|
def main():
|
|
# 0. Parse command line arguments
|
|
parser = CommandLineParser(description="Upscale an image.")
|
|
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, supports_dimension_scale_factor=True)
|
|
parser.add_controlnet_arguments(require_image=True)
|
|
parser.add_output_arguments()
|
|
args = parser.parse_args()
|
|
|
|
# 1. Load the model
|
|
flux = Flux1Controlnet(
|
|
model_config=ModelConfig.dev_controlnet_upscaler(),
|
|
quantize=args.quantize,
|
|
local_path=args.path,
|
|
lora_paths=args.lora_paths,
|
|
lora_scales=args.lora_scales,
|
|
)
|
|
|
|
# 2. Register the optional callbacks
|
|
memory_saver = CallbackManager.register_callbacks(args=args, flux=flux)
|
|
|
|
try:
|
|
# Calculate output dimensions and handle safety warnings
|
|
width, height = _calculate_output_dimensions(args)
|
|
|
|
for seed in args.seed:
|
|
# 3. Generate an upscaled image for each seed value
|
|
image = flux.generate_image(
|
|
seed=seed,
|
|
prompt=get_effective_prompt(args),
|
|
controlnet_image_path=args.controlnet_image_path,
|
|
config=Config(
|
|
num_inference_steps=args.steps,
|
|
height=height,
|
|
width=width,
|
|
controlnet_strength=args.controlnet_strength,
|
|
),
|
|
)
|
|
|
|
# 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 _calculate_output_dimensions(args) -> tuple[int, int]:
|
|
"""Calculate output dimensions from args, handling scale factors and safety warnings."""
|
|
# Image.open is lazy/efficient, just need the dimension metadata
|
|
orig_image = PIL.Image.open(args.controlnet_image_path)
|
|
output_width, output_height = orig_image.size
|
|
|
|
if isinstance(args.height, ScaleFactor):
|
|
output_height: int = args.height.get_scaled_value(orig_image.height) # type: ignore
|
|
|
|
else:
|
|
output_height = args.height # type: ignore
|
|
|
|
if isinstance(args.width, ScaleFactor):
|
|
output_width: int = args.width.get_scaled_value(orig_image.width) # type: ignore
|
|
|
|
else:
|
|
output_width = args.width # type: ignore
|
|
|
|
# Check if dimensions exceed safe limits
|
|
total_pixels = output_height * output_width
|
|
|
|
if total_pixels > ui_defaults.MAX_PIXELS_WARNING_THRESHOLD:
|
|
print(
|
|
f"⚠️ WARNING: The requested dimensions {output_width}x{output_height} "
|
|
f"({total_pixels:,} pixels) exceed max recommended ({ui_defaults.MAX_PIXELS_WARNING_THRESHOLD:,} pixels)."
|
|
)
|
|
print("This generation is likely to exceed the capabilities of this computer and may:")
|
|
print(" ⏳ Take a very long time to complete")
|
|
print(" 🔥 Run out of memory")
|
|
print(" 💥 Cause the program and your Mac to crash")
|
|
|
|
user_input = input("\nPress Enter to continue at your own risk, or type 'n' to cancel: ")
|
|
if user_input.lower() in ["n", "no"]:
|
|
print("🛑 Generation cancelled by user.")
|
|
sys.exit(1)
|
|
|
|
return output_width, output_height
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|