Add FLUX.1 export and runtime support

export/huggingface/diffusers/export.py

@@ -47,12 +47,19 @@ def build_diffusion_metadata(
             behavior.default_cfg_scale if model_info.supports_cfg else 1.0
         ),
         default_guidance_scale=(
-            behavior.default_guidance_scale
-            if model_info.accepts_guidance_input
-            else None
+            model_info.default_guidance_scale
+            if model_info.default_guidance_scale is not None
+            else (
+                behavior.default_guidance_scale
+                if model_info.accepts_guidance_input
+                else None
+            )
         ),
         guidance_input_dim=model_info.guidance_input_dim,
-        default_inference_steps=behavior.default_inference_steps,
+        # Model specific defaults (e.g. Schnell vs Dev) override the pipeline default.
+        default_inference_steps=(
+            model_info.default_inference_steps or behavior.default_inference_steps
+        ),
         text_encoder_dtype=str(config.model.text_encoder_dtype).removeprefix("torch."),
         denoiser_dtype=str(config.model.denoiser_dtype).removeprefix("torch."),
         vae_dtype=str(config.model.vae_dtype).removeprefix("torch."),
@@ -72,7 +79,7 @@ def _resolve_max_sequence_lengths(
     export_spec: DiffusionPipelineExportSpec,
     override: int | None,
 ) -> tuple[int, ...]:
-    lengths = list(export_spec.default_max_sequence_lengths)
+    lengths = list(export_spec.default_max_text_sequence_lengths)
     if override is None:
         return tuple(lengths)
     if override <= 0:
