update sefi model and add rl variant

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -3,7 +3,7 @@
 ## Update for 2026-08-04
 
 - **Models**
-  - [SeFi-Image](https://huggingface.co/SeFi-Image/SeFi-Image-5B-RL) in *Base* and *Turbo* (distilled) variants  
+  - [SeFi-Image](https://huggingface.co/SeFi-Image/SeFi-Image-5B-RL) in *Base*, *Turbo* (distilled) and *RL* (finetuned) variants  
     SeFi is an interesting model that separates generation into semantic and texture latent streams  
     and denoising semantic structure slightly ahead of texture details  
     SeFi comes in sizes with 1B, 2B and 5B params  

data/reference-base.json

@@ -928,27 +928,35 @@
     "date": "2026 July"
   },
   "SeFi-Image 1B Base": {
-    "path": "vladmandic/SeFi-Image-1B-Base",
-    "preview": "vladmandic--SeFi-Image-1B-Base.jpg",
+    "path": "SeFi-Image/SeFi-Image-1B-Base-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-1B-Base-diffusers.jpg",
     "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
     "extras": "sampler: Default",
-    "size": 0,
+    "size": 6.32,
     "date": "2026 July"
   },
   "SeFi-Image 2B Base": {
-    "path": "vladmandic/SeFi-Image-2B-Base",
-    "preview": "vladmandic--SeFi-Image-2B-Base.jpg",
+    "path": "SeFi-Image/SeFi-Image-2B-Base-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-2B-Base-diffusers.jpg",
     "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
     "extras": "sampler: Default",
-    "size": 0,
+    "size": 8.18,
     "date": "2026 July"
   },
   "SeFi-Image 5B Base": {
-    "path": "vladmandic/SeFi-Image-5B-Base",
-    "preview": "vladmandic--SeFi-Image-5B-Base.jpg",
+    "path": "SeFi-Image/SeFi-Image-5B-Base-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-5B-Base-diffusers.jpg",
     "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
     "extras": "sampler: Default",
-    "size": 0,
+    "size": 17.69,
+    "date": "2026 July"
+  },
+  "SeFi-Image 5B RL": {
+    "path": "SeFi-Image/SeFi-Image-5B-RL-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-5B-RL-diffusers.jpg",
+    "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+    "extras": "sampler: Default",
+    "size": 17.69,
     "date": "2026 July"
   }
 }

data/reference-distilled.json

@@ -235,27 +235,27 @@
     "date": "2026 July"
   },
   "SeFi-Image 1B Turbo": {
-    "path": "vladmandic/SeFi-Image-1B-Turbo",
-    "preview": "vladmandic--SeFi-Image-1B-Turbo.jpg",
+    "path": "SeFi-Image/SeFi-Image-1B-turbo-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-1B-turbo-diffusers.jpg",
     "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
     "extras": "sampler: Default",
-    "size": 0,
+    "size": 6.32,
     "date": "2026 July"
   },
   "SeFi-Image 2B Turbo": {
-    "path": "vladmandic/SeFi-Image-2B-Turbo",
-    "preview": "vladmandic--SeFi-Image-2B-Turbo.jpg",
+    "path": "SeFi-Image/SeFi-Image-2B-turbo-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-2B-turbo-diffusers.jpg",
     "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
     "extras": "sampler: Default",
-    "size": 0,
+    "size": 8.18,
     "date": "2026 July"
   },
   "SeFi-Image 5B Turbo": {
-    "path": "vladmandic/SeFi-Image-5B-Turbo",
-    "preview": "vladmandic--SeFi-Image-5B-Turbo.jpg",
+    "path": "SeFi-Image/SeFi-Image-5B-turbo-diffusers",
+    "preview": "SeFi-Image--SeFi-Image-5B-turbo-diffusers.jpg",
     "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
     "extras": "sampler: Default",
-    "size": 0,
+    "size": 17.69,
     "date": "2026 July"
   }
 }

