From dbefc209f54c757f42aa32d5a192800130ffa10a Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Wed, 3 Dec 2025 18:34:59 +0000 Subject: [PATCH] Manually handle flux1 inference --- .../worker/engines/mflux/generator/flux1.py | 167 ++++++++++++++++++ .../engines/mflux/generator/generate.py | 2 +- .../worker/engines/mflux/generator/mflux.py | 48 ----- 3 files changed, 168 insertions(+), 49 deletions(-) create mode 100644 src/exo/worker/engines/mflux/generator/flux1.py delete mode 100644 src/exo/worker/engines/mflux/generator/mflux.py diff --git a/src/exo/worker/engines/mflux/generator/flux1.py b/src/exo/worker/engines/mflux/generator/flux1.py new file mode 100644 index 00000000..02445cbe --- /dev/null +++ b/src/exo/worker/engines/mflux/generator/flux1.py @@ -0,0 +1,167 @@ +import mlx.core as mx +from mflux.callbacks.callbacks import Callbacks +from mflux.config.config import Config +from mflux.config.runtime_config import RuntimeConfig +from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator +from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator +from mflux.models.flux.model.flux_text_encoder.prompt_encoder import PromptEncoder +from mflux.models.flux.variants.txt2img.flux import Flux1 +from mflux.utils.array_util import ArrayUtil +from mflux.utils.exceptions import StopImageGenerationException +from mflux.utils.image_util import ImageUtil +from PIL import Image +from tqdm import tqdm + +from exo.shared.types.api import ImageGenerationTaskParams + + +def parse_size(size_str: str | None) -> tuple[int, int]: + """ + Parse size parameter like '1024x1024' to (width, height) tuple. + """ + if not size_str or size_str == "auto": + size_str = "1024x1024" + + try: + parts = size_str.split("x") + if len(parts) == 2: + width, height = int(parts[0]), int(parts[1]) + return (width, height) + except (ValueError, AttributeError): + pass + + # Default fallback + return (1024, 1024) + + +def _generate_image(model: Flux1, settings: Config, prompt: str, seed: int): + # 0. Create a new runtime config based on the model type and input parameters + config = RuntimeConfig(settings, model.model_config) + time_steps = tqdm(range(config.init_time_step, config.num_inference_steps)) + + # 1. Create the initial latents + latents = LatentCreator.create_for_txt2img_or_img2img( + seed=seed, + height=config.height, + width=config.width, + img2img=Img2Img( + vae=model.vae, + latent_creator=FluxLatentCreator, + image_path=config.image_path, + sigmas=config.scheduler.sigmas, + init_time_step=config.init_time_step, + ), + ) + + # 2. Encode the prompt + prompt_embeds, pooled_prompt_embeds = PromptEncoder.encode_prompt( + prompt=prompt, + prompt_cache=model.prompt_cache, + t5_tokenizer=model.t5_tokenizer, + clip_tokenizer=model.clip_tokenizer, + t5_text_encoder=model.t5_text_encoder, + clip_text_encoder=model.clip_text_encoder, + ) + + # (Optional) Call subscribers for beginning of loop + Callbacks.before_loop( + seed=seed, + prompt=prompt, + latents=latents, + config=config, + ) + + for t in time_steps: + try: + # Scale model input if needed by the scheduler + latents = config.scheduler.scale_model_input(latents, t) + + # 3.t Predict the noise + noise = model.transformer( + t=t, + config=config, + hidden_states=latents, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + ) + + # 4.t Take one denoise step + latents = config.scheduler.step( + model_output=noise, + timestep=t, + sample=latents, + ) + + # (Optional) Call subscribers in-loop + Callbacks.in_loop( + t=t, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) + + # (Optional) Evaluate to enable progress tracking + mx.eval(latents) + + except KeyboardInterrupt: # noqa: PERF203 + Callbacks.interruption( + t=t, + seed=seed, + prompt=prompt, + latents=latents, + config=config, + time_steps=time_steps, + ) + raise StopImageGenerationException( + f"Stopping image generation at step {t + 1}/{len(time_steps)}" + ) from None + + # (Optional) Call subscribers after loop + Callbacks.after_loop( + seed=seed, + prompt=prompt, + latents=latents, + config=config, + ) + + # 7. Decode the latent array and return the image + latents = ArrayUtil.unpack_latents( + latents=latents, height=config.height, width=config.width + ) + decoded = model.vae.decode(latents) + return ImageUtil.to_image( + decoded_latents=decoded, + config=config, + seed=seed, + prompt=prompt, + quantization=model.bits, + lora_paths=model.lora_paths, + lora_scales=model.lora_scales, + image_path=config.image_path, + image_strength=config.image_strength, + generation_time=time_steps.format_dict["elapsed"], + ) + + +def generate_image( + model: Flux1, + task: ImageGenerationTaskParams, +) -> Image.Image: + # Parse parameters + width, height = parse_size(task.size) + quality = task.quality or "medium" + + # TODO: Flux1 only + steps = 2 + if quality == "low": + steps = 1 + elif quality == "high": + steps = 4 + + seed = 2 # TODO: not in OAI API? + config = Config(num_inference_steps=steps, height=height, width=width) + + image = _generate_image(model=model, settings=config, prompt=task.prompt, seed=seed) + return image.image diff --git a/src/exo/worker/engines/mflux/generator/generate.py b/src/exo/worker/engines/mflux/generator/generate.py index 76953ade..612f9177 100644 --- a/src/exo/worker/engines/mflux/generator/generate.py +++ b/src/exo/worker/engines/mflux/generator/generate.py @@ -8,7 +8,7 @@ from PIL import Image from exo.shared.types.api import ImageGenerationTaskParams from exo.shared.types.worker.runner_response import ImageGenerationResponse -from exo.worker.engines.mflux.generator.mflux import generate_image +from exo.worker.engines.mflux.generator.flux1 import generate_image image_generation_stream = mx.new_stream(mx.default_device()) diff --git a/src/exo/worker/engines/mflux/generator/mflux.py b/src/exo/worker/engines/mflux/generator/mflux.py deleted file mode 100644 index 4f978b85..00000000 --- a/src/exo/worker/engines/mflux/generator/mflux.py +++ /dev/null @@ -1,48 +0,0 @@ -from mflux.config.config import Config -from mflux.models.flux.variants.txt2img.flux import Flux1 -from PIL import Image - -from exo.shared.types.api import ImageGenerationTaskParams - - -def parse_size(size_str: str | None) -> tuple[int, int]: - """ - Parse size parameter like '1024x1024' to (width, height) tuple. - """ - if not size_str or size_str == "auto": - size_str = "1024x1024" - - try: - parts = size_str.split("x") - if len(parts) == 2: - width, height = int(parts[0]), int(parts[1]) - return (width, height) - except (ValueError, AttributeError): - pass - - # Default fallback - return (1024, 1024) - - -def generate_image( - model: Flux1, - task: ImageGenerationTaskParams, -) -> Image.Image: - # Parse parameters - width, height = parse_size(task.size) - quality = task.quality or "medium" - - # TODO: Flux1 only - steps = 2 - if quality == "low": - steps = 1 - elif quality == "high": - steps = 4 - - image = model.generate_image( - seed=2, - prompt=task.prompt, - config=Config(num_inference_steps=steps, height=height, width=width), - ) - - return image.image