Add a shortfin pipeline for flux (#876)

Co-authored-by: Ean Garvey <ean.garvey@amd.com>

sharktank/sharktank/layers/configs/llm_configs.py

@@ -302,7 +302,7 @@ class ClipTextConfig:
 
     @staticmethod
     def from_hugging_face_clip_text_model_config(
-        config: "transformers.CLIPTextConfig",
+        config: "transformers.CLIPTextConfig",  # type: ignore
     ) -> "ClipTextConfig":
         return ClipTextConfig(
             vocab_size=config.vocab_size,
@@ -323,7 +323,7 @@ def from_hugging_face_clip_text_model_config(
             dtype=config.torch_dtype or torch.float32,
         )
 
-    def to_hugging_face_clip_text_model_config(self) -> "transformers.CLIPTextConfig":
+    def to_hugging_face_clip_text_model_config(self) -> "transformers.CLIPTextConfig":  # type: ignore
         kwargs = self.to_properties()
         kwargs["torch_dtype"] = kwargs["dtype"]
         del kwargs["dtype"]

sharktank/sharktank/models/clip/export.py

@@ -4,6 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+import functools
 from typing import Optional, Union
 import transformers
 from transformers.models.clip.modeling_clip import (
@@ -18,6 +19,7 @@
 from ...layers.configs import ClipTextConfig
 from .clip import ClipTextModel
 from iree.turbine.aot import FxProgramsBuilder, export
+from sharktank.transforms.dataset import set_float_dtype
 
 
 def hugging_face_clip_attention_to_theta(model: HfCLIPAttention) -> Theta:
@@ -50,8 +52,14 @@ def clip_text_model_to_dataset(model: ClipTextModel) -> Dataset:
     return Dataset(properties=model.config.to_properties(), root_theta=model.theta)
 
 
-def export_clip_text_model_iree_parameters(model: ClipTextModel, output_path: PathLike):
+def export_clip_text_model_iree_parameters(
+    model: ClipTextModel, output_path: PathLike, dtype: torch.dtype = None
+):
     dataset = clip_text_model_to_dataset(model)
+    if dtype:
+        dataset.root_theta = dataset.root_theta.transform(
+            functools.partial(set_float_dtype, dtype=dtype)
+        )
     dataset.save(output_path)
 
 

sharktank/sharktank/models/flux/export.py

@@ -4,6 +4,7 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+import functools
 from os import PathLike
 import os
 from pathlib import Path
@@ -14,6 +15,7 @@
 from .flux import FluxModelV1, FluxParams
 from ...types import Dataset
 from ...utils.hf_datasets import get_dataset
+from sharktank.transforms.dataset import set_float_dtype
 
 flux_transformer_default_batch_sizes = [1]
 
@@ -27,11 +29,16 @@ def export_flux_transformer_model_mlir(
 
 
 def export_flux_transformer_iree_parameters(
-    model: FluxModelV1, parameters_output_path: PathLike
+    model: FluxModelV1, parameters_output_path: PathLike, dtype=None
 ):
     model.theta.rename_tensors_to_paths()
-    # TODO: export properties
-    dataset = Dataset(root_theta=model.theta, properties={})
+    dataset = Dataset(
+        root_theta=model.theta, properties=model.params.to_hugging_face_properties()
+    )
+    if dtype:
+        dataset.root_theta = dataset.root_theta.transform(
+            functools.partial(set_float_dtype, dtype=dtype)
+        )
     dataset.save(parameters_output_path)
 
 

sharktank/sharktank/models/flux/flux.py

@@ -12,8 +12,8 @@
 from typing import Any, Optional
 from collections import OrderedDict
 from copy import copy
+from dataclasses import dataclass, asdict
 import math
-from dataclasses import dataclass
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -49,6 +49,19 @@ class FluxParams:
     qkv_bias: bool
     guidance_embed: bool
 
+    def to_hugging_face_properties(self) -> dict[str, Any]:
+        hparams = {
+            "in_channels": self.in_channels,
+            "pooled_projection_dim": self.vec_in_dim,
+            "joint_attention_dim": self.context_in_dim,
+            "num_attention_heads": self.num_heads,
+            "num_layers": self.depth,
+            "num_single_layers": self.depth_single_blocks,
+            "attention_head_dim": sum(self.axes_dim),
+            "guidance_embeds": self.guidance_embed,
+        }
+        return {"hparams": hparams}
+
     @staticmethod
     def from_hugging_face_properties(properties: dict[str, Any]) -> "FluxParams":
         p = properties["hparams"]
@@ -175,6 +188,7 @@ def forward(
                     "Didn't get guidance strength for guidance distilled model."
                 )
             vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
+
         vec = vec + self.vector_in(y)
 
         txt = self.txt_in(txt)

sharktank/sharktank/models/vae/model.py

@@ -74,15 +74,6 @@ def forward(
                 "latent_embeds": latent_embeds,
             },
         )
-        if not self.hp.use_post_quant_conv:
-            sample = rearrange(
-                sample,
-                "b (h w) (c ph pw) -> b c (h ph) (w pw)",
-                h=math.ceil(1024 / 16),
-                w=math.ceil(1024 / 16),
-                ph=2,
-                pw=2,
-            )
         sample = sample / self.hp.scaling_factor + self.hp.shift_factor
 
         if self.hp.use_post_quant_conv:

sharktank/sharktank/pipelines/flux/README.md

@@ -0,0 +1,10 @@
+# Flux.1 dynamo exports
+
+### Quick Start
+
+All the exports in this directory are done through `export.py`, with the CLI syntax as follows:
+```shell
+python -m sharktank.pipelines.flux.export_parameters --dtype <fp32/fp16/bf16> --input-dir <input-dir> --output-dir <output-dir>
+
+python -m sharktank.pipelines.flux.export_components --model="flux-dev" --component=<clip/vae/t5xxl/mmdit/scheduler> --precision=<fp32/fp16/bf16>
+```

sharktank/sharktank/pipelines/flux/__init__.py

@@ -0,0 +1,7 @@
+"""Flux text-to-image generation pipeline."""
+
+from .flux_pipeline import FluxPipeline
+
+__all__ = [
+    "FluxPipeline",
+]