Add flux1 pipeline

[kkimmk]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @kkimmk, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly expands the MAX framework's capabilities by introducing a complete image generation workflow through the Flux1 diffusion pipeline. It integrates several advanced deep learning models (VAE, CLIP, T5, and a custom Flux Transformer) and a specialized diffusion scheduler. The changes enable users to generate high-quality images from textual prompts with fine-grained control over various parameters, while also enhancing core utilities for model compilation and image handling. This lays a robust foundation for future generative AI applications within the MAX ecosystem.

Highlights

• New Flux1 Image Generation Pipeline: A comprehensive image generation pipeline, FluxPipeline, has been integrated, leveraging various sub-models including VAE, CLIP, T5, and a custom Flux Transformer model for text-to-image synthesis.
• Dedicated Diffusion Entrypoint and CLI: A new DiffusionPipeline entrypoint is introduced for image generation tasks, complemented by new CLI commands under max.entrypoints.pipelines to facilitate image generation with configurable parameters like prompts, dimensions, and inference steps.
• Core Model Architectures Implemented: Detailed implementations for AutoencoderKL (VAE), CLIPTextModel, and T5EncoderModel are added, complete with their respective sub-modules such as ResnetBlock2D, VAEAttention, CLIPAttention, CLIPMLP, T5LayerNorm, and T5Attention.
• Custom Diffusion Scheduler: A FlowMatchEulerDiscreteScheduler has been added, providing a native Modular implementation of the Euler scheduler, crucial for controlling the denoising process in diffusion models.
• Image Processing and Utility Enhancements: A VaeImageProcessor is introduced for robust image preprocessing and postprocessing, alongside experimental CompileWrapper and max_compile utilities for optimizing graph compilation, and a finfo class for DType numerical properties, including bfloat16 support.
• Configuration System Updates: The configuration system (max.config and max.pipelines.lib.config) has been updated to support the new image generation pipeline, including handling of model_index.json for Diffusers-style repositories and optimizing memory estimation by skipping KV cache profiling for image generation tasks.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This is a large but well-structured pull request that adds support for the FLUX.1 diffusion model. The changes span from new model architecture implementations, a new diffusion pipeline, CLI extensions, to various supporting utilities. The code is generally of high quality. My review focuses on improving robustness, correctness, and maintainability in a few key areas. I've identified a potential bug in the CLIP text encoder's pooling logic, suggested improvements for command-line argument parsing and exception handling, and recommended using pathlib for more robust path manipulation.

[gemini-code-assist]

The special handling for self.eos_token_id == 2 using ops.argmax(input_ids, ...) is brittle and likely incorrect. It assumes that the token with the highest ID in the sequence is the EOS token, which is not a safe assumption. This can lead to incorrect pooling by selecting the wrong token's hidden state.

The logic in the else block, which explicitly finds the eos_token_id, is much more robust. I recommend using this robust logic for all cases to prevent potential bugs.

        eos_token_indices = ops.argmax(
            ops.equal(input_ids, self.eos_token_id).cast(DType.int32),
            axis=-1,
        ).cast(DType.int32)

[gemini-code-assist]

Using type=bool with argparse can lead to unexpected behavior. For instance, python script.py --use-torch-randn False would result in True because bool('False') evaluates to True.

For boolean flags, it's better to use action=argparse.BooleanOptionalAction (for Python 3.9+) which automatically creates --use-torch-randn and --no-use-torch-randn flags. This makes the CLI behavior explicit and less error-prone.

Suggested change

	parser.add_argument("--use-torch-randn", type=bool, default=True)
	parser.add_argument("--use-torch-randn", action=argparse.BooleanOptionalAction, default=True)

[gemini-code-assist]

Catching a generic Exception is too broad and can hide unexpected errors. It's better to catch more specific exceptions that you expect to handle, such as json.JSONDecodeError for parsing issues and IOError for file reading problems.

Suggested change

	except Exception as e:
	raise ValueError(
	f"Failed to load configuration from {config_path}: {e}"
	) from e
	except (json.JSONDecodeError, IOError) as e:
	raise ValueError(
	f"Failed to load configuration from {config_path}: {e}"
	) from e

[gemini-code-assist]

Using weight_path.split("/")[0] for path manipulation is not platform-agnostic and can fail on operating systems that use a different path separator (like Windows). It's more robust to use pathlib for this, as it handles path components correctly across different OS environments.

Suggested change

	if weight_path.split("/")[0] == name
	if Path(weight_path).parts[0] == name