Address SeFi review feedback

Author: jmliu206

docs/source/en/api/pipelines/sefi.md

@@ -17,14 +17,27 @@ It denoises a semantic latent stream slightly ahead of a texture latent stream,
 to images.
 
 The public checkpoints on the Hub are gated and distributed under a non-commercial license. Accept the checkpoint
-license on the Hub before downloading them. Original SeFi-Image artifacts can be converted with
-`scripts/convert_sefi_to_diffusers.py` before loading them with [`SeFiPipeline`].
+license on the Hub before downloading them. Standalone Diffusers checkpoints are available for the following models:
+
+| Model | Type |
+|---|---|
+| [`SeFi-Image/SeFi-Image-1B-Base-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-1B-Base-diffusers) | Base |
+| [`SeFi-Image/SeFi-Image-1B-turbo-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-1B-turbo-diffusers) | Turbo |
+| [`SeFi-Image/SeFi-Image-2B-Base-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-2B-Base-diffusers) | Base |
+| [`SeFi-Image/SeFi-Image-2B-turbo-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-2B-turbo-diffusers) | Turbo |
+| [`SeFi-Image/SeFi-Image-5B-Base-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-5B-Base-diffusers) | Base |
+| [`SeFi-Image/SeFi-Image-5B-turbo-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-5B-turbo-diffusers) | Turbo |
+| [`SeFi-Image/SeFi-Image-5B-RL-diffusers`](https://huggingface.co/SeFi-Image/SeFi-Image-5B-RL-diffusers) | RL |
+
+Original SeFi-Image artifacts can also be converted with `scripts/convert_sefi_to_diffusers.py`.
 
 ```python
 import torch
 from diffusers import SeFiPipeline
 
-pipe = SeFiPipeline.from_pretrained("./sefi-1b-base-diffusers", torch_dtype=torch.bfloat16)
+pipe = SeFiPipeline.from_pretrained(
+    "SeFi-Image/SeFi-Image-1B-Base-diffusers", dtype=torch.bfloat16
+)
 pipe.to("cuda")
 
 image = pipe("A red apple on a wooden table.").images[0]

scripts/convert_sefi_to_diffusers.py

@@ -23,7 +23,7 @@
 from huggingface_hub import snapshot_download
 from safetensors.torch import load_file
 
-from diffusers import __version__
+from diffusers import FlowMatchEulerDiscreteScheduler, __version__
 from diffusers.models import SeFiTransformer2DModel
 
 
@@ -106,6 +106,16 @@ def parse_args():
     parser.add_argument("--output", required=True, help="Output Diffusers checkpoint folder.")
     parser.add_argument("--cache-dir", default=None, help="Optional Hugging Face cache directory.")
     parser.add_argument("--token", default=None, help="Optional Hugging Face token for gated checkpoints.")
+    parser.add_argument(
+        "--source-repo-id",
+        default=None,
+        help="Original Hub repo id to record in the converted model card when converting a local checkpoint.",
+    )
+    parser.add_argument(
+        "--target-repo-id",
+        default=None,
+        help="Converted Hub repo id to use in the generated model card example.",
+    )
     parser.add_argument(
         "--variant",
         choices=["base", "rl", "turbo"],
@@ -144,6 +154,66 @@ def save_json(path: Path, payload: dict):
         handle.write("\n")
 
 
+def save_model_card(output: Path, source_repo_id: str | None, target_repo_id: str, variant: str):
+    metadata = [
+        "---",
+        "license: cc-by-nc-4.0",
+        "library_name: diffusers",
+        "pipeline_tag: text-to-image",
+        "gated: true",
+        "tags:",
+        "- sefi-image",
+        "- semantic-first-diffusion",
+        "- safetensors",
+    ]
+    if source_repo_id is not None:
+        metadata.append(f"base_model: {source_repo_id}")
+    metadata.append("---")
+
+    source_link = (
+        f"[`{source_repo_id}`](https://huggingface.co/{source_repo_id})"
+        if source_repo_id is not None
+        else "the original SeFi-Image checkpoint"
+    )
+    inference_call = (
+        """image = pipe(
+    \"A red apple on a wooden table.\",
+    num_inference_steps=4,
+    guidance_scale=1.0,
+).images[0]"""
+        if variant == "turbo"
+        else 'image = pipe("A red apple on a wooden table.").images[0]'
+    )
+    card = (
+        "\n".join(metadata)
+        + f"""
+
+# SeFi-Image Diffusers checkpoint
+
+This repository is a Diffusers-format conversion of {source_link}. The original checkpoint is not modified by the
+conversion. Refer to the source model card for model details, limitations, and responsible-use guidance.
+
+```python
+import torch
+from diffusers import SeFiPipeline
+
+pipe = SeFiPipeline.from_pretrained(
+    \"{target_repo_id}\", dtype=torch.bfloat16
+).to(\"cuda\")
+{inference_call}
+image.save(\"sefi.png\")
+```
+
+## License
+
+The checkpoint is distributed under the Creative Commons Attribution-NonCommercial 4.0 International license
+(CC BY-NC 4.0). It is for non-commercial use only.
+"""
+    )
+    with open(output / "README.md", "w", encoding="utf-8") as handle:
+        handle.write(card)
+
+
 def infer_variant(checkpoint: str, config: dict, explicit_variant: str | None) -> str:
     if explicit_variant is not None:
         return explicit_variant
@@ -165,9 +235,7 @@ def default_guidance_scale(variant: str) -> float:
 
 
 def texture_vae_config_path(root: Path, texture_vae_name: str) -> Path:
-    if texture_vae_name == "sd1.5":
-        return root / "vae" / "config.json"
-    if texture_vae_name in {"flux1", "flux2"}:
+    if texture_vae_name in {"sd1.5", "flux1", "flux2"}:
         return root / "vae" / "config.json"
     raise ValueError(f"Unsupported texture VAE: {texture_vae_name}")
 
@@ -198,7 +266,6 @@ def build_transformer_config(root: Path, sefi_config: dict) -> dict:
 
     transformer_config["in_channels"] = total_channels
     transformer_config["out_channels"] = total_channels
-    transformer_config["text_input_dim"] = text_dim
     if int(transformer_config["joint_attention_dim"]) != text_dim:
         raise ValueError(
             "Text dimension mismatch: "
@@ -225,6 +292,38 @@ def load_transformer_state_dict(transformer_dir: Path) -> dict[str, torch.Tensor
     raise FileNotFoundError(f"No supported transformer weights found under {transformer_dir}.")
 
 
+def convert_transformer_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    converted = {}
+    for key, value in state_dict.items():
+        if key.startswith("backbone."):
+            converted_key = key.removeprefix("backbone.")
+        elif key.startswith("dual_time_embed."):
+            converted_key = key
+        else:
+            raise ValueError(f"Unexpected transformer key in the original SeFi checkpoint: {key}")
+
+        if converted_key in converted:
+            raise ValueError(f"Transformer key collision after conversion: {converted_key}")
+        converted[converted_key] = value
+    return converted
+
+
+def convert_scheduler(root: Path, output: Path):
+    scheduler_config = load_json(root / "scheduler" / "scheduler_config.json")
+    source_scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+    scheduler_config["shift"] = 1.0
+    scheduler_config["use_dynamic_shifting"] = False
+    scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+
+    torch.testing.assert_close(scheduler.sigmas, source_scheduler.sigmas, rtol=0.0, atol=0.0)
+    torch.testing.assert_close(scheduler.timesteps, source_scheduler.timesteps, rtol=0.0, atol=0.0)
+    num_train_timesteps = int(scheduler.config.num_train_timesteps)
+    expected_sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps)
+    torch.testing.assert_close(scheduler.sigmas, expected_sigmas, rtol=0.0, atol=1e-7)
+    torch.testing.assert_close(scheduler.timesteps, expected_sigmas * num_train_timesteps, rtol=0.0, atol=1e-4)
+    scheduler.save_pretrained(output / "scheduler")
+
+
 def copy_tokenizer_files(src: Path, dst: Path):
     weight_patterns = {
         "model*.safetensors",
@@ -253,14 +352,28 @@ def main():
     variant = infer_variant(args.checkpoint, sefi_config, args.variant)
     transformer_config = build_transformer_config(root, sefi_config)
 
-    transformer = SeFiTransformer2DModel(**transformer_config)
-    state_dict = load_transformer_state_dict(root / "transformer")
-    missing, unexpected = transformer.load_state_dict(state_dict, strict=False)
-    if missing or unexpected:
-        raise ValueError(f"Transformer state dict mismatch. Missing={missing[:20]}, unexpected={unexpected[:20]}")
+    state_dict = convert_transformer_state_dict(load_transformer_state_dict(root / "transformer"))
+    with torch.device("meta"):
+        transformer = SeFiTransformer2DModel(**transformer_config)
+    transformer.load_state_dict(state_dict, strict=True, assign=True)
+    expected_transformer_keys = set(state_dict)
+    expected_transformer_dtypes = {key: value.dtype for key, value in state_dict.items()}
+    floating_dtypes = {value.dtype for value in state_dict.values() if value.is_floating_point()}
+    if len(floating_dtypes) != 1:
+        raise ValueError(f"Expected one floating-point transformer dtype, got {sorted(map(str, floating_dtypes))}.")
+    transformer_dtype = floating_dtypes.pop()
     transformer.save_pretrained(output / "transformer", safe_serialization=True)
+    del transformer, state_dict
+
+    reloaded_transformer = SeFiTransformer2DModel.from_pretrained(output / "transformer", dtype=transformer_dtype)
+    if set(reloaded_transformer.state_dict()) != expected_transformer_keys:
+        raise ValueError("Transformer state dict keys changed after the save/load round trip.")
+    round_trip_dtypes = {key: value.dtype for key, value in reloaded_transformer.state_dict().items()}
+    if round_trip_dtypes != expected_transformer_dtypes:
+        raise ValueError("Transformer state dict dtypes changed after the save/load round trip.")
+    del reloaded_transformer
 
-    copytree(root / "scheduler", output / "scheduler")
+    convert_scheduler(root, output)
     copytree(root / "vae", output / "vae")
 
     text_encoder_name = sefi_config["model"]["text_encoder"]["model_name"]
@@ -299,6 +412,10 @@ def main():
         "max_sequence_length": int(model_config["text_encoder"].get("max_length", 1024)),
     }
     save_json(output / "model_index.json", model_index)
+    source_repo_id = args.source_repo_id
+    if source_repo_id is None and not Path(args.checkpoint).expanduser().exists():
+        source_repo_id = args.checkpoint
+    save_model_card(output, source_repo_id, args.target_repo_id or output.name, variant)
     shutil.copy2(root / "sefi_config.yaml", output / "sefi_config.yaml")
     print(f"Saved SeFi-Image Diffusers checkpoint to {output}")
 

src/diffusers/models/transformers/transformer_sefi.py

@@ -19,10 +19,11 @@
 import torch.nn as nn
 
 from ...configuration_utils import ConfigMixin, register_to_config
-from ...utils import BaseOutput, apply_lora_scale
+from ...utils import BaseOutput
 from ..embeddings import TimestepEmbedding, Timesteps
 from ..modeling_utils import ModelMixin
-from .transformer_flux2 import Flux2Transformer2DModel
+from ..normalization import AdaLayerNormContinuous
+from .transformer_flux2 import Flux2Modulation, Flux2PosEmbed, Flux2SingleTransformerBlock, Flux2TransformerBlock
 
 
 @dataclass
@@ -105,8 +106,6 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
             RoPE theta.
         eps (`float`, defaults to `1e-6`):
             Normalization epsilon.
-        text_input_dim (`int`, *optional*):
-            Expected text embedding dimension. Defaults to `joint_attention_dim`.
     """
 
     _supports_gradient_checkpointing = True
@@ -130,44 +129,60 @@ def __init__(
         axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
         rope_theta: int = 2000,
         eps: float = 1e-6,
-        text_input_dim: int | None = None,
     ):
         super().__init__()
 
-        text_input_dim = joint_attention_dim if text_input_dim is None else text_input_dim
-        if int(text_input_dim) != int(joint_attention_dim):
-            raise ValueError(
-                f"`text_input_dim` must match `joint_attention_dim`, got {text_input_dim} and {joint_attention_dim}."
-            )
-
         self.out_channels = out_channels or in_channels
         self.inner_dim = num_attention_heads * attention_head_dim
-        self.backbone = Flux2Transformer2DModel(
-            patch_size=patch_size,
-            in_channels=in_channels,
-            out_channels=out_channels,
-            num_layers=num_layers,
-            num_single_layers=num_single_layers,
-            attention_head_dim=attention_head_dim,
-            num_attention_heads=num_attention_heads,
-            joint_attention_dim=joint_attention_dim,
-            timestep_guidance_channels=timestep_guidance_channels,
-            mlp_ratio=mlp_ratio,
-            axes_dims_rope=axes_dims_rope,
-            rope_theta=rope_theta,
-            eps=eps,
-            guidance_embeds=False,
-        )
-        # The reference SeFi transformer deletes Flux2's timestep/guidance embedder and stores only the dual embedder.
-        self.backbone.time_guidance_embed = nn.Identity()
+
+        self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
         self.dual_time_embed = SeFiDualTimestepEmbeddings(
             in_channels=timestep_guidance_channels,
             embedding_dim=self.inner_dim,
             bias=False,
         )
+
+        self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False)
+
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Flux2TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                Flux2SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False
+        )
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False)
+
         self.gradient_checkpointing = False
 
-    @apply_lora_scale("joint_attention_kwargs")
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -211,26 +226,26 @@ def forward(
         timestep_tex = timestep_tex.to(hidden_states.dtype) * 1000
         temb = self.dual_time_embed(timestep_sem, timestep_tex)
 
-        double_stream_mod_img = self.backbone.double_stream_modulation_img(temb)
-        double_stream_mod_txt = self.backbone.double_stream_modulation_txt(temb)
-        single_stream_mod = self.backbone.single_stream_modulation(temb)
+        double_stream_mod_img = self.double_stream_modulation_img(temb)
+        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
+        single_stream_mod = self.single_stream_modulation(temb)
 
-        hidden_states = self.backbone.x_embedder(hidden_states)
-        encoder_hidden_states = self.backbone.context_embedder(encoder_hidden_states)
+        hidden_states = self.x_embedder(hidden_states)
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
 
         if img_ids.ndim == 3:
             img_ids = img_ids[0]
         if txt_ids.ndim == 3:
             txt_ids = txt_ids[0]
 
-        image_rotary_emb = self.backbone.pos_embed(img_ids)
-        text_rotary_emb = self.backbone.pos_embed(txt_ids)
+        image_rotary_emb = self.pos_embed(img_ids)
+        text_rotary_emb = self.pos_embed(txt_ids)
         concat_rotary_emb = (
             torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
             torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
         )
 
-        for block in self.backbone.transformer_blocks:
+        for block in self.transformer_blocks:
             if torch.is_grad_enabled() and self.gradient_checkpointing:
                 encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
                     block,
@@ -253,7 +268,7 @@ def forward(
 
         hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
 
-        for block in self.backbone.single_transformer_blocks:
+        for block in self.single_transformer_blocks:
             if torch.is_grad_enabled() and self.gradient_checkpointing:
                 hidden_states = self._gradient_checkpointing_func(
                     block,
@@ -273,8 +288,8 @@ def forward(
                 )
 
         hidden_states = hidden_states[:, num_txt_tokens:, ...]
-        hidden_states = self.backbone.norm_out(hidden_states, temb)
-        output = self.backbone.proj_out(hidden_states)
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
 
         if not return_dict:
             return (output,)