modules/processing.py

@@ -423,13 +423,13 @@ def print_stats():
     if shared.opts.sdnq_dequantize_compile:
         from modules.timer_sdnq import update_sdnq_attention_timers
         update_sdnq_attention_timers()
-        if timer.autotune.get_total() > 0.001:
+        if timer.autotune.get_total() > 0.1:
             log.debug(f'Processed: autotune={timer.autotune.dct(min_time=0)}')
 
     if devices.triton_ok:
         from modules.sd_models_compile import update_compile_times
         update_compile_times()
-        if timer.dynamo.get_total() > 0.001:
+        if timer.dynamo.get_total() > 0.1:
             log.debug(f'Processed: dynamo={timer.dynamo.dct(min_time=2.0, no_total=True)}')
 
 

pipelines/sefi/convert_sefi_to_diffusers.py

@@ -23,7 +23,7 @@ import yaml
 from huggingface_hub import snapshot_download
 from safetensors.torch import load_file
 
-from diffusers import __version__
+from diffusers import FlowMatchEulerDiscreteScheduler, __version__
 from transformer_sefi 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
 
-    copytree(root / "scheduler", output / "scheduler")
+    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
+
+    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}")
 

pipelines/sefi/pipeline_sefi.py

@@ -19,7 +19,7 @@ from transformers import Qwen2Tokenizer, Qwen3VLForConditionalGeneration
 
 from diffusers.models import AutoencoderKL, AutoencoderKLFlux2
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils import is_torch_xla_available, replace_example_docstring
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -35,16 +35,15 @@ else:
     XLA_AVAILABLE = False
 
 
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
 EXAMPLE_DOC_STRING = """
     Examples:
         ```py
         >>> 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]
         >>> image.save("sefi.png")
@@ -135,13 +134,8 @@ class SeFiPipeline(DiffusionPipeline):
         )
         if isinstance(text_encoder_hidden_layers, str):
             text_encoder_hidden_layers = tuple(int(layer) for layer in text_encoder_hidden_layers.split(","))
-        semantic_channels = 16 if semantic_channels is None else semantic_channels
-        if texture_vae_name is None:
-            texture_vae_name = "flux2" if vae is not None and hasattr(vae, "bn") else "sd1.5"
-        default_guidance_scale = 4.0 if default_guidance_scale is None else default_guidance_scale
-        default_num_inference_steps = 50 if default_num_inference_steps is None else default_num_inference_steps
-        text_encoder_hidden_layers = (9, 18, 27) if text_encoder_hidden_layers is None else text_encoder_hidden_layers
-        max_sequence_length = 1024 if max_sequence_length is None else max_sequence_length
+        elif text_encoder_hidden_layers is not None:
+            text_encoder_hidden_layers = tuple(text_encoder_hidden_layers)
         self.register_to_config(
             semantic_channels=semantic_channels,
             texture_vae_name=texture_vae_name,
@@ -150,11 +144,11 @@ class SeFiPipeline(DiffusionPipeline):
             default_num_inference_steps=default_num_inference_steps,
             delta_t=delta_t,
             timestep_shift_alpha=timestep_shift_alpha,
-            text_encoder_hidden_layers=tuple(text_encoder_hidden_layers),
+            text_encoder_hidden_layers=text_encoder_hidden_layers,
             max_sequence_length=max_sequence_length,
         )
 
-        self.semantic_channels = int(semantic_channels)
+        self.semantic_channels = semantic_channels
         self.texture_vae_name = str(texture_vae_name).lower()
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
         self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
@@ -184,24 +178,23 @@ class SeFiPipeline(DiffusionPipeline):
     def num_timesteps(self):
         return self._num_timesteps
 
-    @staticmethod
-    def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
-        B, L, _ = x.shape
-        out_ids = []
-
-        for i in range(B):
-            t = torch.arange(1) if t_coord is None else t_coord[i]
-            h = torch.arange(1)
-            w = torch.arange(1)
-            l = torch.arange(L)
-
-            coords = torch.cartesian_prod(t, h, w, l)
-            out_ids.append(coords)
-
-        return torch.stack(out_ids)
+    @property
+    def current_timestep(self):
+        return self._current_timestep
 
     @staticmethod
-    def _prepare_latent_ids(latents: torch.Tensor):
+    def _prepare_text_ids(x: torch.Tensor):
+        batch_size, sequence_length, _ = x.shape
+        text_ids = torch.cartesian_prod(
+            torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(sequence_length)
+        )
+        return text_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+    @staticmethod
+    # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_latent_ids
+    def _prepare_latent_ids(
+        latents: torch.Tensor,  # (B, C, H, W)
+    ):
         r"""
         Generates 4D position coordinates (T, H, W, L) for latent tensors.
 
