diff --git a/export/huggingface/diffusers/export.py b/export/huggingface/diffusers/export.py index 65118a53471..77948576f9f 100644 --- a/export/huggingface/diffusers/export.py +++ b/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 } diff --git a/export/huggingface/diffusers/pipeline_discovery.py b/export/huggingface/diffusers/pipeline_discovery.py index 0680f713ce9..21d61f6aeec 100644 --- a/export/huggingface/diffusers/pipeline_discovery.py +++ b/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, ), ) diff --git a/export/huggingface/diffusers/pipeline_specs/__init__.py b/export/huggingface/diffusers/pipeline_specs/__init__.py index d765b95c426..667f7d6df11 100644 --- a/export/huggingface/diffusers/pipeline_specs/__init__.py +++ b/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", diff --git a/export/huggingface/diffusers/pipeline_specs/base.py b/export/huggingface/diffusers/pipeline_specs/base.py index aafd9202555..864e2cdbdec 100644 --- a/export/huggingface/diffusers/pipeline_specs/base.py +++ b/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.""" diff --git a/export/huggingface/diffusers/pipeline_specs/flux.py b/export/huggingface/diffusers/pipeline_specs/flux.py new file mode 100644 index 00000000000..629d2eb3490 --- /dev/null +++ b/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(), + } + ) diff --git a/export/huggingface/diffusers/pipeline_specs/sd15.py b/export/huggingface/diffusers/pipeline_specs/sd15.py index 12fe78b4b43..57fa3fed75b 100644 --- a/export/huggingface/diffusers/pipeline_specs/sd15.py +++ b/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 diff --git a/export/huggingface/diffusers/pipeline_specs/sdxl.py b/export/huggingface/diffusers/pipeline_specs/sdxl.py index f8107ddc30b..82776103a4e 100644 --- a/export/huggingface/diffusers/pipeline_specs/sdxl.py +++ b/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 diff --git a/export/huggingface/diffusers/wrappers.py b/export/huggingface/diffusers/wrappers.py index 2d90e24ea0d..6893c8567fe 100644 --- a/export/huggingface/diffusers/wrappers.py +++ b/export/huggingface/diffusers/wrappers.py @@ -65,6 +65,121 @@ def forward(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model(input_ids).last_hidden_state +class CLIPPooledTextEncoder(torch.nn.Module): + """Return the pooled CLIP representation used by FLUX.""" + + def __init__(self, model: torch.nn.Module) -> None: + super().__init__() + self.model = model + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + output = self.model(input_ids) + pooled = getattr(output, "pooler_output", None) + return output[1] if pooled is None else pooled + + +class T5TextEncoder(torch.nn.Module): + """Return T5's final hidden state.""" + + def __init__(self, model: torch.nn.Module) -> None: + super().__init__() + self.model = model + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model(input_ids).last_hidden_state + + +class BoundedFluxPosEmbed(torch.nn.Module): + """Select precomputed FLUX RoPE values for static or bounded dynamic IDs.""" + + def __init__( + self, + position_embedding: torch.nn.Module, + max_positions: int, + ) -> None: + super().__init__() + self.position_embedding = position_embedding + # FLUX.1 axes are (type, row, column); type is always 0 for text-to-image. + axes_dims = tuple(position_embedding.axes_dim) + ids = ( + torch.arange(max_positions) + .reshape(max_positions, 1) + .expand(max_positions, len(axes_dims)) + ) + with torch.no_grad(): + cos, sin = position_embedding(ids) + for index, (axis_cos, axis_sin) in enumerate( + zip(cos.split(axes_dims, dim=-1), sin.split(axes_dims, dim=-1)) + ): + self.register_buffer(f"cos_{index}", axis_cos, persistent=True) + self.register_buffer(f"sin_{index}", axis_sin, persistent=True) + self.axes = len(axes_dims) + + def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + indices = ids.to(torch.int64) + cos = torch.cat( + [ + getattr(self, f"cos_{axis}")[indices[:, axis]] + for axis in range(self.axes) + ], + dim=-1, + ) + sin = torch.cat( + [ + getattr(self, f"sin_{axis}")[indices[:, axis]] + for axis in range(self.axes) + ], + dim=-1, + ) + return cos, sin + + +def freeze_flux_pos_embed( + model: torch.nn.Module, + *, + max_positions: int, +) -> None: + """Freeze Diffusers' float64 FLUX RoPE construction for export.""" + position_embedding = getattr(model, "pos_embed", None) + if position_embedding is None: + return + if isinstance(position_embedding, BoundedFluxPosEmbed): + if position_embedding.cos_0.shape[0] >= max_positions: + return + position_embedding = position_embedding.position_embedding + model.pos_embed = BoundedFluxPosEmbed(position_embedding, max_positions) + + +class FluxDenoiser(torch.nn.Module): + """Flatten the keyword-heavy Diffusers FLUX transformer interface.""" + + def __init__(self, model: torch.nn.Module, guidance_embeds: bool) -> None: + super().__init__() + self.model = model + self.guidance_embeds = guidance_embeds + + def forward( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + prompt_embeddings: torch.Tensor, + pooled_embedding: torch.Tensor, + text_ids: torch.Tensor, + image_ids: torch.Tensor, + guidance: torch.Tensor | None = None, + ) -> torch.Tensor: + return self.model( + hidden_states=latents, + timestep=timestep, + guidance=guidance if self.guidance_embeds else None, + pooled_projections=pooled_embedding, + encoder_hidden_states=prompt_embeddings, + txt_ids=text_ids, + img_ids=image_ids, + return_dict=False, + )[0] + + class SD15Denoiser(torch.nn.Module): def __init__(self, model: torch.nn.Module) -> None: super().__init__() diff --git a/export/tests/test_diffusers.py b/export/tests/test_diffusers.py index 2f34eedcbce..4cc16babd73 100644 --- a/export/tests/test_diffusers.py +++ b/export/tests/test_diffusers.py @@ -33,11 +33,16 @@ from executorch.export.huggingface.diffusers.pipeline_specs import ( DiffusionComponents, DiffusionModelInfo, + FluxPipelineExportSpec, ResolvedDiffusionPipeline, SD15PipelineExportSpec, SDXLPipelineExportSpec, ) -from executorch.export.huggingface.diffusers.wrappers import _guidance_scale_embedding +from executorch.export.huggingface.diffusers.wrappers import ( + _guidance_scale_embedding, + BoundedFluxPosEmbed, + freeze_flux_pos_embed, +) from executorch.extension.diffusers.pipeline_behaviors import SD15PipelineBehavior @@ -89,6 +94,16 @@ def decode(self, latents): return SimpleNamespace(sample=latents[:, :3]) +class FluxTransformer(torch.nn.Module): + def __init__(self, guidance_embeds: bool) -> None: + super().__init__() + self.config = SimpleNamespace(in_channels=16, guidance_embeds=guidance_embeds) + + def forward(self, hidden_states, **kwargs): + image_ids = kwargs["img_ids"][:, :1].unsqueeze(0) + return (hidden_states + image_ids * 0,) + + def config( height: int, width: int, @@ -149,7 +164,7 @@ def setUp(self): def test_sequence_length_override_changes_only_supported_encoder(self): spec = SimpleNamespace( name="sd3", - default_max_sequence_lengths=(77, 77, 256), + default_max_text_sequence_lengths=(77, 77, 256), max_sequence_length_override_index=2, ) self.assertEqual( @@ -369,6 +384,147 @@ def step(self, prediction, timestep, latents, return_dict=False): self.assertEqual(image.shape, (1, 3, 1, 1)) torch.testing.assert_close(image, torch.full_like(image, 0.5)) + def test_portable_flux_pte_runs_through_diffusion_runner(self): + from executorch.extension.diffusers.runner import DiffusionRunner + from executorch.runtime import Runtime, Verification + + class Tokenizer: + def __init__(self, length): + self.model_max_length = length + + def __call__(self, prompt, *, max_length, **kwargs): + return SimpleNamespace( + input_ids=torch.zeros(1, max_length, dtype=torch.long) + ) + + class Scheduler: + config = SimpleNamespace( + prediction_type=None, + base_image_seq_len=1, + max_image_seq_len=4, + base_shift=0.5, + max_shift=1.0, + ) + + def set_timesteps(self, steps, *, sigmas, mu): + self.timesteps = torch.arange(steps - 1, -1, -1) + + def step(self, prediction, timestep, latents, return_dict=False): + return (latents - prediction,) + + components = DiffusionComponents( + { + "text_encoder": TextEncoder(3), + "text_encoder_2": TextEncoder(5), + "denoiser": FluxTransformer(False), + "vae": VAE(), + } + ) + export_config = config(16, 16, (77, 8)) + export_spec = FluxPipelineExportSpec() + model_info = export_spec.resolve_model_info( + components, + DiffusionModelInfo( + scheduler_name="FakeScheduler", + max_sequence_lengths=(77, 8), + zero_negative_prompt=False, + ), + ) + pipeline = ResolvedDiffusionPipeline(export_spec, model_info) + edge_program = lower_programs( + export_programs(components, export_spec, export_config), + pipeline, + export_config, + ) + + with TemporaryDirectory() as directory: + output = Path(directory) / "tiny_flux.pte" + write_pte(edge_program, output) + program = Runtime.get().load_program( + str(output), verification=Verification.Minimal + ) + image = DiffusionRunner( + program, [Tokenizer(77), Tokenizer(8)], Scheduler() + ).generate("cat", steps=1, seed=9) + del program + + self.assertEqual(image.shape, (1, 3, 2, 2)) + torch.testing.assert_close(image, torch.full_like(image, 0.5)) + + def test_dynamic_portable_flux_pte_runs_at_a_different_size(self): + from executorch.extension.diffusers.runner import DiffusionRunner + from executorch.runtime import Runtime, Verification + + class Tokenizer: + def __init__(self, length): + self.model_max_length = length + + def __call__(self, prompt, *, max_length, **kwargs): + return SimpleNamespace( + input_ids=torch.zeros(1, max_length, dtype=torch.long) + ) + + class Scheduler: + config = SimpleNamespace( + prediction_type=None, + base_image_seq_len=1, + max_image_seq_len=64, + base_shift=0.5, + max_shift=1.0, + ) + + def set_timesteps(self, steps, *, sigmas, mu): + self.timesteps = torch.arange(steps - 1, -1, -1) + + def step(self, prediction, timestep, latents, return_dict=False): + return (latents - prediction,) + + components = DiffusionComponents( + { + "text_encoder": TextEncoder(3), + "text_encoder_2": TextEncoder(5), + "denoiser": FluxTransformer(False), + "vae": VAE(), + } + ) + export_config = config( + 64, + 64, + (77, 8), + dynamic_shapes=True, + min_size=32, + max_size=96, + ) + export_spec = FluxPipelineExportSpec() + model_info = export_spec.resolve_model_info( + components, + DiffusionModelInfo( + scheduler_name="FakeScheduler", + max_sequence_lengths=(77, 8), + zero_negative_prompt=False, + ), + ) + pipeline = ResolvedDiffusionPipeline(export_spec, model_info) + edge_program = lower_programs( + export_programs(components, export_spec, export_config), + pipeline, + export_config, + ) + + with TemporaryDirectory() as directory: + output = Path(directory) / "dynamic_flux.pte" + write_pte(edge_program, output) + program = Runtime.get().load_program( + str(output), verification=Verification.Minimal + ) + image = DiffusionRunner( + program, [Tokenizer(77), Tokenizer(8)], Scheduler() + ).generate("cat", steps=1, height=32, width=48, seed=9) + del program + + self.assertEqual(image.shape, (1, 3, 4, 6)) + torch.testing.assert_close(image, torch.full_like(image, 0.5)) + def test_sd15_exports_dynamic_spatial_shapes(self): components = DiffusionComponents( { @@ -421,6 +577,239 @@ def test_pipeline_specs_are_discovered(self): self.assertIsInstance(sd15, SD15PipelineExportSpec) self.assertIsInstance(sdxl, SDXLPipelineExportSpec) + flux = pipeline_spec_from_model_index( + { + "_class_name": "FluxPipeline", + "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"], + "tokenizer": ["transformers", "CLIPTokenizer"], + "tokenizer_2": ["transformers", "T5TokenizerFast"], + "text_encoder": ["transformers", "CLIPTextModel"], + "text_encoder_2": ["transformers", "T5EncoderModel"], + "transformer": ["diffusers", "FluxTransformer2DModel"], + "vae": ["diffusers", "AutoencoderKL"], + } + ) + self.assertIsInstance(flux, FluxPipelineExportSpec) + + def test_flux_examples_use_packed_latents_and_position_ids(self): + components = DiffusionComponents( + { + "text_encoder": TextEncoder(3), + "text_encoder_2": TextEncoder(5), + "denoiser": FluxTransformer(True), + "vae": VAE(), + } + ) + inputs = FluxPipelineExportSpec().example_inputs( + components, config(64, 80, (77, 128)) + )["denoise"] + self.assertEqual(inputs[0].shape, (1, 20, 16)) + self.assertEqual(inputs[2].shape, (1, 128, 5)) + self.assertEqual(inputs[4].shape, (128, 3)) + self.assertEqual(inputs[5].shape, (20, 3)) + self.assertEqual(len(inputs), 7) + + def test_flux_resolves_schnell_and_dev_defaults(self): + def resolve(guidance_embeds): + components = DiffusionComponents( + { + "denoiser": FluxTransformer(guidance_embeds), + "vae": VAE(), + } + ) + return FluxPipelineExportSpec().resolve_model_info( + components, + DiffusionModelInfo( + scheduler_name="FlowMatchEulerDiscreteScheduler", + max_sequence_lengths=(77, 512), + zero_negative_prompt=False, + ), + ) + + schnell, dev = resolve(False), resolve(True) + self.assertEqual( + ( + schnell.default_inference_steps, + schnell.accepts_guidance_input, + schnell.default_guidance_scale, + ), + (4, False, 0.0), + ) + self.assertEqual( + ( + dev.default_inference_steps, + dev.accepts_guidance_input, + dev.default_guidance_scale, + ), + (28, True, 3.5), + ) + + def test_flux_rejects_invalid_static_and_dynamic_sizes(self): + components = DiffusionComponents( + { + "denoiser": FluxTransformer(False), + "vae": VAE(), + } + ) + with self.assertRaisesRegex(ValueError, "multiples of 16"): + FluxPipelineExportSpec().wrap_components( + components, config(64, 72, (77, 512)) + ) + with self.assertRaisesRegex(ValueError, "dynamic size bounds"): + FluxPipelineExportSpec().wrap_components( + components, + config( + 64, + 64, + (77, 512), + dynamic_shapes=True, + min_size=48, + max_size=100, + ), + ) + + def test_flux_exports_dynamic_image_tokens_and_decode_shape(self): + components = DiffusionComponents( + { + "text_encoder": TextEncoder(3), + "text_encoder_2": TextEncoder(5), + "denoiser": FluxTransformer(False), + "vae": VAE(), + } + ) + export_config = config( + 64, + 64, + (77, 8), + dynamic_shapes=True, + min_size=32, + max_size=96, + ) + programs = export_programs(components, FluxPipelineExportSpec(), export_config) + denoise = programs["denoise"].module() + output = denoise( + torch.zeros(1, 15, 16), + torch.zeros(1), + torch.zeros(1, 8, 5), + torch.zeros(1, 3), + torch.zeros(8, 3), + torch.zeros(15, 3), + ) + self.assertEqual(output.shape, (1, 15, 16)) + decoded = programs["decode"].module()(torch.zeros(1, 4, 6, 10)) + self.assertEqual(decoded.shape, (1, 3, 6, 10)) + + def test_flux_position_embeddings_are_frozen_without_losing_precision(self): + class PositionEmbedding(torch.nn.Module): + axes_dim = (2, 2, 2) + + def forward(self, ids): + frequencies = torch.arange(2, dtype=torch.float64) + angles = torch.cat( + [ + ids[:, axis : axis + 1].double() * frequencies + for axis in range(3) + ], + dim=-1, + ) + return angles.cos().float(), angles.sin().float() + + model = torch.nn.Module() + model.pos_embed = PositionEmbedding() + ids = torch.tensor([[0.0, 0.0, 0.0], [0.0, 1.0, 3.0]]) + expected = model.pos_embed(ids) + + freeze_flux_pos_embed(model, max_positions=4) + + self.assertIsInstance(model.pos_embed, BoundedFluxPosEmbed) + actual = model.pos_embed(ids) + torch.testing.assert_close(actual[0], expected[0]) + torch.testing.assert_close(actual[1], expected[1]) + self.assertEqual(actual[0].dtype, torch.float32) + + def test_flux_rope_tables_cover_static_and_dynamic_spatial_bounds(self): + class PositionEmbedding(torch.nn.Module): + axes_dim = (2, 2, 2) + + def forward(self, ids): + values = ids.repeat_interleave(2, dim=-1).float() + return values, -values + + for dynamic, bound in ((False, 5), (True, 6)): + with self.subTest(dynamic=dynamic): + transformer = FluxTransformer(False) + transformer.pos_embed = PositionEmbedding() + components = DiffusionComponents( + { + "text_encoder": TextEncoder(3), + "text_encoder_2": TextEncoder(5), + "denoiser": transformer, + "vae": VAE(), + } + ) + FluxPipelineExportSpec().wrap_components( + components, + config( + 64, + 80, + (77, 128), + dynamic_shapes=dynamic, + min_size=32, + max_size=96, + ), + ) + self.assertIsInstance(transformer.pos_embed, BoundedFluxPosEmbed) + self.assertEqual(transformer.pos_embed.cos_2.shape, (bound, 2)) + ids = torch.tensor([[0.0, 3.0, bound - 1]]) + actual = transformer.pos_embed(ids) + expected = PositionEmbedding()(ids) + torch.testing.assert_close(actual, expected) + + def test_bounded_flux_position_embeddings_gather_dynamic_ids(self): + class PositionEmbedding(torch.nn.Module): + axes_dim = (2, 2, 2) + + def forward(self, ids): + parts = [ids[:, axis : axis + 1].repeat(1, 2) for axis in range(3)] + values = torch.cat(parts, dim=-1).float() + return values, -values + + embedding = BoundedFluxPosEmbed(PositionEmbedding(), 8) + ids = torch.tensor([[0.0, 2.0, 3.0], [0.0, 4.0, 1.0]]) + cos, sin = embedding(ids) + expected = torch.tensor( + [[0.0, 0.0, 2.0, 2.0, 3.0, 3.0], [0.0, 0.0, 4.0, 4.0, 1.0, 1.0]] + ) + torch.testing.assert_close(cos, expected) + torch.testing.assert_close(sin, -expected) + + def test_flux_rope_tables_preserve_coordinates_above_bfloat16_precision(self): + class PositionEmbedding(torch.nn.Module): + axes_dim = (2, 2, 2) + + def forward(self, ids): + values = ids.repeat_interleave(2, dim=-1).float() + return values, -values + + embedding = BoundedFluxPosEmbed(PositionEmbedding(), 260) + ids = torch.tensor([[0, 257, 259]]) + torch.testing.assert_close(embedding(ids), PositionEmbedding()(ids)) + + def test_flux_rope_tables_expand_when_reusing_components(self): + class PositionEmbedding(torch.nn.Module): + axes_dim = (2, 2, 2) + + def forward(self, ids): + values = ids.repeat_interleave(2, dim=-1).float() + return values, -values + + model = torch.nn.Module() + model.pos_embed = PositionEmbedding() + freeze_flux_pos_embed(model, max_positions=2) + freeze_flux_pos_embed(model, max_positions=4) + ids = torch.tensor([[0, 1, 3]]) + torch.testing.assert_close(model.pos_embed(ids), PositionEmbedding()(ids)) + def test_model_info_is_discovered_from_hugging_face(self): pipeline_config = { "_class_name": "StableDiffusionXLPipeline", diff --git a/extension/diffusers/README.md b/extension/diffusers/README.md index 4b3243f963a..e7da8353d82 100644 --- a/extension/diffusers/README.md +++ b/extension/diffusers/README.md @@ -1,9 +1,10 @@ # Diffusers This directory provides experimental export and Python runtime support for -Hugging Face checkpoints using `StableDiffusionPipeline` or -`StableDiffusionXLPipeline`. The exported PTE stores the text encoder, denoiser -and VAE methods together with the metadata needed by the runner. +Hugging Face checkpoints using `StableDiffusionPipeline`, +`StableDiffusionXLPipeline`, or `FluxPipeline`. The exported PTE stores +the text encoder, denoiser and VAE methods together with the metadata needed +by the runner. All Python APIs in this package are experimental and may change without notice. @@ -31,6 +32,9 @@ python -m executorch.export.huggingface.diffusers.export \ Use `stabilityai/stable-diffusion-xl-base-1.0` to export SDXL. +Use `black-forest-labs/FLUX.1-schnell` or `black-forest-labs/FLUX.1-dev` to +export FLUX. + ### Run ```bash diff --git a/extension/diffusers/pipeline_behaviors/__init__.py b/extension/diffusers/pipeline_behaviors/__init__.py index 52b79dda512..3e5f312831c 100644 --- a/extension/diffusers/pipeline_behaviors/__init__.py +++ b/extension/diffusers/pipeline_behaviors/__init__.py @@ -7,16 +7,19 @@ """Supported diffusion pipeline specifications and registry.""" from .base import DiffusionPipelineBehavior +from .flux import FluxPipelineBehavior from .sd15 import SD15PipelineBehavior from .sdxl import SDXLPipelineBehavior PIPELINE_BEHAVIORS: dict[str, DiffusionPipelineBehavior] = { + FluxPipelineBehavior.name: FluxPipelineBehavior(), SDXLPipelineBehavior.name: SDXLPipelineBehavior(), SD15PipelineBehavior.name: SD15PipelineBehavior(), } __all__ = [ "DiffusionPipelineBehavior", + "FluxPipelineBehavior", "PIPELINE_BEHAVIORS", "SD15PipelineBehavior", "SDXLPipelineBehavior", diff --git a/extension/diffusers/pipeline_behaviors/base.py b/extension/diffusers/pipeline_behaviors/base.py index e597094c0fc..b63272be721 100644 --- a/extension/diffusers/pipeline_behaviors/base.py +++ b/extension/diffusers/pipeline_behaviors/base.py @@ -47,7 +47,7 @@ def create_latents( *, height: int, width: int, - vae_scale_factor: int, + vae_scale_factor: int = 8, generator: torch.Generator, dtype: torch.dtype, scheduler: Any, @@ -90,6 +90,7 @@ def prepare_schedule( *, height: int, width: int, + vae_scale_factor: int = 8, ) -> torch.Tensor: """Configure the scheduler and return its inference timesteps.""" scheduler.set_timesteps(steps) @@ -117,6 +118,7 @@ def prepare_denoiser_inputs( width: int, dtype: torch.dtype, guidance_scale: float | None, + vae_scale_factor: int = 8, ) -> tuple[torch.Tensor, ...]: """Assemble one invocation of the exported denoiser method.""" ... diff --git a/extension/diffusers/pipeline_behaviors/flux.py b/extension/diffusers/pipeline_behaviors/flux.py new file mode 100644 index 00000000000..f12194d63f8 --- /dev/null +++ b/extension/diffusers/pipeline_behaviors/flux.py @@ -0,0 +1,160 @@ +# 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. + +"""FLUX text-to-image runtime behavior.""" + +from typing import Any + +import torch +from executorch.exir._warnings import experimental + +from .base import DiffusionPipelineBehavior + + +def pack_latents(latents: torch.Tensor) -> torch.Tensor: + """Pack 2x2 spatial neighborhoods into the channel dimension.""" + batch, channels, height, width = latents.shape + if height % 2 or width % 2: + raise ValueError("FLUX latent height and width must be even") + return ( + latents.reshape(batch, channels, height // 2, 2, width // 2, 2) + .permute(0, 2, 4, 1, 3, 5) + .reshape(batch, height * width // 4, channels * 4) + ) + + +def unpack_latents( + latents: torch.Tensor, height: int, width: int, vae_scale_factor: int +) -> torch.Tensor: + """Restore packed FLUX latents to the VAE's NCHW layout.""" + latent_height = height // vae_scale_factor + latent_width = width // vae_scale_factor + batch, _, packed_channels = latents.shape + if packed_channels % 4: + raise ValueError("FLUX packed latent channels must be divisible by four") + return ( + latents.reshape( + batch, latent_height // 2, latent_width // 2, packed_channels // 4, 2, 2 + ) + .permute(0, 3, 1, 4, 2, 5) + .reshape(batch, packed_channels // 4, latent_height, latent_width) + ) + + +def image_position_ids(height: int, width: int, device=None) -> torch.Tensor: + """Create FLUX (type, row, column) image position IDs.""" + rows = torch.arange(height, device=device).reshape(height, 1).expand(height, width) + columns = torch.arange(width, device=device).reshape(1, width).expand(height, width) + return torch.stack((torch.zeros_like(rows), rows, columns), dim=-1).reshape(-1, 3) + + +@experimental("This API is experimental and may change without notice.") +class FluxPipelineBehavior(DiffusionPipelineBehavior): + """Runtime behavior shared by FLUX.1-schnell and FLUX.1-dev.""" + + name = "flux" + supports_cfg = True + accepts_guidance_input = True + supports_negative_prompt = True + default_cfg_scale = 1.0 + default_guidance_scale = 3.5 + default_inference_steps = 28 + default_zero_negative_prompt = False + tokenizer_slots = ("tokenizer", "tokenizer_2") + text_encoder_slots = ("text_encoder", "text_encoder_2") + text_encoder_input_names = (("input_ids",), ("input_ids",)) + + def assemble_conditioning(self, encoder_outputs): + (pooled,), (prompt,) = encoder_outputs + return prompt, pooled + + def should_use_cfg(self, cfg_scale: float, negative_prompt: str | None) -> bool: + return cfg_scale > 1 and negative_prompt is not None + + def create_latents( + self, + exported_shape, + *, + height, + width, + vae_scale_factor, + generator, + dtype, + scheduler, + ): + channels = exported_shape[-1] // 4 + noise = torch.randn( + (1, channels, height // vae_scale_factor, width // vae_scale_factor), + generator=generator, + dtype=dtype, + ) + return pack_latents(noise) + + def prepare_encoded_latents(self, mean, std, *, height, width, generator): + raise ValueError("FLUX img2img is not supported") + + def prepare_img2img_latents(self, scheduler, image_latents, timesteps, **kwargs): + raise ValueError("FLUX img2img is not supported") + + def prepare_schedule( + self, + scheduler: Any, + steps: int, + *, + height: int, + width: int, + vae_scale_factor: int, + ) -> torch.Tensor: + image_sequence_length = (height // vae_scale_factor // 2) * ( + width // vae_scale_factor // 2 + ) + config = scheduler.config + base_length = getattr(config, "base_image_seq_len", 256) + max_length = getattr(config, "max_image_seq_len", 4096) + base_shift = getattr(config, "base_shift", 0.5) + max_shift = getattr(config, "max_shift", 1.15) + mu = base_shift + (image_sequence_length - base_length) * ( + max_shift - base_shift + ) / (max_length - base_length) + sigmas = ( + None + if getattr(config, "use_flow_sigmas", False) + else torch.linspace(1.0, 1.0 / steps, steps, dtype=torch.float64).tolist() + ) + scheduler.set_timesteps(steps, sigmas=sigmas, mu=mu) + return scheduler.timesteps + + def prepare_denoiser_inputs( + self, + latents, + timestep, + conditioning, + *, + height, + width, + vae_scale_factor, + dtype, + guidance_scale, + ): + prompt, pooled = conditioning + text_ids = torch.zeros(prompt.shape[1], 3, dtype=torch.int64) + image_ids = image_position_ids( + height // vae_scale_factor // 2, width // vae_scale_factor // 2 + ) + inputs = ( + latents, + timestep.to(dtype) / 1000, + prompt, + pooled, + text_ids, + image_ids, + ) + if guidance_scale is not None: + return (*inputs, torch.tensor([guidance_scale], dtype=dtype)) + return inputs + + def prepare_decode_latents(self, latents, *, height, width, vae_scale_factor): + return unpack_latents(latents, height, width, vae_scale_factor) diff --git a/extension/diffusers/pipeline_behaviors/sd15.py b/extension/diffusers/pipeline_behaviors/sd15.py index 9aad34bf8f7..73c17997769 100644 --- a/extension/diffusers/pipeline_behaviors/sd15.py +++ b/extension/diffusers/pipeline_behaviors/sd15.py @@ -43,6 +43,7 @@ def prepare_denoiser_inputs( width: int, dtype: torch.dtype, guidance_scale: float | None, + vae_scale_factor: int = 8, ) -> tuple[torch.Tensor, ...]: inputs = (latents, timestep, *conditioning) if guidance_scale is None: diff --git a/extension/diffusers/pipeline_behaviors/sdxl.py b/extension/diffusers/pipeline_behaviors/sdxl.py index ccca936663a..56076a20c29 100644 --- a/extension/diffusers/pipeline_behaviors/sdxl.py +++ b/extension/diffusers/pipeline_behaviors/sdxl.py @@ -44,6 +44,7 @@ def prepare_denoiser_inputs( width: int, dtype: torch.dtype, guidance_scale: float | None, + vae_scale_factor: int = 8, ) -> tuple[torch.Tensor, ...]: # SDXL time_ids: original size, crop coordinates, and target size. inputs = ( diff --git a/extension/diffusers/runner.py b/extension/diffusers/runner.py index a3029fa9a44..4449c838527 100644 --- a/extension/diffusers/runner.py +++ b/extension/diffusers/runner.py @@ -280,6 +280,22 @@ def _resolve_dimensions( ) return output_height, output_width + def _resolve_guidance_scale(self, guidance_scale: float | None) -> float | None: + guidance_scale = ( + self.metadata.default_guidance_scale + if guidance_scale is None + else guidance_scale + ) + if guidance_scale is not None and ( + not math.isfinite(guidance_scale) or guidance_scale < 0 + ): + raise ValueError("guidance_scale must be finite and non-negative") + if guidance_scale is not None and not self.metadata.accepts_guidance_input: + if guidance_scale != 0: + raise ValueError("this program does not accept a guidance input") + guidance_scale = None + return guidance_scale + def generate( self, prompt: str, @@ -296,23 +312,13 @@ def generate( ) -> torch.Tensor: steps = self.metadata.default_inference_steps if steps is None else steps cfg_scale = self.metadata.default_cfg_scale if cfg_scale is None else cfg_scale - guidance_scale = ( - self.metadata.default_guidance_scale - if guidance_scale is None - else guidance_scale - ) if steps <= 0: raise ValueError("steps must be positive") if not math.isfinite(cfg_scale) or cfg_scale < 0: raise ValueError("cfg_scale must be finite and non-negative") if cfg_scale > 1 and not self.metadata.supports_cfg: raise ValueError("this program does not support CFG") - if guidance_scale is not None and ( - not math.isfinite(guidance_scale) or guidance_scale < 0 - ): - raise ValueError("guidance_scale must be finite and non-negative") - if guidance_scale is not None and not self.metadata.accepts_guidance_input: - raise ValueError("this program does not accept a guidance input") + guidance_scale = self._resolve_guidance_scale(guidance_scale) if not 0 < strength <= 1: raise ValueError("strength must be in (0, 1]") if negative_prompt is not None and not self.metadata.supports_negative_prompt: @@ -332,6 +338,7 @@ def generate( steps, height=output_height, width=output_width, + vae_scale_factor=self.metadata.vae_scale_factor, ) dtype = getattr(torch, self.metadata.denoiser_dtype) vae_dtype = getattr(torch, self.metadata.vae_dtype) @@ -394,6 +401,7 @@ def generate( conditioning, height=output_height, width=output_width, + vae_scale_factor=self.metadata.vae_scale_factor, dtype=dtype, guidance_scale=guidance_scale, ), diff --git a/extension/diffusers/test/test_diffusers.py b/extension/diffusers/test/test_diffusers.py index e085838a03a..052e2dc9ad3 100644 --- a/extension/diffusers/test/test_diffusers.py +++ b/extension/diffusers/test/test_diffusers.py @@ -22,9 +22,15 @@ from executorch.export.huggingface.diffusers.pipeline_specs import DiffusionModelInfo from executorch.extension.diffusers.metadata import DiffusionMetadata, METADATA_METHODS from executorch.extension.diffusers.pipeline_behaviors import ( + FluxPipelineBehavior, SD15PipelineBehavior, SDXLPipelineBehavior, ) +from executorch.extension.diffusers.pipeline_behaviors.flux import ( + image_position_ids, + pack_latents, + unpack_latents, +) from executorch.extension.diffusers.runner import ( _cast_floating_tensors, DiffusionRunner, @@ -166,6 +172,82 @@ def test_model_guidance_scale_preserves_float32_precision(self): ) torch.testing.assert_close(inputs[-1], torch.tensor([7.3])) + def test_flux_true_cfg_requires_an_explicit_negative_prompt(self): + behavior = FluxPipelineBehavior() + self.assertFalse(behavior.should_use_cfg(4.0, None)) + self.assertTrue(behavior.should_use_cfg(4.0, "blurry")) + + +class FluxBehaviorTest(unittest.TestCase): + def test_position_ids_preserve_coordinates_with_bfloat16_denoiser(self): + inputs = FluxPipelineBehavior().prepare_denoiser_inputs( + torch.zeros(1, 260, 16, dtype=torch.bfloat16), + torch.zeros(1), + (torch.zeros(1, 8, 5), torch.zeros(1, 3)), + height=4160, + width=16, + vae_scale_factor=8, + dtype=torch.bfloat16, + guidance_scale=None, + ) + self.assertEqual(inputs[5][-1].tolist(), [0, 259, 0]) + self.assertEqual(inputs[5].dtype, torch.int64) + + def test_packed_latents_round_trip(self): + latents = torch.arange(1 * 4 * 4 * 6).reshape(1, 4, 4, 6) + packed = pack_latents(latents) + self.assertEqual(packed.shape, (1, 6, 16)) + torch.testing.assert_close(unpack_latents(packed, 32, 48, 8), latents) + + def test_position_ids_encode_rows_and_columns(self): + torch.testing.assert_close( + image_position_ids(2, 2), + torch.tensor([[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1]]), + ) + + def test_flow_matching_schedule_uses_dynamic_shift(self): + class Scheduler: + config = SimpleNamespace( + base_image_seq_len=4, + max_image_seq_len=16, + base_shift=0.5, + max_shift=1.1, + ) + + def set_timesteps(self, steps, *, sigmas, mu): + self.call = steps, sigmas, mu + self.timesteps = torch.arange(steps) + + scheduler = Scheduler() + FluxPipelineBehavior().prepare_schedule( + scheduler, 4, height=32, width=64, vae_scale_factor=8 + ) + self.assertEqual(scheduler.call[0], 4) + self.assertAlmostEqual(scheduler.call[1][-1], 0.25) + self.assertAlmostEqual(scheduler.call[2], 0.7) + + def test_flow_matching_schedule_uses_scheduler_flow_sigmas(self): + scheduler = mock.Mock() + scheduler.config = SimpleNamespace(use_flow_sigmas=True) + FluxPipelineBehavior().prepare_schedule( + scheduler, 4, height=32, width=64, vae_scale_factor=8 + ) + self.assertIsNone(scheduler.set_timesteps.call_args.kwargs["sigmas"]) + + def test_denoiser_inputs_scale_timestep_and_add_guidance(self): + inputs = FluxPipelineBehavior().prepare_denoiser_inputs( + torch.zeros(1, 8, 16), + torch.tensor([500.0]), + (torch.zeros(1, 9, 5), torch.zeros(1, 3)), + height=32, + width=64, + vae_scale_factor=8, + dtype=torch.float32, + guidance_scale=3.5, + ) + torch.testing.assert_close(inputs[1], torch.tensor([0.5])) + torch.testing.assert_close(inputs[-1], torch.tensor([3.5])) + class PipelineSpecTest(unittest.TestCase): def setUp(self):