From 808a25749f92a057222daeaf76d980924812e25a Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Tue, 4 Aug 2026 12:40:55 +0200 Subject: [PATCH] update sefi model and add rl variant Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 +- data/reference-base.json | 26 ++-- data/reference-distilled.json | 18 +-- modules/processing.py | 4 +- pipelines/sefi/convert_sefi_to_diffusers.py | 139 ++++++++++++++++++-- pipelines/sefi/pipeline_sefi.py | 117 +++++++--------- pipelines/sefi/transformer_sefi.py | 97 ++++++++------ 7 files changed, 262 insertions(+), 141 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fb2ded039..d3ea72705 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/data/reference-base.json b/data/reference-base.json index 760540bf3..96b18d90d 100644 --- a/data/reference-base.json +++ b/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" } } diff --git a/data/reference-distilled.json b/data/reference-distilled.json index ceb3a5d32..96f70d453 100644 --- a/data/reference-distilled.json +++ b/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" } } diff --git a/modules/processing.py b/modules/processing.py index 9940f77a0..fafd4d9f6 100644 --- a/modules/processing.py +++ b/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)}') diff --git a/pipelines/sefi/convert_sefi_to_diffusers.py b/pipelines/sefi/convert_sefi_to_diffusers.py index 52d6c2565..fca38a5a5 100644 --- a/pipelines/sefi/convert_sefi_to_diffusers.py +++ b/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}") diff --git a/pipelines/sefi/pipeline_sefi.py b/pipelines/sefi/pipeline_sefi.py index 120b4690c..3dc18d1cc 100644 --- a/pipelines/sefi/pipeline_sefi.py +++ b/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: diff --git a/pipelines/sefi/transformer_sefi.py b/pipelines/sefi/transformer_sefi.py index 96ec29309..75904c79c 100644 --- a/pipelines/sefi/transformer_sefi.py +++ b/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,)