@@ -231,6 +224,7 @@ class SeFiPipeline(DiffusionPipeline):
         return latent_ids
 
     @staticmethod
+    # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpatchify_latents
     def _unpatchify_latents(latents):
         batch_size, num_channels_latents, height, width = latents.shape
         latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
@@ -239,6 +233,7 @@ class SeFiPipeline(DiffusionPipeline):
         return latents
 
     @staticmethod
+    # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._pack_latents
     def _pack_latents(latents):
         """
         pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels)
@@ -250,15 +245,14 @@ class SeFiPipeline(DiffusionPipeline):
         return latents
 
     @staticmethod
-    def _unpack_latents_with_ids(
-        x: torch.Tensor, x_ids: torch.Tensor, height: int | None = None, width: int | None = None
-    ):
+    # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpack_latents_with_ids
+    def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch.Tensor]:
         """
         using position ids to scatter tokens into place
         """
         x_list = []
         for data, pos in zip(x, x_ids):
-            _, ch = data.shape
+            _, ch = data.shape  # noqa: F841
             h_ids = pos[:, 1].to(torch.int64)
             w_ids = pos[:, 2].to(torch.int64)
 
@@ -306,34 +300,19 @@ class SeFiPipeline(DiffusionPipeline):
 
     def _build_chat_text(self, prompt: str) -> str:
         messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
-        try:
-            return self.tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True,
-                enable_thinking=False,
-            )
-        except TypeError:
-            return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        return self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
 
     def _align_text_encoder_rotary_dtype(self, device: torch.device):
-        text_encoder = self.text_encoder
-        if text_encoder is None:
-            return
-
-        try:
-            text_encoder_dtype = next(text_encoder.parameters()).dtype
-        except StopIteration:
-            return
-
-        text_model = text_encoder.model if hasattr(text_encoder, "model") else text_encoder
-        language_model = getattr(text_model, "language_model", None)
-        rotary_emb = getattr(language_model, "rotary_emb", None)
-        if rotary_emb is not None:
-            # Qwen3-VL stores RoPE inverse frequencies as non-persistent buffers. `from_pretrained(torch_dtype=...)`
-            # can leave them in fp32 even when text weights are bf16, while the reference SeFi wrapper casts the whole
-            # text encoder module. Keep these buffers aligned before text encoding.
-            rotary_emb.to(device=device, dtype=text_encoder_dtype)
+        text_encoder_dtype = next(self.text_encoder.parameters()).dtype
+        # Qwen3-VL stores RoPE inverse frequencies as non-persistent buffers. `from_pretrained(dtype=...)`
+        # can leave them in fp32 even when text weights are bf16, while the reference SeFi wrapper casts the whole
+        # text encoder module. Keep these buffers aligned before text encoding.
+        self.text_encoder.model.language_model.rotary_emb.to(device=device, dtype=text_encoder_dtype)
 
     def _get_qwen3vl_prompt_embeds(
         self,
@@ -390,10 +369,10 @@ class SeFiPipeline(DiffusionPipeline):
     ):
         device = device or self._execution_device
         dtype = dtype or (self.transformer.dtype if self.transformer is not None else self.text_encoder.dtype)
-        max_sequence_length = max_sequence_length or self.config.max_sequence_length
-        text_encoder_hidden_layers = text_encoder_hidden_layers or tuple(self.config.text_encoder_hidden_layers)
 
         if prompt_embeds is None:
+            max_sequence_length = max_sequence_length or self.config.max_sequence_length
+            text_encoder_hidden_layers = text_encoder_hidden_layers or tuple(self.config.text_encoder_hidden_layers)
             prompt_embeds = self._get_qwen3vl_prompt_embeds(
                 prompt=prompt,
                 device=device,
@@ -608,7 +587,7 @@ class SeFiPipeline(DiffusionPipeline):
             dtype=torch.float32,
         )
         u_shifted_unit = _apply_timestep_shift_unit_interval(u_base_unit, self.config.timestep_shift_alpha)
-        _, base_sigmas_schedule = self._timesteps_and_sigmas(u_shifted_unit, n_dim=1, dtype=torch.float32)
+        base_timesteps_schedule, _ = self._timesteps_and_sigmas(u_shifted_unit, n_dim=1, dtype=torch.float32)
         u_sem_raw_schedule = u_shifted_unit * (1.0 + float(self.config.delta_t))
 
         self._num_timesteps = num_inference_steps
@@ -631,7 +610,7 @@ class SeFiPipeline(DiffusionPipeline):
                 _, sigmas_sem_next = self._timesteps_and_sigmas(u_sem_next, latents.ndim, latents.dtype)
                 _, sigmas_tex_next = self._timesteps_and_sigmas(u_tex_next, latents.ndim, latents.dtype)
 
-                self._current_timestep = base_sigmas_schedule[i]
+                self._current_timestep = base_timesteps_schedule[i]
                 packed_latents = self._pack_latents(latents)
                 noise_pred = self.transformer(
                     hidden_states=packed_latents,
@@ -643,8 +622,8 @@ class SeFiPipeline(DiffusionPipeline):
                     joint_attention_kwargs=self.attention_kwargs,
                     return_dict=False,
                 )[0]
-                pred_cond = noise_pred[:, : packed_latents.size(1)]
-                pred_cond = self._unpack_latents_with_ids(pred_cond, latent_ids)
+                noise_pred = noise_pred[:, : packed_latents.size(1)]
+                noise_pred = self._unpack_latents_with_ids(noise_pred, latent_ids)
 
                 if self.do_classifier_free_guidance:
                     pred_uncond = self.transformer(
@@ -659,9 +638,9 @@ class SeFiPipeline(DiffusionPipeline):
                     )[0]
                     pred_uncond = pred_uncond[:, : packed_latents.size(1)]
                     pred_uncond = self._unpack_latents_with_ids(pred_uncond, latent_ids)
-                    velocity = _combine_guided_velocity(pred_uncond, pred_cond, guidance_scale)
+                    velocity = _combine_guided_velocity(pred_uncond, noise_pred, guidance_scale)
                 else:
-                    velocity = pred_cond
+                    velocity = noise_pred
 
                 vel_sem = velocity[:, : self.semantic_channels]
                 vel_tex = velocity[:, self.semantic_channels :]
@@ -686,6 +665,8 @@ class SeFiPipeline(DiffusionPipeline):
 
                 progress_bar.update()
 
+        self._current_timestep = None
+
         if output_type == "latent":
             image = latents
         else:

pipelines/sefi/transformer_sefi.py

@@ -19,10 +19,11 @@ import torch
 import torch.nn as nn
 
 from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.utils import BaseOutput, apply_lora_scale
+from diffusers.utils import BaseOutput
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
+from diffusers.models.normalization import AdaLayerNormContinuous
+from diffusers.models.transformers.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 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
         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 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
         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 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
 
         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 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
                 )
 
         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,)