@@ -121,7 +128,9 @@ def export_programs(
     with torch.no_grad():
         return {
             name: torch.export.export(
-                modules[name], inputs[name], dynamic_shapes=dynamic_shapes[name]
+                modules[name],
+                inputs[name],
+                dynamic_shapes=dynamic_shapes[name],
             )
             for name in modules
         }

export/huggingface/diffusers/pipeline_discovery.py

@@ -60,7 +60,7 @@ def pipeline_spec_from_model_index(
             f"{pipeline_name} requires text encoder components "
             f"{defaults.text_encoder_slots}, found {text_encoder_slots}"
         )
-    if len(export_spec.default_max_sequence_lengths) != len(text_encoder_slots):
+    if len(export_spec.default_max_text_sequence_lengths) != len(text_encoder_slots):
         raise ValueError(
             f"{pipeline_name} must define one maximum sequence length per "
             "text encoder"
@@ -103,7 +103,7 @@ def resolve_huggingface_pipeline(
         model_info=DiffusionModelInfo(
             scheduler_name=scheduler_name,
             scheduler_prediction_type=prediction_type,
-            max_sequence_lengths=export_spec.default_max_sequence_lengths,
+            max_sequence_lengths=export_spec.default_max_text_sequence_lengths,
             zero_negative_prompt=zero_negative_prompt,
         ),
     )

export/huggingface/diffusers/pipeline_specs/__init__.py

@@ -12,10 +12,12 @@
     DiffusionPipelineExportSpec,
     ResolvedDiffusionPipeline,
 )
+from .flux import FluxPipelineExportSpec
 from .sd15 import SD15PipelineExportSpec
 from .sdxl import SDXLPipelineExportSpec
 
 HF_PIPELINE_EXPORT_SPECS: dict[str, type[DiffusionPipelineExportSpec]] = {
+    "FluxPipeline": FluxPipelineExportSpec,
     "StableDiffusionPipeline": SD15PipelineExportSpec,
     "StableDiffusionXLPipeline": SDXLPipelineExportSpec,
 }
@@ -25,6 +27,7 @@
     "DiffusionModelInfo",
     "DiffusionPipelineExportSpec",
     "HF_PIPELINE_EXPORT_SPECS",
+    "FluxPipelineExportSpec",
     "ResolvedDiffusionPipeline",
     "SD15PipelineExportSpec",
     "SDXLPipelineExportSpec",

export/huggingface/diffusers/pipeline_specs/base.py

@@ -40,7 +40,7 @@ class DiffusionPipelineExportSpec(ABC):
     """Hugging Face loading and graph-export contract for one pipeline."""
 
     pipeline_behavior: DiffusionPipelineBehavior
-    default_max_sequence_lengths: tuple[int, ...]
+    default_max_text_sequence_lengths: tuple[int, ...]
     max_sequence_length_override_index: int | None
     default_size: tuple[int, int]
 
@@ -101,6 +101,8 @@ class DiffusionModelInfo:
     supports_cfg: bool = True
     accepts_guidance_input: bool = False
     guidance_input_dim: int = 0
+    default_inference_steps: int | None = None
+    default_guidance_scale: float | None = None
 
     def __post_init__(self) -> None:
         """Validate discovered model facts."""

export/huggingface/diffusers/pipeline_specs/flux.py

@@ -0,0 +1,221 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Hugging Face export specification for FLUX.1."""
+
+from dataclasses import replace
+from typing import Any
+
+import torch
+
+from executorch.extension.diffusers.pipeline_behaviors import FluxPipelineBehavior
+
+from ..wrappers import (
+    CLIPPooledTextEncoder,
+    FluxDenoiser,
+    freeze_flux_pos_embed,
+    T5TextEncoder,
+    vae_scale_factor,
+    VAEDecoder,
+)
+from .base import DiffusionComponents, DiffusionPipelineExportSpec
+
+
+class FluxPipelineExportSpec(DiffusionPipelineExportSpec):
+    """Compose FLUX text encoders, transformer, and VAE for export."""
+
+    pipeline_behavior = FluxPipelineBehavior()
+    default_max_text_sequence_lengths = (77, 512)
+    max_sequence_length_override_index = 1
+    default_size = (1024, 1024)
+
+    def _validate(self, components, config) -> int:
+        if (
+            len(config.model.max_sequence_lengths) != 2
+            or config.model.max_sequence_lengths[0] != 77
+        ):
+            raise ValueError("FLUX requires CLIP length 77 and one T5 sequence length")
+        scale = vae_scale_factor(components.require("vae"))
+        spatial_multiple = scale * 2
+        if (
+            config.model.height % spatial_multiple
+            or config.model.width % spatial_multiple
+        ):
+            raise ValueError(
+                f"FLUX height and width must be multiples of {spatial_multiple}"
+            )
+        if config.export.dynamic_shapes and (
+            config.export.min_size % spatial_multiple
+            or config.export.max_size % spatial_multiple
+        ):
+            raise ValueError(
+                f"FLUX dynamic size bounds must be multiples of {spatial_multiple}"
+            )
+        return scale
+
+    def wrap_components(self, components, config):
+        scale = self._validate(components, config)
+        transformer = components.require("denoiser")
+        vae = components.require("vae")
+        freeze_flux_pos_embed(
+            transformer,
+            max_positions=(
+                config.export.max_size
+                if config.export.dynamic_shapes
+                else max(config.model.height, config.model.width)
+            )
+            // scale
+            // 2,
+        )
+        return {
+            "text_encoder_0": CLIPPooledTextEncoder(components.require("text_encoder")),
+            "text_encoder_1": T5TextEncoder(components.require("text_encoder_2")),
+            "denoise": FluxDenoiser(
+                transformer, bool(getattr(transformer.config, "guidance_embeds", False))
+            ),
+            "decode": VAEDecoder(vae, config.model.vae_dtype),
+        }
+
+    def example_inputs(self, components, config):
+        scale = self._validate(components, config)
+        clip_ids = torch.zeros(1, 77, dtype=torch.long)
+        t5_ids = torch.zeros(1, config.model.max_sequence_lengths[1], dtype=torch.long)
+        with torch.no_grad():
+            pooled = CLIPPooledTextEncoder(components.require("text_encoder"))(
+                clip_ids
+            ).to(config.model.denoiser_dtype)
+            prompt = T5TextEncoder(components.require("text_encoder_2"))(t5_ids).to(
+                config.model.denoiser_dtype
+            )
+        transformer = components.require("denoiser")
+        packed_channels = int(transformer.config.in_channels)
+        sequence_length = (config.model.height // scale // 2) * (
+            config.model.width // scale // 2
+        )
+        latents = torch.zeros(
+            1, sequence_length, packed_channels, dtype=config.model.denoiser_dtype
+        )
+        inputs = self.pipeline_behavior.prepare_denoiser_inputs(
+            latents,
+            torch.zeros(1),
+            (prompt, pooled),
+            height=config.model.height,
+            width=config.model.width,
+            vae_scale_factor=scale,
+            dtype=config.model.denoiser_dtype,
+            guidance_scale=(
+                self.pipeline_behavior.default_guidance_scale
+                if getattr(transformer.config, "guidance_embeds", False)
+                else None
+            ),
+        )
+        vae_channels = packed_channels // 4
+        return {
+            "text_encoder_0": (clip_ids,),
+            "text_encoder_1": (t5_ids,),
+            "denoise": inputs,
+            "decode": (
+                torch.zeros(
+                    1,
+                    vae_channels,
+                    config.model.height // scale,
+                    config.model.width // scale,
+                    dtype=config.model.denoiser_dtype,
+                ),
+            ),
+        }
+
+    def dynamic_shapes(self, components, config):
+        scale = self._validate(components, config)
+        modules = self.wrap_components(components, config)
+        if not config.export.dynamic_shapes:
+            return {name: None for name in modules}
+        packed_multiple = scale * 2
+        image_tokens = torch.export.Dim(
+            "image_tokens",
+            min=(config.export.min_size // packed_multiple) ** 2,
+            max=(config.export.max_size // packed_multiple) ** 2,
+        )
+        latent_height = torch.export.Dim(
+            "latent_height",
+            min=config.export.min_size // scale,
+            max=config.export.max_size // scale,
+        )
+        latent_width = torch.export.Dim(
+            "latent_width",
+            min=config.export.min_size // scale,
+            max=config.export.max_size // scale,
+        )
+        denoise_shapes = (
+            {1: image_tokens},
+            None,
+            None,
+            None,
+            None,
+            {0: image_tokens},
+        )
+        if getattr(components.require("denoiser").config, "guidance_embeds", False):
+            denoise_shapes = (*denoise_shapes, None)
+        return {
+            "text_encoder_0": None,
+            "text_encoder_1": None,
+            "denoise": denoise_shapes,
+            "decode": ({2: latent_height, 3: latent_width},),
+        }
+
+    def resolve_model_info(self, components, model_info):
+        scale = vae_scale_factor(components.require("vae"))
+        guidance_embeds = bool(
+            getattr(components.require("denoiser").config, "guidance_embeds", False)
+        )
+        return replace(
+            model_info,
+            vae_scale_factor=scale,
+            dynamic_spatial_multiple=scale * 2,
+            supports_cfg=True,
+            accepts_guidance_input=guidance_embeds,
+            guidance_input_dim=1 if guidance_embeds else 0,
+            default_inference_steps=28 if guidance_embeds else 4,
+            default_guidance_scale=3.5 if guidance_embeds else 0.0,
+        )
+
+    def load_components(self, model_id, revision, config):
+        try:
+            from diffusers import AutoencoderKL, FluxTransformer2DModel
+            from transformers import CLIPTextModel, T5EncoderModel
+        except ImportError as error:
+            raise ImportError(
+                "Loading FLUX requires diffusers and transformers"
+            ) from error
+        common: dict[str, Any] = {"revision": revision}
+        return DiffusionComponents(
+            {
+                "text_encoder": CLIPTextModel.from_pretrained(
+                    model_id,
+                    subfolder="text_encoder",
+                    torch_dtype=config.model.text_encoder_dtype,
+                    **common,
+                ).eval(),
+                "text_encoder_2": T5EncoderModel.from_pretrained(
+                    model_id,
+                    subfolder="text_encoder_2",
+                    torch_dtype=config.model.text_encoder_dtype,
+                    **common,
+                ).eval(),
+                "denoiser": FluxTransformer2DModel.from_pretrained(
+                    model_id,
+                    subfolder="transformer",
+                    torch_dtype=config.model.denoiser_dtype,
+                    **common,
+                ).eval(),
+                "vae": AutoencoderKL.from_pretrained(
+                    model_id,
+                    subfolder="vae",
+                    torch_dtype=config.model.vae_dtype,
+                    **common,
+                ).eval(),
+            }
+        )

export/huggingface/diffusers/pipeline_specs/sd15.py

@@ -29,7 +29,7 @@ class SD15PipelineExportSpec(StableDiffusionPipelineExportSpec):
     """Compose SD1.5 components into exported PTE methods."""
 
     pipeline_behavior = SD15PipelineBehavior()
-    default_max_sequence_lengths = (77,)
+    default_max_text_sequence_lengths = (77,)
     max_sequence_length_override_index = None
     default_size = (512, 512)
 
@@ -39,7 +39,7 @@ def _validate(
         scale = vae_scale_factor(components.require("vae"))
         if config.model.height % scale or config.model.width % scale:
             raise ValueError(f"height and width must be multiples of {scale}")
-        if config.model.max_sequence_lengths != self.default_max_sequence_lengths:
+        if config.model.max_sequence_lengths != self.default_max_text_sequence_lengths:
             raise ValueError("SD 1.5 requires a sequence length of 77")
         return scale
 
@@ -86,6 +86,7 @@ def example_inputs(
                 (prompt,),
                 height=config.model.height,
                 width=config.model.width,
+                vae_scale_factor=scale,
                 dtype=config.model.denoiser_dtype,
                 guidance_scale=(
                     self.pipeline_behavior.default_guidance_scale

export/huggingface/diffusers/pipeline_specs/sdxl.py

@@ -29,7 +29,7 @@ class SDXLPipelineExportSpec(StableDiffusionPipelineExportSpec):
     """Compose SDXL components into exported PTE methods."""
 
     pipeline_behavior = SDXLPipelineBehavior()
-    default_max_sequence_lengths = (77, 77)
+    default_max_text_sequence_lengths = (77, 77)
     max_sequence_length_override_index = None
     default_size = (1024, 1024)
 
@@ -39,7 +39,7 @@ def _validate(
         scale = vae_scale_factor(components.require("vae"))
         if config.model.height % scale or config.model.width % scale:
             raise ValueError(f"height and width must be multiples of {scale}")
-        if config.model.max_sequence_lengths != self.default_max_sequence_lengths:
+        if config.model.max_sequence_lengths != self.default_max_text_sequence_lengths:
             raise ValueError("SDXL requires sequence lengths of (77, 77)")
         return scale
 
@@ -94,6 +94,7 @@ def example_inputs(
                 (prompt, pooled),
                 height=config.model.height,
                 width=config.model.width,
+                vae_scale_factor=scale,
                 dtype=config.model.denoiser_dtype,
                 guidance_scale=(
                     self.pipeline_behavior.default_guidance_scale