Upload RotEx inference code, demos, documentation, and pretrained checkpoints

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo/illumination_variation_space.jpg filter=lfs diff=lfs merge=lfs -text
+demo/RotEx_Proposed.png filter=lfs diff=lfs merge=lfs -text
+demo/VLA_Traj.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,306 @@
----
-license: apache-2.0
----
+# [Official] RotEx: A Rotational Expert Network for Realistic Scene Synthesis Under Illumination Variation
+
+<p align="center">
+    📄 <b>Paper: coming soon</b>&nbsp;&nbsp; | &nbsp;&nbsp;
+    💻 <a href="https://github.com/val-utehy/RotEx"><b>Source Code</b></a>&nbsp;&nbsp; | &nbsp;&nbsp;
+    🤗 <a href="https://huggingface.co/val-utehy/RotEx"><b>Hugging Face</b></a>
+</p>
+
+<p align="center">
+    <img src="demo/RotEx_Proposed.png" alt="Overview of RotEx and its double-stream and single-stream transformer blocks" width="1100"/>
+</p>
+<p align="center">
+    <em>Architecture of RotEx: content and style streams, depth conditioning, and illumination-dependent expert routing.</em>
+</p>
+
+> The Hugging Face commands below describe the intended weight layout; they become usable once the checkpoints are uploaded.
+
+## Introduction
+
+**RotEx** is a transformer-based image-to-image framework for synthesizing scenes under varying illumination conditions. It combines content latents, depth information, style features, and a continuous time-of-day variable to generate illumination variations of an input scene.
+
+The model operates in the latent space of a FLUX autoencoder. DINOv2 supplies style features, Depth Anything V2 supplies depth maps, and time-conditioned mixture-of-experts modules adapt the synthesis process to the requested illumination. The release includes both multi-step inference and a distilled one-step inference path.
+
+### Key Features
+
+- **Continuous illumination control:** Specify a time-of-day value and map it to an angle through `phi = hour * 2π / 24`.
+- **Illumination-dependent experts:** Route style and image features through experts conditioned on the illumination angle.
+- **Content and depth conditioning:** Use an encoded content image and its estimated depth map during synthesis.
+- **Two inference modes:** Use a full checkpoint with a multi-step Heun solver, or a distilled checkpoint with a single model evaluation.
+- **Image and video processing:** Generate image sequences across hours or vary the target hour across the frames of a video.
+- **Reference-style conditioning:** Use a separate style image to condition illumination synthesis.
+
+## Demo
+
+### Illumination Variation
+
+<p align="center">
+    <img src="demo/illumination_variation_space.jpg" alt="Scene synthesis across the illumination variation space" width="1000"/>
+</p>
+<p align="center">
+    <em>Illustration of the illumination variation space.</em>
+</p>
+
+### Downstream Trajectory Comparison
+
+<p align="center">
+    <img src="demo/VLA_Traj.png" alt="AutoVLA trajectory comparison under different image translation methods" width="1100"/>
+</p>
+<p align="center">
+    <em>Qualitative comparison of AutoVLA rollouts using daylight reference images and images produced by different translation methods.</em>
+</p>
+
+## Installation and Requirements
+
+The supplied inference code uses CUDA and imports FlashAttention directly. Use a compatible NVIDIA GPU and a Linux environment with the CUDA toolkit required by your FlashAttention build.
+
+| Component | Requirement |
+| --- | --- |
+| Python | 3.10 or later |
+| PyTorch | 2.6.0 or later, as specified in `requirements.txt` |
+| torchvision | A version compatible with the installed PyTorch release; `requirements.txt` specifies 0.19.0 or later |
+| FlashAttention | 2.0.0 or later, with a build compatible with PyTorch, CUDA, and the GPU |
+| Default inference precision | `torch.bfloat16` |
+| Default image resolution | 384 × 384 |
+
+### 1. Clone the Repository
+
+```shell
+git clone https://github.com/val-utehy/RotEx.git
+cd RotEx
+```
+
+All commands below run from the directory containing `config.py`, `flux/`, and the inference scripts. If using this code inside the BenchRotEx workspace, that directory is `B/RotEx`.
+
+### 2. Create an Environment
+
+```shell
+conda create -n rotex python=3.10 -y
+conda activate rotex
+python -m pip install --upgrade pip
+```
+
+### 3. Install PyTorch and Dependencies
+
+Install a matching CUDA-enabled PyTorch and torchvision pair using the [official PyTorch installation instructions](https://pytorch.org/get-started/locally/). Ensure the selected versions satisfy the requirements above.
+
+After installing PyTorch:
+
+```shell
+python -m pip install packaging ninja psutil
+python -m pip install "flash-attn>=2.0.0" --no-build-isolation
+python -m pip install -r requirements.txt
+```
+
+See the [FlashAttention installation instructions](https://github.com/Dao-AILab/flash-attention#installation-and-features) for supported CUDA/GPU combinations and build requirements.
+
+Verify the environment:
+
+```shell
+python -c "import torch, flash_attn; print('PyTorch:', torch.__version__); print('CUDA:', torch.cuda.is_available()); print('BF16:', torch.cuda.is_bf16_supported())"
+```
+
+### 4. Prepare the Auxiliary Models
+
+The inference scripts load these pretrained components:
+
+| Component | Source | Local behavior |
+| --- | --- | --- |
+| FLUX autoencoder | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev), `ae.safetensors` | Loaded from `checkpoints/ae.safetensors`; downloaded there if missing |
+| DINOv2 Base | [facebook/dinov2-base](https://huggingface.co/facebook/dinov2-base) | Downloaded through Transformers and stored in the Hugging Face cache |
+| Depth Anything V2 Small | [depth-anything/Depth-Anything-V2-Small-hf](https://huggingface.co/depth-anything/Depth-Anything-V2-Small-hf) | Downloaded through Transformers and stored in the Hugging Face cache |
+
+An internet connection is needed for the first download. If access to an auxiliary model requires authentication or approval, complete the access process on its model page and log in with `hf auth login`.
+
+To download the autoencoder explicitly:
+
+```shell
+hf download black-forest-labs/FLUX.1-dev ae.safetensors --local-dir checkpoints
+```
+
+## Pretrained Models
+
+RotEx weights will be distributed through [Hugging Face](https://huggingface.co/val-utehy/RotEx). Download the checkpoint matching the inference mode you intend to run.
+
+| Checkpoint | Purpose | Expected local path |
+| --- | --- | --- |
+| Full model, best checkpoint | Multi-step image and video inference | `checkpoints/full/best.pt` |
+| Full model, final checkpoint | Alternative full-model checkpoint | `checkpoints/full/last.pt` |
+| Distilled model, best checkpoint | One-step inference | `distill_checkpoints/best_distill.pt` |
+
+Once the model repository is published with this directory structure, download the required weights from the project root:
+
+```shell
+hf download val-utehy/RotEx checkpoints/full/best.pt --local-dir .
+hf download val-utehy/RotEx distill_checkpoints/best_distill.pt --local-dir .
+```
+
+The final full-model checkpoint is optional:
+
+```shell
+hf download val-utehy/RotEx checkpoints/full/last.pt --local-dir .
+```
+
+The `hf` command is provided by `huggingface_hub`. See the [Hugging Face CLI guide](https://huggingface.co/docs/huggingface_hub/guides/cli) for download and authentication options.
+
+Additional placement instructions are available in [checkpoints/README.md](checkpoints/README.md), [checkpoints/full/README.md](checkpoints/full/README.md), and [distill_checkpoints/README.md](distill_checkpoints/README.md).
+
+> [!NOTE]
+> The scripts still contain paths from the original development environment. Pass explicit paths to the multi-step scripts and update the configuration constants in `infer_1step_style.py`. The model configuration is defined in `config.py`; these scripts do not use an `opt.pkl` file.
+
+## Usage Guide
+
+### 1. Model Training
+
+This release contains model components and inference scripts. A standalone training entry point and a distillation training entry point are not included in this directory.
+
+### 2. Choosing an Inference Script
+
+| Script | Checkpoint | Style input | Output |
+| --- | --- | --- | --- |
+| `infer_1step_style.py` | Distilled | One fixed reference image | One PNG per image and hour for each configured input/output pair; hours 1–24 by default |
+| `infer_multistep.py` | Full | Every image in a style directory | One labeled strip per content/style pair; hours 0–24 by default |
+| `infer_multistep_vid.py` | Full | One fixed reference image | One MP4 with the target hour interpolated across frames |
+
+All image-directory scripts read `.jpg`, `.jpeg`, and `.png` files directly inside the configured folders. They do not search subdirectories recursively.
+
+### 3. Multi-Step Image Inference
+
+Prepare separate content and style directories, for example:
+
+```text
+inputs/
+├── content/
+│   ├── scene_001.png
+│   └── scene_002.jpg
+└── style/
+    └── night_reference.jpg
+```
+
+Run:
+
+```shell
+python infer_multistep.py \
+    --checkpoint checkpoints/full/best.pt \
+    --content_dir inputs/content \
+    --style_dir inputs/style \
+    --output_dir outputs/multistep \
+    --denoise_steps 50 \
+    --image_size 384 \
+    --depth_size 48 \
+    --num_seeds 1 \
+    --device cuda \
+    --dtype bfloat16
+```
+
+Each content image is combined with every style image. Results are saved as `<content_stem>__sty_<style_stem>.png`. Each result is a horizontal strip containing the content image, style reference, depth visualization, and generated images for hours 0 through 24. Existing output files are skipped.
+
+| Argument | Default | Description |
+| --- | --- | --- |
+| `--denoise_steps` | `50` | Number of Heun integration steps; each step evaluates the model twice |
+| `--image_size` | `384` | Square image resolution used for content and style preprocessing |
+| `--depth_size` | `48` | Depth-map resolution supplied to the model |
+| `--num_seeds` | `1` | Number of independently sampled outputs to average per hour, in decoded image space |
+| `--device` | `cuda` | Device used for inference |
+| `--dtype` | `bfloat16` | Use `bfloat16` for the supplied default GPU path |
+
+The hour list is set inside `run_inference()` and is not exposed as a command-line argument. The legacy `--noise_alpha` argument is unused; initialization is sampled from Gaussian noise. The image solver clamps latents to `[-20, 20]` during integration and `[-15, 15]` before decoding.
+
+### 4. One-Step Image Inference with a Fixed Style
+
+Open `infer_1step_style.py` and configure the distilled checkpoint, style image, and dataset paths:
+
+```python
+CHECKPOINT = "distill_checkpoints/best_distill.pt"
+FIXED_STYLE_PATH = "inputs/style/night_reference.jpg"
+DATA_PAIRS = [
+    {"input": "inputs/content", "output": "outputs/1step_style"},
+]
+HOURS = list(range(1, 25))
+```
+
+Run:
+
+```shell
+python infer_1step_style.py
+```
+
+The reference style is encoded once and reused for all configured image folders. Outputs follow the `<image_stem>_hXX.png` naming convention. The script writes to the configured output paths and can overwrite existing images with the same names.
+
+### 5. Multi-Step Video Inference
+
+```shell
+python infer_multistep_vid.py \
+    --checkpoint checkpoints/full/best.pt \
+    --video_path inputs/video.mp4 \
+    --style_path inputs/style/night_reference.jpg \
+    --output_dir outputs/video \
+    --denoise_steps 50 \
+    --image_size 384 \
+    --depth_size 48 \
+    --max_height 384 \
+    --hour_start 0 \
+    --hour_end 24 \
+    --device cuda \
+    --dtype bfloat16 \
+    --no_audio
+```
+
+The script encodes the style reference once, reuses a base noise tensor across frames, and interpolates the target hour from `--hour_start` to `--hour_end` over the input video. The output is named `<video_stem>_chronflux.mp4` and uses the input video's reported frame rate.
+
+`--image_size` controls the square resolution used for model inference. `--max_height` controls output resizing, with output dimensions aligned to multiples of 16. The video solver uses latent clamps of `[-8, 8]` during integration and `[-6, 6]` before decoding, which differ from the image solver.
+
+The example produces a silent video. Audio transfer is optional and requires MoviePy, which is not included in `requirements.txt`; the script falls back to saving the video without audio if transfer fails.
+
+### 6. Illumination Control and Reproducibility
+
+The scripts map hours to the model's illumination angle as follows:
+
+```python
+phi = hour * (2 * math.pi / 24.0)
+```
+
+Hours 0 and 24 correspond to the endpoints of one full cycle. The image scripts use discrete hour lists, while the video script supplies intermediate hour values across frames. To render fractional hours in an image script, edit its hour list; the one-step filename formatting must also be adapted if the values are not integers.
+
+The image scripts sample new Gaussian noise for each output and do not expose a seed argument. Repeated runs can therefore produce different results. Increasing `--num_seeds` in the multi-step image script averages decoded outputs and increases computation; it is separate from increasing the number of solver steps.
+
+## Repository Structure
+
+```text
+RotEx/
+├── README.md
+├── config.py
+├── requirements.txt
+├── infer_1step_style.py
+├── infer_multistep.py
+├── infer_multistep_vid.py
+├── flux/                         # Model, autoencoder, encoders, and utilities
+├── demo/                         # Architecture and qualitative figures
+├── checkpoints/
+│   ├── README.md
+│   ├── ae.safetensors            # Download separately
+│   └── full/
+│       ├── README.md
+│       ├── best.pt               # Download from the RotEx model repository
+│       └── last.pt               # Optional alternative full checkpoint
+└── distill_checkpoints/
+    ├── README.md
+    └── best_distill.pt           # Download from the RotEx model repository
+```
+
+The Python implementation retains internal names such as `ChronFlux` and the `q1_paper` configuration preset. These are the identifiers used by the current inference scripts.
+
+## Citation
+
+Paper details and the official BibTeX entry will be added when available.
+
+## References
+
+- [FLUX](https://github.com/black-forest-labs/flux)
+- [DINOv2](https://github.com/facebookresearch/dinov2)
+- [Depth Anything V2](https://github.com/DepthAnything/Depth-Anything-V2)
+- [FlashAttention](https://github.com/Dao-AILab/flash-attention)
+
+## License
+
+A project license has not yet been included in this release. The license will be specified when the repository is published. Consult the respective upstream repositories and model cards for the terms governing pretrained components.

checkpoints/README.md

@@ -0,0 +1,53 @@
+# Pretrained Checkpoints
+
+RotEx model weights will be hosted on **Hugging Face**:
+
+[RotEx model repository](https://huggingface.co/val-utehy/RotEx)
+
+The commands below assume the uploaded files preserve the directory structure shown here.
+
+## Directory Layout
+
+```text
+checkpoints/
+├── README.md
+├── ae.safetensors
+└── full/
+    ├── README.md
+    ├── best.pt
+    └── last.pt
+```
+
+`best.pt` and `last.pt` are full-model checkpoints for `infer_multistep.py` and `infer_multistep_vid.py`. Use the distilled model in [`../distill_checkpoints/`](../distill_checkpoints/README.md) for `infer_1step_style.py`.
+
+## Download the Full Model
+
+Run from the project root, the directory containing `config.py`:
+
+```shell
+hf download val-utehy/RotEx checkpoints/full/best.pt --local-dir .
+```
+
+Optionally download the final full-model checkpoint:
+
+```shell
+hf download val-utehy/RotEx checkpoints/full/last.pt --local-dir .
+```
+
+Pass `--checkpoint checkpoints/full/best.pt` to either multi-step inference script. See [`full/README.md`](full/README.md) for examples.
+
+## Download the Autoencoder
+
+The FLUX autoencoder is distributed separately by its original provider:
+
+[black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
+
+```shell
+hf download black-forest-labs/FLUX.1-dev ae.safetensors --local-dir checkpoints
+```
+
+If authentication is required, complete the model's access process and run `hf auth login` first. The inference loader also attempts to download `ae.safetensors` automatically when it is absent from this directory.
+
+DINOv2 and Depth Anything V2 are loaded through Transformers into the Hugging Face cache; their weights do not need to be copied into this folder.
+
+See the [main README](../README.md) for installation and inference instructions.

checkpoints/full/README.md

@@ -0,0 +1,49 @@
+# Full-Model Weights
+
+Download full-model checkpoints from the [RotEx Hugging Face repository](https://huggingface.co/val-utehy/RotEx).
+
+Downloads become available after the files are uploaded.
+
+| File | Use |
+| --- | --- |
+| `best.pt` | Best full-model checkpoint |
+| `last.pt` | Final full-model checkpoint; optional alternative |
+
+Run from the project root:
+
+```shell
+hf download val-utehy/RotEx checkpoints/full/best.pt --local-dir .
+```
+
+To download the alternative checkpoint:
+
+```shell
+hf download val-utehy/RotEx checkpoints/full/last.pt --local-dir .
+```
+
+The expected local paths are `checkpoints/full/best.pt` and `checkpoints/full/last.pt`.
+
+## Image Inference
+
+```shell
+python infer_multistep.py \
+    --checkpoint checkpoints/full/best.pt \
+    --content_dir inputs/content \
+    --style_dir inputs/style \
+    --output_dir outputs/multistep
+```
+
+## Video Inference
+
+```shell
+python infer_multistep_vid.py \
+    --checkpoint checkpoints/full/best.pt \
+    --video_path inputs/video.mp4 \
+    --style_path inputs/style/night_reference.jpg \
+    --output_dir outputs/video \
+    --no_audio
+```
+
+These checkpoints are intended for the multi-step scripts. For one-step inference, download `distill_checkpoints/best_distill.pt` and follow the [distilled-checkpoint instructions](../../distill_checkpoints/README.md).
+
+See the [main README](../../README.md) for dependencies, auxiliary models, and inference options.

checkpoints/full/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b24f810ca16eb2a6d1ab5eed60d5ea0ba9ead1d0ebeb827e1daaf77305670c7b
+size 8396220961

checkpoints/full/last.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b24f810ca16eb2a6d1ab5eed60d5ea0ba9ead1d0ebeb827e1daaf77305670c7b
+size 8396220961

config.py

@@ -0,0 +1,118 @@
+"""
+Configuration system for enhanced RIFT model.
+"""
+from dataclasses import dataclass, field
+from typing import Literal
+
+@dataclass
+class EncoderConfig:
+    type: Literal["siglip", "dinov2", "evaclip"] = "dinov2"
+    model_size: Literal["small", "base", "large"] = "base"
+    output_dim: int = 512
+    depth_aware: bool = False
+
+@dataclass
+class LossConfig:
+    """
+    Configuration for loss functions and weights.
+    """
+    recon_weight: float = 10.0          # Charbonnier Loss
+    color_match_weight: float = 5.0     # ColorMatchingLoss
+    day_preserve_weight: float = 3.0    # DayPreservationLoss
+    moe_weight: float = 0.05            # MoE Load Balancing
+    brightness_cont_weight: float = 0.1
+    color_cont_weight: float = 0.1
+
+    use_charbonnier: bool = True
+
+    use_vgg: bool = True
+    vgg_weight: float = 0.5
+    # Tuning guide:
+    #   0.1  → subtle sharpening, safe
+    #   0.5  → noticeable edge/texture improvement  ← default
+
+    use_lpips: bool = False
+    lpips_weight: float = 0.0
+
+    use_l1: bool = False
+    use_gan: bool = False
+    gan_weight: float = 0.0
+
+    # Legacy (disabled)
+    use_smooth_transition: bool = False
+    smooth_weight: float = 0.0
+    use_content_preservation: bool = False
+    content_weight: float = 0.0
+    use_brightness_preservation: bool = False
+    brightness_weight: float = 0.0
+    min_brightness: float = 0.0
+    max_brightness: float = 0.0
+    charbonnier_eps: float = 1e-4
+
+@dataclass
+class AttentionConfig:
+    use_cross_attention: bool = True
+    use_adain: bool = True
+    use_depth_attention: bool = True
+    use_multi_scale_attention: bool = False
+
+@dataclass
+class MoEConfig:
+    enabled: bool = True
+    num_experts: int = 4
+    top_k: int = 2
+    temperature: float = 0.3
+    noise_std: float = 0.1
+    load_balance_coef: float = 0.01
+
+@dataclass
+class SamplingConfig:
+    method: Literal["euler", "heun", "consistency"] = "heun"
+    num_steps: int = 30
+    use_flow_matching: bool = True
+    use_consistency: bool = False
+    use_cfg: bool = True
+    cfg_scale: float = 7.5
+    sigma_min: float = 0.002
+    sigma_max: float = 80.0
+
+@dataclass
+class ModelConfig:
+    in_channels: int = 16
+    hidden_size: int = 1024
+    num_double_layers: int = 8
+    num_single_layers: int = 16
+    num_heads: int = 16
+    style_dim: int = 512
+    use_text: bool = False
+    text_dim: int = 1152
+    encoder: EncoderConfig = field(default_factory=EncoderConfig)
+    loss: LossConfig = field(default_factory=LossConfig)
+    attention: AttentionConfig = field(default_factory=AttentionConfig)
+    moe: MoEConfig = field(default_factory=MoEConfig)
+    sampling: SamplingConfig = field(default_factory=SamplingConfig)
+
+@dataclass
+class TrainingConfig:
+    content_dir: str = "/workspace/ChronFLUX/rafael_ds_ver2/train/day"
+    style_dir: str = "/workspace/ChronFLUX/rafael_ds_ver2/train/night"
+    image_size: int = 384
+    batch_size: int = 1
+    grad_accum_steps: int = 4
+    epochs: int = 43
+    lr_max: float = 1e-7
+    lr_min: float = 1e-5
+    weight_decay: float = 1e-4
+    save_interval_epochs: int = 1
+    device: str = "cuda"
+    dtype: str = "bfloat16"
+
+def get_q1_paper_config() -> ModelConfig:
+    return ModelConfig()
+
+PRESETS = {
+    "q1_paper": get_q1_paper_config,
+}
+
+def get_config(preset: str = "q1_paper") -> ModelConfig:
+    return PRESETS.get(preset, get_q1_paper_config)()

distill_checkpoints/README.md

@@ -0,0 +1,35 @@
+# Distilled One-Step Weights
+
+Download the distilled checkpoint from the [RotEx Hugging Face repository](https://huggingface.co/val-utehy/RotEx).
+
+Downloads become available after the checkpoint is uploaded.
+
+Run from the project root, the directory containing `config.py`:
+
+```shell
+hf download val-utehy/RotEx distill_checkpoints/best_distill.pt --local-dir .
+```
+
+The expected local path is:
+
+```text
+distill_checkpoints/best_distill.pt
+```
+
+## Usage
+
+Set the checkpoint constant in `infer_1step_style.py`:
+
+```python
+CHECKPOINT = "distill_checkpoints/best_distill.pt"
+```
+
+Configure the input/output paths in the script, then run:
+
+```shell
+python infer_1step_style.py
+```
+
+`infer_1step_style.py` requires a fixed reference image configured through `FIXED_STYLE_PATH` and input/output folders configured through `DATA_PAIRS`.
+
+The script also needs the FLUX autoencoder, DINOv2, and Depth Anything V2 described in the [main README](../README.md). Full-model weights for multi-step inference belong in [`../checkpoints/full/`](../checkpoints/full/README.md).

distill_checkpoints/best_distill.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee1b4ef1320479571288b9f00cbd4ddd89e6e7829e0f9d1d104546f821d487ab
+size 2798766921

flux/__init__.py

@@ -0,0 +1,13 @@
+try:
+    from ._version import (
+        version as __version__,
+        version_tuple,
+    )
+except ImportError:
+    __version__ = "unknown"
+    version_tuple = (0, 0, "unknown", "noinfo")
+
+from pathlib import Path
+
+PACKAGE = __package__.replace("_", "-")
+PACKAGE_ROOT = Path(__file__).parent

flux/chronflux.py

@@ -0,0 +1,330 @@
+import torch
+from torch import nn
+from einops import rearrange, repeat
+from flux.modules.layers import EmbedND, timestep_embedding, MLPEmbedder
+from flux.utils.math import attention, apply_rope
+from flux.modules.attention import AdaIN1d, DepthControlledAttention
+from flux.modules.moe import MoEStyleAdapter, MoEFeatureProcessor
+import torch.utils.checkpoint as checkpoint
+from typing import Optional, Dict, Any
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = int(hidden_features * 2 / 3)
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+        self.w3 = nn.Linear(hidden_features, out_features)
+        self.silu = nn.SiLU()
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.silu(x1) * x2
+        return self.w3(hidden)
+
+class AdaLNModulation(nn.Module):
+    def __init__(self, hidden_size, components=6):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(hidden_size, components * hidden_size, bias=True)
+        self.components = components
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, vec):
+        chunk = self.linear(self.silu(vec))
+        return chunk.chunk(self.components, dim=1)
+
+class QKNorm(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = nn.LayerNorm(dim, elementwise_affine=True, eps=1e-6)
+        self.key_norm = nn.LayerNorm(dim, elementwise_affine=True, eps=1e-6)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.query_norm(q), self.key_norm(k)
+
+class ChronFluxDoubleStreamBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_depth_attn=False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.use_depth_attn = use_depth_attn
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_qkv = nn.Linear(hidden_size, hidden_size * 3)
+        self.img_proj = nn.Linear(hidden_size, hidden_size)
+        self.sty_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.sty_qkv = nn.Linear(hidden_size, hidden_size * 3)
+        self.sty_proj = nn.Linear(hidden_size, hidden_size)
+        self.img_qk_norm = QKNorm(self.head_dim)
+        self.sty_qk_norm = QKNorm(self.head_dim)
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        self.sty_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.sty_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        self.adaLN = AdaLNModulation(hidden_size, components=12)
+        if use_depth_attn:
+            self.depth_attn = DepthControlledAttention(hidden_size, num_heads=num_heads)
+        self.use_adain = False
+
+    def set_adain(self, use_adain, style_dim):
+        self.use_adain = use_adain
+        if use_adain:
+            self.adain = AdaIN1d(self.hidden_size, style_dim)
+
+    def forward(self, img, sty, vec, pe, depth_map=None, style_vec=None):
+        if self.use_adain and style_vec is not None:
+            adain_style = style_vec
+            if style_vec.dim() == 3:
+                adain_style = style_vec.mean(dim=1)
+            img = self.adain(img, adain_style)
+
+        params = self.adaLN(vec)
+        (shift_i_a, scale_i_a, gate_i_a, shift_i_m, scale_i_m, gate_i_m,
+         shift_s_a, scale_s_a, gate_s_a, shift_s_m, scale_s_m, gate_s_m) = params
+
+        img_mod = self.img_norm1(img) * (1 + scale_i_a.unsqueeze(1)) + shift_i_a.unsqueeze(1)
+        sty_mod = self.sty_norm1(sty) * (1 + scale_s_a.unsqueeze(1)) + shift_s_a.unsqueeze(1)
+        qkv_i = self.img_qkv(img_mod)
+        qkv_s = self.sty_qkv(sty_mod)
+        q_i, k_i, v_i = rearrange(qkv_i, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q_s, k_s, v_s = rearrange(qkv_s, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q_i, k_i = self.img_qk_norm(q_i, k_i)
+        q_s, k_s = self.sty_qk_norm(q_s, k_s)
+        dtype = img.dtype
+        q_i, k_i, v_i = q_i.to(dtype), k_i.to(dtype), v_i.to(dtype)
+        q_s, k_s, v_s = q_s.to(dtype), k_s.to(dtype), v_s.to(dtype)
+        L_s_tokens = sty.shape[1]
+        pe_s = pe[:, :, :L_s_tokens, ...]
+        pe_i = pe[:, :, L_s_tokens:, ...]
+        q_i, k_i = apply_rope(q_i, k_i, pe_i)
+        q_s, k_s = apply_rope(q_s, k_s, pe_s)
+        q_i, k_i = q_i.to(dtype), k_i.to(dtype)
+        q_s, k_s = q_s.to(dtype), k_s.to(dtype)
+        q_i = rearrange(q_i, "B H L D -> B L H D")
+        k_i = rearrange(k_i, "B H L D -> B L H D")
+        v_i = rearrange(v_i, "B H L D -> B L H D")
+        q_s = rearrange(q_s, "B H L D -> B L H D")
+        k_s = rearrange(k_s, "B H L D -> B L H D")
+        v_s = rearrange(v_s, "B H L D -> B L H D")
+        full_k = torch.cat([k_s, k_i], dim=1)
+        full_v = torch.cat([v_s, v_i], dim=1)
+        img_attn_out = self.img_proj(attention(q_i, full_k, full_v))
+        sty_attn_out = self.sty_proj(attention(q_s, full_k, full_v))
+        img = img + gate_i_a.unsqueeze(1) * img_attn_out
+        sty = sty + gate_s_a.unsqueeze(1) * sty_attn_out
+        if self.use_depth_attn and depth_map is not None:
+            img = img + self.depth_attn(img, depth_map)
+        img_mod = self.img_norm2(img) * (1 + scale_i_m.unsqueeze(1)) + shift_i_m.unsqueeze(1)
+        img = img + gate_i_m.unsqueeze(1) * self.img_mlp(img_mod)
+        sty_mod = self.sty_norm2(sty) * (1 + scale_s_m.unsqueeze(1)) + shift_s_m.unsqueeze(1)
+        sty = sty + gate_s_m.unsqueeze(1) * self.sty_mlp(sty_mod)
+        return img, sty
+
+class ChronFluxSingleStreamBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_moe=False, moe_params: Optional[Dict[str, Any]] = None):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.use_moe = use_moe
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.qkv = nn.Linear(hidden_size, hidden_size * 3)
+        self.proj = nn.Linear(hidden_size, hidden_size)
+        self.qk_norm = QKNorm(self.head_dim)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        if use_moe:
+            params = moe_params or {}
+            num_experts = params.get('num_experts', 3) # Default safe value
+            # Extract only valid args for MoEFeatureProcessor
+            valid_keys = ['custom_centers', 'custom_widths']
+            kwargs = {k: params[k] for k in valid_keys if k in params}
+
+            self.moe_processor = MoEFeatureProcessor(hidden_size, num_experts=num_experts, **kwargs)
+            self.mlp = None
+        else:
+            self.mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+            self.moe_processor = None
+        self.adaLN = AdaLNModulation(hidden_size, components=6)
+
+    def forward(self, x, vec, pe, phi=None):
+        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaLN(vec)
+        x_mod = self.norm1(x) * (1 + scale_a.unsqueeze(1)) + shift_a.unsqueeze(1)
+        qkv = self.qkv(x_mod)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.qk_norm(q, k)
+        dtype = x.dtype
+        q, k, v = q.to(dtype), k.to(dtype), v.to(dtype)
+        q, k = apply_rope(q, k, pe)
+        q, k = q.to(dtype), k.to(dtype)
+        q = rearrange(q, "B H L D -> B L H D")
+        k = rearrange(k, "B H L D -> B L H D")
+        v = rearrange(v, "B H L D -> B L H D")
+        x_attn = self.proj(attention(q, k, v))
+        x = x + gate_a.unsqueeze(1) * x_attn
+        x_mod = self.norm2(x) * (1 + scale_m.unsqueeze(1)) + shift_m.unsqueeze(1)
+
+        aux_loss = 0.0
+        if self.use_moe and phi is not None:
+            moe_out, aux = self.moe_processor(x_mod, phi, return_aux=True)
+            x = x + gate_m.unsqueeze(1) * moe_out
+            if aux is not None:
+                aux_loss = aux['load_balance_loss']
+        else:
+            x = x + gate_m.unsqueeze(1) * self.mlp(x_mod)
+
+        return x, aux_loss
+
+class ChronFlux(nn.Module):
+    def __init__(self, in_channels=16, hidden_size=1024, num_double_layers=8,
+                 num_single_layers=16, num_heads=16, style_dim=512,
+                 use_text=False, text_dim=1152,
+                 use_moe=True, use_depth_attn=True, use_adain=False,
+                 moe_params: Optional[Dict[str, Any]] = None):
+        super().__init__()
+        self.in_channels = in_channels
+        self.gradient_checkpointing = False
+        self.use_text = use_text
+        self.style_dim = style_dim
+        self.use_moe = use_moe
+        self.use_depth_attn = use_depth_attn
+        self.use_adain = use_adain
+        self.img_in = nn.Linear(16 * 3, hidden_size)
+        self.depth_proj = nn.Conv2d(1, 16, kernel_size=1)
+        self.style_in = nn.Linear(style_dim, hidden_size)
+        self.default_style = nn.Parameter(torch.randn(1, style_dim))
+        self.time_mlp = nn.Sequential(
+            nn.Linear(256, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size)
+        )
+        if self.use_text:
+            self.txt_mlp = nn.Sequential(
+                nn.Linear(text_dim, hidden_size),
+                nn.SiLU(),
+                nn.Linear(hidden_size, hidden_size)
+            )
+            self.style_adapter = nn.Sequential(
+                nn.Linear(text_dim, hidden_size),
+                nn.SiLU(),
+                nn.Linear(hidden_size, style_dim)
+            )
+            nn.init.zeros_(self.style_adapter[-1].weight)
+            nn.init.zeros_(self.style_adapter[-1].bias)
+        self.phi_mlp = nn.Sequential(
+            nn.Linear(5, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size)
+        )
+
+        # Pass MoE params
+        if self.use_moe:
+            params = moe_params or {}
+            num_experts = params.get('num_experts', 3)
+            valid_keys = ['custom_centers', 'custom_widths']
+            kwargs = {k: params[k] for k in valid_keys if k in params}
+            self.moe_style_adapter = MoEStyleAdapter(style_dim, hidden_dim=style_dim * 2, num_experts=num_experts, **kwargs)
+
+        self.vector_in = MLPEmbedder(hidden_size, hidden_size)
+        self.pe_embedder = EmbedND(dim=hidden_size // num_heads, theta=10000, axes_dim=[16, 24, 24])
+        self.double_blocks = nn.ModuleList()
+        for _ in range(num_double_layers):
+            block = ChronFluxDoubleStreamBlock(hidden_size, num_heads, use_depth_attn=use_depth_attn)
+            block.set_adain(use_adain, style_dim)
+            self.double_blocks.append(block)
+        self.single_blocks = nn.ModuleList([
+            ChronFluxSingleStreamBlock(
+                hidden_size, num_heads,
+                use_moe=(use_moe and i >= num_single_layers // 2),
+                moe_params=moe_params
+            )
+            for i in range(num_single_layers)
+        ])
+        self.final_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.final_layer = nn.Linear(hidden_size, in_channels)
+
+    def enable_gradient_checkpointing(self):
+        self.gradient_checkpointing = True
+
+    def forward(self, x_t, content_latent, depth_map, style_vec, timesteps, phi, txt_emb=None):
+        B, C, H, W = x_t.shape
+        moe_loss = torch.tensor(0.0, device=x_t.device, dtype=x_t.dtype)
+
+        content_normalized = content_latent
+        depth_small = torch.nn.functional.interpolate(depth_map, size=(H, W), mode="nearest")
+        depth_emb = self.depth_proj(depth_small)
+        img_cat = torch.cat([x_t, content_normalized, depth_emb], dim=1)
+        img_flat = rearrange(img_cat, 'b c h w -> b (h w) c')
+        img = self.img_in(img_flat)
+
+        if style_vec is None:
+            style_vec = self.default_style.expand(B, -1).to(x_t.dtype)
+            if self.use_text and txt_emb is not None:
+                style_delta = self.style_adapter(txt_emb)
+                style_vec = style_vec + style_delta
+
+        if self.use_moe:
+            style_vec, aux = self.moe_style_adapter(style_vec, phi, return_aux=True)
+            if aux is not None:
+                moe_loss += aux['load_balance_loss']
+
+        style = self.style_in(style_vec)
+        if style.dim() == 2:
+            style = style.unsqueeze(1)
+        L_s = style.shape[1]
+
+        t_emb = timestep_embedding(timesteps, 256).to(x_t.device).to(x_t.dtype)
+        t_vec = self.time_mlp(t_emb)
+        phi_norm = (phi % (2 * torch.pi)) / (2 * torch.pi)
+        phi_input = torch.cat([
+            torch.sin(phi.unsqueeze(1)),
+            torch.cos(phi.unsqueeze(1)),
+            torch.sin(2 * phi.unsqueeze(1)),
+            torch.cos(2 * phi.unsqueeze(1)),
+            phi_norm.unsqueeze(1)
+        ], dim=1)
+        phi_vec = self.phi_mlp(phi_input)
+        vec_input = t_vec + phi_vec
+        if self.use_text and txt_emb is not None:
+            txt_vec = self.txt_mlp(txt_emb)
+            vec_input = vec_input + txt_vec
+        vec = self.vector_in(vec_input)
+
+        img_ids = torch.zeros(H, W, 3).to(x_t.device)
+        img_ids[..., 1] = img_ids[..., 1] + torch.arange(H, device=x_t.device)[:, None]
+        img_ids[..., 2] = img_ids[..., 2] + torch.arange(W, device=x_t.device)[None, :]
+        if self.training:
+            img_ids += torch.randn_like(img_ids) * 0.5
+        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=B)
+        pe = self.pe_embedder(img_ids)
+        dtype = pe.dtype
+        device = pe.device
+        pe_dim = pe.shape[-3]
+        style_pe = torch.zeros(B, 1, L_s, pe_dim, 2, 2, device=device, dtype=dtype)
+        full_pe = torch.cat([style_pe, pe], dim=2)
+
+        for block in self.double_blocks:
+            if self.gradient_checkpointing and self.training:
+                img, style = checkpoint.checkpoint(block, img, style, vec, full_pe, depth_map, style_vec, use_reentrant=False)
+            else:
+                img, style = block(img, style, vec, full_pe, depth_map, style_vec=style_vec)
+
+        x = torch.cat([style, img], dim=1)
+
+        for block in self.single_blocks:
+            if self.gradient_checkpointing and self.training:
+                x, block_loss = checkpoint.checkpoint(block, x, vec, full_pe, phi, use_reentrant=False)
+            else:
+                x, block_loss = block(x, vec, full_pe, phi)
+            moe_loss += block_loss
+
+        x = x[:, style.shape[1]:, :]
+        x = self.final_norm(x)
+        x = self.final_layer(x)
+        x = rearrange(x, 'b (h w) c -> b c h w', h=H, w=W)
+
+        return x, moe_loss

flux/modules/adaptive_blocks.py

@@ -0,0 +1,253 @@
+"""
+Adaptive Flux Blocks with enhanced AdaLN integration.
+These blocks serve as building blocks for the SwiftFlux architecture with
+advanced style transfer control and optional MoE-compatible components.
+Author: https://github.com/datnguyen-tien204
+"""
+import torch
+from torch import nn
+from einops import rearrange
+from flux.modules.normalization import AdaLNModulation
+from flux.utils.math import attention, apply_rope
+from flux.utils.logger_utils import ConsoleLogger
+
+
+# Constants
+SWIGLU_HIDDEN_RATIO = 2 / 3  # Hidden dimension ratio for SwiGLU
+
+
+class QKNorm(nn.Module):
+    """Query-Key normalization for stable attention"""
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = nn.LayerNorm(dim, elementwise_affine=True, eps=1e-6)
+        self.key_norm = nn.LayerNorm(dim, elementwise_affine=True, eps=1e-6)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.query_norm(q), self.key_norm(k)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation for efficient FFN"""
+    def __init__(self, in_features, hidden_features=None, out_features=None):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        hidden_features = int(hidden_features * SWIGLU_HIDDEN_RATIO)
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+        self.w3 = nn.Linear(hidden_features, out_features)
+        self.silu = nn.SiLU()
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.silu(x1) * x2
+        return self.w3(hidden)
+
+
+class AdaptiveFluxBlock(nn.Module):
+    """
+    Base Adaptive Flux Block with AdaLN integration.
+    Serves as foundation for both double-stream and single-stream variants.
+    
+    Features:
+    - AdaLN for layer-wise style and conditioning control
+    - QKNorm for attention stability
+    - SwiGLU for efficient feed-forward
+    - Optional sparse activation support for MoE compatibility
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_sparse=False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.use_sparse = use_sparse
+        # Note: adaLN should be initialized in subclasses
+        
+    def apply_adaptive_norm(self, x, scale, shift):
+        """Apply AdaLN normalization with scale and shift"""
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class AdaptiveDoubleStreamBlock(AdaptiveFluxBlock):
+    """
+    Adaptive Double-Stream Block with enhanced AdaLN.
+    Processes image and style streams separately with cross-attention.
+    
+    Improvements:
+    - Dedicated AdaLN for both streams
+    - Enhanced normalization for better style consistency
+    - Support for sparse activation patterns
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_sparse=False):
+        super().__init__(hidden_size, num_heads, mlp_ratio, use_sparse)
+        
+        # Image stream components
+        self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_qkv = nn.Linear(hidden_size, hidden_size * 3)
+        self.img_proj = nn.Linear(hidden_size, hidden_size)
+        self.img_qk_norm = QKNorm(self.head_dim)
+        
+        self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.img_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        
+        # Style stream components
+        self.sty_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.sty_qkv = nn.Linear(hidden_size, hidden_size * 3)
+        self.sty_proj = nn.Linear(hidden_size, hidden_size)
+        self.sty_qk_norm = QKNorm(self.head_dim)
+        
+        self.sty_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.sty_mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        
+        # AdaLN with 12 components (6 for img, 6 for style)
+        self.adaLN = AdaLNModulation(hidden_size, components=12)
+    
+    def forward(self, img, sty, vec, pe):
+        """
+        Forward pass with dual-stream processing
+        
+        Args:
+            img: Image features [B, L_img, D]
+            sty: Style features [B, L_sty, D]
+            vec: Conditioning vector [B, D]
+            pe: Position embeddings
+        
+        Returns:
+            Updated (img, sty) features
+        """
+        # Generate all modulation parameters
+        params = self.adaLN(vec)
+        (shift_i_a, scale_i_a, gate_i_a, shift_i_m, scale_i_m, gate_i_m,
+         shift_s_a, scale_s_a, gate_s_a, shift_s_m, scale_s_m, gate_s_m) = params
+        
+        # Image stream - attention
+        img_mod = self.apply_adaptive_norm(self.img_norm1(img), scale_i_a, shift_i_a)
+        sty_mod = self.apply_adaptive_norm(self.sty_norm1(sty), scale_s_a, shift_s_a)
+        
+        qkv_i = self.img_qkv(img_mod)
+        qkv_s = self.sty_qkv(sty_mod)
+        
+        q_i, k_i, v_i = rearrange(qkv_i, "B L (K H D) -> K B L H D", K=3, H=self.num_heads)
+        q_s, k_s, v_s = rearrange(qkv_s, "B L (K H D) -> K B L H D", K=3, H=self.num_heads)
+        
+        q_i, k_i = self.img_qk_norm(q_i, k_i)
+        q_s, k_s = self.sty_qk_norm(q_s, k_s)
+        
+        # Preserve dtype for Flash Attention
+        dtype = img.dtype
+        q_i, k_i, v_i = q_i.to(dtype), k_i.to(dtype), v_i.to(dtype)
+        q_s, k_s, v_s = q_s.to(dtype), k_s.to(dtype), v_s.to(dtype)
+        
+        # Apply RoPE if position embeddings provided
+        if pe is not None:
+            pe = pe.transpose(1, 2)
+            pe_s = pe[:, :1, ...]
+            pe_i = pe[:, 1:, ...]
+            
+            q_i, k_i = apply_rope(q_i, k_i, pe_i)
+            q_s, k_s = apply_rope(q_s, k_s, pe_s)
+            
+            q_i, k_i = q_i.to(dtype), k_i.to(dtype)
+            q_s, k_s = q_s.to(dtype), k_s.to(dtype)
+        
+        # Cross-attention between streams
+        full_k = torch.cat([k_s, k_i], dim=1)
+        full_v = torch.cat([v_s, v_i], dim=1)
+        
+        img_attn_out = self.img_proj(attention(q_i, full_k, full_v))
+        sty_attn_out = self.sty_proj(attention(q_s, full_k, full_v))
+        
+        # Apply gating and residual
+        img = img + gate_i_a.unsqueeze(1) * img_attn_out
+        sty = sty + gate_s_a.unsqueeze(1) * sty_attn_out
+        
+        # MLP blocks with AdaLN
+        img_mod = self.apply_adaptive_norm(self.img_norm2(img), scale_i_m, shift_i_m)
+        img = img + gate_i_m.unsqueeze(1) * self.img_mlp(img_mod)
+        
+        sty_mod = self.apply_adaptive_norm(self.sty_norm2(sty), scale_s_m, shift_s_m)
+        sty = sty + gate_s_m.unsqueeze(1) * self.sty_mlp(sty_mod)
+        
+        return img, sty
+
+
+class AdaptiveSingleStreamBlock(AdaptiveFluxBlock):
+    """
+    Adaptive Single-Stream Block with enhanced AdaLN.
+    Processes unified features with self-attention.
+    
+    Improvements:
+    - Enhanced AdaLN control for unified stream
+    - Better style propagation through the network
+    - Optional sparse activation for MoE
+    """
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, use_sparse=False):
+        super().__init__(hidden_size, num_heads, mlp_ratio, use_sparse)
+        
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.qkv = nn.Linear(hidden_size, hidden_size * 3)
+        self.proj = nn.Linear(hidden_size, hidden_size)
+        self.qk_norm = QKNorm(self.head_dim)
+        
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = SwiGLU(hidden_size, int(hidden_size * mlp_ratio))
+        
+        # AdaLN with 6 components (scale, shift, gate for attention and MLP)
+        self.adaLN = AdaLNModulation(hidden_size, components=6)
+    
+    def forward(self, x, vec, pe):
+        """
+        Forward pass with single-stream processing
+        
+        Args:
+            x: Input features [B, L, D]
+            vec: Conditioning vector [B, D]
+            pe: Position embeddings
+        
+        Returns:
+            Updated features
+        """
+        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaLN(vec)
+        
+        # Attention with AdaLN
+        x_mod = self.apply_adaptive_norm(self.norm1(x), scale_a, shift_a)
+        qkv = self.qkv(x_mod)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B L H D", K=3, H=self.num_heads)
+        q, k = self.qk_norm(q, k)
+        
+        dtype = x.dtype
+        q, k, v = q.to(dtype), k.to(dtype), v.to(dtype)
+        
+        # Apply RoPE if position embeddings provided
+        if pe is not None:
+            pe_reshaped = pe.transpose(1, 2)
+            q, k = apply_rope(q, k, pe_reshaped)
+            q, k = q.to(dtype), k.to(dtype)
+        
+        x_attn = self.proj(attention(q, k, v))
+        x = x + gate_a.unsqueeze(1) * x_attn
+        
+        # MLP with AdaLN
+        x_mod = self.apply_adaptive_norm(self.norm2(x), scale_m, shift_m)
+        x = x + gate_m.unsqueeze(1) * self.mlp(x_mod)
+        
+        return x
+
+
+# For backward compatibility, create aliases
+SwiftDoubleStreamBlock = AdaptiveDoubleStreamBlock
+SwiftSingleStreamBlock = AdaptiveSingleStreamBlock
+
+
+__all__ = [
+    'AdaptiveFluxBlock',
+    'AdaptiveDoubleStreamBlock',
+    'AdaptiveSingleStreamBlock',
+    'SwiftDoubleStreamBlock',
+    'SwiftSingleStreamBlock',
+    'QKNorm',
+    'SwiGLU',
+]

flux/modules/advanced_encoders.py

@@ -0,0 +1,256 @@
+"""
+Advanced encoder modules for enhanced feature extraction.
+Supports DINOv2, EVA-CLIP, and other state-of-the-art vision encoders.
+Author: https://github.com/datnguyen-tien204
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel, AutoImageProcessor, AutoImageProcessor, AutoModelForDepthEstimation
+from typing import Optional, Literal
+from flux.utils.logger_utils import ConsoleLogger
+
+class DepthEstimator:
+    def __init__(self, device):
+        self.device = device
+        ConsoleLogger.info("Loading Depth Anything V2 Small...")
+        self.proc = AutoImageProcessor.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf")
+        self.model = AutoModelForDepthEstimation.from_pretrained("depth-anything/Depth-Anything-V2-Small-hf").to(device).eval()
+
+    @torch.no_grad()
+    def predict(self, images, size=384):
+        imgs_255 = ((images * 0.5 + 0.5) * 255).clamp(0, 255).to(torch.uint8)
+        inputs = self.proc(images=imgs_255, return_tensors="pt", size={"height": 518, "width": 518}).to(self.device)
+        depth = self.model(**inputs).predicted_depth
+        depth = F.interpolate(depth.unsqueeze(1), size=(size, size), mode="bicubic")
+        for i in range(depth.shape[0]):
+            depth[i] = (depth[i] - depth[i].min()) / (depth[i].max() - depth[i].min() + 1e-6)
+        return depth
+
+class DINOv2Encoder(nn.Module):
+    """
+    DINOv2 encoder for enhanced feature extraction.
+    Better semantic understanding and structural preservation.
+    """
+    def __init__(self, device: str, output_dim: int = 512, 
+                 model_size: Literal["small", "base", "large"] = "base"):
+        super().__init__()
+        self.device = device
+        self.output_dim = output_dim
+        
+        # Select model based on size
+        model_names = {
+            "small": "facebook/dinov2-small",
+            "base": "facebook/dinov2-base",
+            "large": "facebook/dinov2-large"
+        }
+        model_name = model_names[model_size]
+        
+        print(f"Loading DINOv2 encoder: {model_name}")
+        self.model = AutoModel.from_pretrained(model_name).to(device)
+        self.model.eval()
+        
+        # Get hidden size from model config
+        hidden_size = self.model.config.hidden_size
+        
+        # Projection layer to match output_dim
+        self.proj = nn.Linear(hidden_size, output_dim, bias=False).to(device)
+        
+    def forward(self, x: torch.Tensor, return_spatial: bool = False) -> torch.Tensor:
+        """
+        Args:
+            x: Input images [B, 3, H, W] in range [-1, 1]
+            return_spatial: Whether to return all patch tokens [B, L, D]
+        Returns:
+            features: [B, output_dim] or [B, L, output_dim]
+        """
+        # Normalize to [0, 1] for DINOv2
+        x_normalized = (x + 1.0) / 2.0
+        
+        # Resize to 224x224 (DINOv2 standard input size)
+        x_resized = F.interpolate(x_normalized, size=(224, 224), 
+                                 mode='bilinear', align_corners=False)
+        
+        with torch.no_grad():
+            outputs = self.model(pixel_values=x_resized, return_dict=True)
+            if return_spatial:
+                # Return all patch tokens (excluding CLS if possible, though Flux handle sequences well)
+                features = outputs.last_hidden_state
+            else:
+                # Use CLS token (first token) as global feature
+                features = outputs.last_hidden_state[:, 0]
+            
+        return self.proj(features)
+
+
+class EVACLIPEncoder(nn.Module):
+    """
+    EVA-CLIP encoder for enhanced vision-language feature extraction.
+    State-of-the-art performance on visual understanding tasks.
+    """
+    def __init__(self, device: str, output_dim: int = 512,
+                 model_size: Literal["base", "large"] = "large"):
+        super().__init__()
+        self.device = device
+        self.output_dim = output_dim
+        
+        # EVA-CLIP models
+        model_names = {
+            "base": "BAAI/EVA-CLIP-8B",
+            "large": "BAAI/EVA-CLIP-18B"
+        }
+        
+        # Note: Using a more accessible model for now
+        # In production, use actual EVA-CLIP if available
+        model_name = "laion/CLIP-ViT-H-14-laion2B-s32B-b79K"
+        
+        print(f"Loading EVA-CLIP style encoder: {model_name}")
+        try:
+            from transformers import CLIPVisionModel
+            self.model = CLIPVisionModel.from_pretrained(model_name).to(device)
+        except Exception as e:
+            print(f"Warning: Could not load {model_name}, falling back to openai/clip-vit-large-patch14")
+            from transformers import CLIPVisionModel
+            self.model = CLIPVisionModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
+        
+        self.model.eval()
+        
+        # Get hidden size
+        hidden_size = self.model.config.hidden_size
+        
+        # Projection layer
+        self.proj = nn.Linear(hidden_size, output_dim, bias=False).to(device)
+        
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input images [B, 3, H, W] in range [-1, 1]
+        Returns:
+            features: [B, output_dim]
+        """
+        # Normalize to [0, 1]
+        x_normalized = (x + 1.0) / 2.0
+        
+        # Resize to 224x224
+        x_resized = F.interpolate(x_normalized, size=(224, 224),
+                                 mode='bilinear', align_corners=False)
+        
+        with torch.no_grad():
+            outputs = self.model(pixel_values=x_resized, return_dict=True)
+            features = outputs.pooler_output
+            
+        return self.proj(features)
+
+
+class DepthAwareEncoder(nn.Module):
+    """
+    Depth-aware feature encoder that combines visual features with depth information.
+    """
+    def __init__(self, base_encoder: nn.Module, depth_channels: int = 1, 
+                 output_dim: int = 512):
+        super().__init__()
+        self.base_encoder = base_encoder
+        self.output_dim = output_dim
+        
+        # Depth processing branch
+        self.depth_conv = nn.Sequential(
+            nn.Conv2d(depth_channels, 32, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d(1)
+        )
+        
+        # Fusion layer
+        self.fusion = nn.Sequential(
+            nn.Linear(output_dim + 64, output_dim),
+            nn.ReLU(inplace=True),
+            nn.Linear(output_dim, output_dim)
+        )
+        
+    def forward(self, x: torch.Tensor, depth: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: Input images [B, 3, H, W]
+            depth: Depth maps [B, 1, H, W] (optional)
+        Returns:
+            features: [B, output_dim]
+        """
+        # Get base visual features
+        visual_features = self.base_encoder(x)
+        
+        if depth is None:
+            return visual_features
+            
+        # Process depth
+        depth_features = self.depth_conv(depth)
+        depth_features = depth_features.view(depth_features.size(0), -1)
+        
+        # Concatenate and fuse
+        combined = torch.cat([visual_features, depth_features], dim=1)
+        fused_features = self.fusion(combined)
+        
+        return fused_features
+
+
+class EncoderFactory:
+    """
+    Factory class for creating different types of encoders.
+    """
+    @staticmethod
+    def create_encoder(
+        encoder_type: Literal["siglip", "dinov2", "evaclip"],
+        device: str,
+        output_dim: int = 512,
+        depth_aware: bool = False,
+        **kwargs
+    ) -> nn.Module:
+        """
+        Create an encoder of the specified type.
+        
+        Args:
+            encoder_type: Type of encoder to create
+            device: Device to load the model on
+            output_dim: Output dimension of the encoder
+            depth_aware: Whether to wrap encoder with depth-aware processing
+            **kwargs: Additional arguments for specific encoders
+        
+        Returns:
+            encoder: Configured encoder module
+        """
+        if encoder_type == "siglip":
+            from transformers import SiglipVisionModel
+            
+            class SigLIPStyleEncoder(nn.Module):
+                def __init__(self, device, output_dim=512):
+                    super().__init__()
+                    model_name = "google/siglip-so400m-patch14-384"
+                    self.model = SiglipVisionModel.from_pretrained(model_name).to(device)
+                    self.model.eval()
+                    self.proj = nn.Linear(self.model.config.hidden_size, output_dim, bias=False).to(device)
+                
+                def forward(self, x):
+                    x_resized = F.interpolate(x, size=(384, 384), mode='bilinear', align_corners=False)
+                    with torch.no_grad():
+                        outputs = self.model(pixel_values=x_resized)
+                        features = outputs.last_hidden_state.mean(dim=1)
+                    return self.proj(features)
+            
+            encoder = SigLIPStyleEncoder(device, output_dim)
+            
+        elif encoder_type == "dinov2":
+            model_size = kwargs.get("model_size", "base")
+            encoder = DINOv2Encoder(device, output_dim, model_size)
+            
+        elif encoder_type == "evaclip":
+            model_size = kwargs.get("model_size", "large")
+            encoder = EVACLIPEncoder(device, output_dim, model_size)
+            
+        else:
+            raise ValueError(f"Unknown encoder type: {encoder_type}")
+        
+        # Wrap with depth-aware processing if requested
+        if depth_aware:
+            encoder = DepthAwareEncoder(encoder, depth_channels=1, output_dim=output_dim)
+        
+        return encoder

flux/modules/attention.py

@@ -0,0 +1,332 @@
+"""
+Advanced attention mechanisms for better semantic and style alignment.
+Includes Cross-Attention, AdaIN, and Depth-Controlled Attention.
+Author: https://github.com/datnguyen-tien204
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from typing import Optional
+
+
+class AdaptiveInstanceNorm2d(nn.Module):
+    """
+    Adaptive Instance Normalization (AdaIN) for style transfer.
+    Aligns mean and variance of content features to match style features.
+    """
+    def __init__(self, num_features: int, style_dim: int):
+        """
+        Args:
+            num_features: Number of feature channels
+            style_dim: Dimension of style vector
+        """
+        super().__init__()
+        self.num_features = num_features
+        
+        # Style modulation parameters
+        self.style_scale = nn.Linear(style_dim, num_features)
+        self.style_shift = nn.Linear(style_dim, num_features)
+        
+        # Initialize to identity transformation
+        nn.init.constant_(self.style_scale.weight, 0)
+        nn.init.constant_(self.style_scale.bias, 1)
+        nn.init.constant_(self.style_shift.weight, 0)
+        nn.init.constant_(self.style_shift.bias, 0)
+    
+    def forward(self, x: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input features [B, C, H, W]
+            style: Style vector [B, style_dim]
+        Returns:
+            normalized: Style-adapted features [B, C, H, W]
+        """
+        # Instance normalization
+        mean = x.mean(dim=[2, 3], keepdim=True)
+        std = x.std(dim=[2, 3], keepdim=True) + 1e-5
+        normalized = (x - mean) / std
+        
+        # Style modulation
+        scale = self.style_scale(style).unsqueeze(-1).unsqueeze(-1)
+        shift = self.style_shift(style).unsqueeze(-1).unsqueeze(-1)
+        
+        return normalized * scale + shift
+
+
+class AdaIN1d(nn.Module):
+    """
+    AdaIN for 1D features (sequence data).
+    """
+    def __init__(self, num_features: int, style_dim: int):
+        super().__init__()
+        self.num_features = num_features
+
+        self.norm = nn.LayerNorm(num_features, elementwise_affine=False)
+        self.style_scale = nn.Linear(style_dim, num_features)
+        self.style_shift = nn.Linear(style_dim, num_features)
+
+        nn.init.constant_(self.style_scale.weight, 0)
+        nn.init.constant_(self.style_scale.bias, 1)
+        nn.init.constant_(self.style_shift.weight, 0)
+        nn.init.constant_(self.style_shift.bias, 0)
+
+    def forward(self, x: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input features [B, L, C]
+            style: Style vector [B, style_dim] OR [B, L_s, style_dim]
+        Returns:
+            normalized: Style-adapted features [B, L, C]
+        """
+        if style.dim() == 3:
+            style = style.mean(dim=1)
+
+        normalized = self.norm(x)
+        scale = self.style_scale(style).unsqueeze(1)
+        shift = self.style_shift(style).unsqueeze(1)
+
+        return normalized * scale + shift
+
+
+
+class CrossAttention(nn.Module):
+    """
+    Cross-Attention mechanism for semantic alignment between content and style.
+    """
+    def __init__(self, query_dim: int, key_dim: int, num_heads: int = 8, 
+                 dropout: float = 0.0):
+        """
+        Args:
+            query_dim: Dimension of query features
+            key_dim: Dimension of key/value features
+            num_heads: Number of attention heads
+            dropout: Dropout probability
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = query_dim // num_heads
+        assert query_dim % num_heads == 0, "query_dim must be divisible by num_heads"
+        
+        self.query = nn.Linear(query_dim, query_dim)
+        self.key = nn.Linear(key_dim, query_dim)
+        self.value = nn.Linear(key_dim, query_dim)
+        
+        self.proj = nn.Linear(query_dim, query_dim)
+        self.dropout = nn.Dropout(dropout)
+    
+    def forward(self, query: torch.Tensor, key: torch.Tensor, 
+                value: Optional[torch.Tensor] = None,
+                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            query: Query features [B, L_q, query_dim]
+            key: Key features [B, L_k, key_dim]
+            value: Value features [B, L_k, key_dim] (if None, uses key)
+            mask: Attention mask [B, L_q, L_k] (optional)
+        Returns:
+            output: Attended features [B, L_q, query_dim]
+        """
+        if value is None:
+            value = key
+        
+        B, L_q, _ = query.shape
+        L_k = key.shape[1]
+        
+        # Linear projections
+        q = self.query(query)
+        k = self.key(key)
+        v = self.value(value)
+        
+        # Reshape for multi-head attention
+        q = rearrange(q, 'b l (h d) -> b h l d', h=self.num_heads)
+        k = rearrange(k, 'b l (h d) -> b h l d', h=self.num_heads)
+        v = rearrange(v, 'b l (h d) -> b h l d', h=self.num_heads)
+        
+        # Scaled dot-product attention
+        scale = self.head_dim ** -0.5
+        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
+        
+        if mask is not None:
+            attn = attn.masked_fill(mask.unsqueeze(1) == 0, float('-inf'))
+        
+        attn = F.softmax(attn, dim=-1)
+        attn = self.dropout(attn)
+        
+        # Apply attention to values
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h l d -> b l (h d)')
+        
+        # Output projection
+        out = self.proj(out)
+        
+        return out
+
+
+class DepthControlledAttention(nn.Module):
+    """
+    Depth-controlled attention mechanism for structure-aware feature learning.
+    Uses depth information to modulate attention patterns.
+    """
+    def __init__(self, feature_dim: int, depth_dim: int = 64, num_heads: int = 8):
+        """
+        Args:
+            feature_dim: Dimension of input features
+            depth_dim: Dimension of depth features
+            num_heads: Number of attention heads
+        """
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = feature_dim // num_heads
+        
+        # Process depth information
+        self.depth_encoder = nn.Sequential(
+            nn.Conv2d(1, depth_dim, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(depth_dim, depth_dim, kernel_size=3, padding=1)
+        )
+        
+        # Depth-modulated attention
+        self.qkv = nn.Linear(feature_dim, feature_dim * 3)
+        self.depth_gate = nn.Linear(depth_dim, feature_dim)
+        self.proj = nn.Linear(feature_dim, feature_dim)
+        
+    def forward(self, x: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input features [B, L, C]
+            depth: Depth map [B, 1, H, W]
+        Returns:
+            output: Depth-modulated features [B, L, C]
+        """
+        B, L, C = x.shape
+        
+        # Encode depth
+        depth_feat = self.depth_encoder(depth)  # [B, depth_dim, H, W]
+        H, W = depth_feat.shape[-2:]
+        depth_feat = rearrange(depth_feat, 'b c h w -> b (h w) c')
+        
+        # Generate Q, K, V
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, 'b l (three h d) -> three b h l d', 
+                           three=3, h=self.num_heads)
+        
+        # Compute attention
+        scale = self.head_dim ** -0.5
+        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
+        
+        # Modulate attention with depth
+        if depth_feat.shape[1] == L:
+            depth_gate = self.depth_gate(depth_feat)
+            depth_gate = rearrange(depth_gate, 'b l (h d) -> b h l d', h=self.num_heads)
+            depth_mask = torch.sigmoid(depth_gate.mean(dim=-1, keepdim=True))
+            attn = attn * depth_mask
+        
+        attn = F.softmax(attn, dim=-1)
+        
+        # Apply attention
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h l d -> b l (h d)')
+        out = self.proj(out)
+        
+        return out
+
+
+class StyleAttentionBlock(nn.Module):
+    """
+    Combined attention block with AdaIN and Cross-Attention.
+    """
+    def __init__(self, feature_dim: int, style_dim: int, num_heads: int = 8):
+        super().__init__()
+        
+        # Layer normalization
+        self.norm1 = nn.LayerNorm(feature_dim)
+        self.norm2 = nn.LayerNorm(feature_dim)
+        
+        # Cross-attention for style conditioning
+        self.cross_attn = CrossAttention(feature_dim, style_dim, num_heads)
+        
+        # AdaIN for style transfer
+        self.adain = AdaIN1d(feature_dim, style_dim)
+        
+        # Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(feature_dim, feature_dim * 4),
+            nn.GELU(),
+            nn.Linear(feature_dim * 4, feature_dim)
+        )
+    
+    def forward(self, x: torch.Tensor, style: torch.Tensor, 
+                style_features: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Args:
+            x: Input features [B, L, C]
+            style: Style vector [B, style_dim]
+            style_features: Optional style feature sequence [B, L_s, style_dim]
+        Returns:
+            output: Style-conditioned features [B, L, C]
+        """
+        # Cross-attention with style features
+        if style_features is not None:
+            attn_out = self.cross_attn(self.norm1(x), style_features)
+            x = x + attn_out
+        
+        # AdaIN style transfer
+        x = self.adain(x, style)
+        
+        # Feed-forward
+        x = x + self.ffn(self.norm2(x))
+        
+        return x
+
+
+class MultiScaleAttention(nn.Module):
+    """
+    Multi-scale attention for capturing features at different resolutions.
+    """
+    def __init__(self, feature_dim: int, num_scales: int = 3, num_heads: int = 8):
+        super().__init__()
+        self.num_scales = num_scales
+        
+        self.attentions = nn.ModuleList([
+            nn.MultiheadAttention(feature_dim, num_heads, batch_first=True)
+            for _ in range(num_scales)
+        ])
+        
+        self.fusion = nn.Linear(feature_dim * num_scales, feature_dim)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Input features [B, L, C]
+        Returns:
+            output: Multi-scale attended features [B, L, C]
+        """
+        B, L, C = x.shape
+        
+        outputs = []
+        for i, attn in enumerate(self.attentions):
+            # Downsample for different scales
+            if i > 0:
+                scale = 2 ** i
+                # Simple downsampling by reshaping
+                L_down = L // scale
+                if L_down > 0:
+                    x_down = x[:, ::scale, :]
+                else:
+                    x_down = x
+                out, _ = attn(x_down, x_down, x_down)
+                # Upsample back
+                if out.shape[1] != L:
+                    out = F.interpolate(out.transpose(1, 2), size=L, mode='linear', 
+                                      align_corners=False).transpose(1, 2)
+            else:
+                out, _ = attn(x, x, x)
+            
+            outputs.append(out)
+        
+        # Concatenate and fuse
+        multi_scale = torch.cat(outputs, dim=-1)
+        fused = self.fusion(multi_scale)
+        
+        return fused

flux/modules/autoencoder.py

@@ -0,0 +1,251 @@
+# Flux project: https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/autoencoder.py
+from dataclasses import dataclass
+import torch
+from einops import rearrange
+from torch import Tensor, nn
+
+@dataclass
+class AutoEncoderParams:
+    resolution: int
+    in_channels: int
+    ch: int
+    out_ch: int
+    ch_mult: list[int]
+    num_res_blocks: int
+    z_channels: int
+    scale_factor: float
+    shift_factor: float
+
+def swish(x: Tensor) -> Tensor:
+    return x * torch.sigmoid(x)
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1)
+
+    def attention(self, h_: Tensor) -> Tensor:
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+        b, c, h, w = q.shape
+        q = rearrange(q, "b c h w -> b 1 (h w) c").contiguous()
+        k = rearrange(k, "b c h w -> b 1 (h w) c").contiguous()
+        v = rearrange(v, "b c h w -> b 1 (h w) c").contiguous()
+        h_ = nn.functional.scaled_dot_product_attention(q, k, v)
+        return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x + self.proj_out(self.attention(x))
+
+class ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if self.in_channels != self.out_channels:
+            self.nin_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = swish(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = swish(h)
+        h = self.conv2(h)
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: Tensor):
+        pad = (0, 1, 0, 1)
+        x = nn.functional.pad(x, pad, mode="constant", value=0)
+        x = self.conv(x)
+        return x
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor):
+        x = nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+class Encoder(nn.Module):
+    def __init__(self, resolution: int, in_channels: int, ch: int, ch_mult: list[int], num_res_blocks: int, z_channels: int):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.conv_in = nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1)
+        curr_res = resolution
+        in_ch_mult = (1,) + tuple(ch_mult)
+        self.in_ch_mult = in_ch_mult
+        self.down = nn.ModuleList()
+        block_in = self.ch
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch * in_ch_mult[i_level]
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions - 1:
+                down.downsample = Downsample(block_in)
+                curr_res = curr_res // 2
+            self.down.append(down)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, 2 * z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions - 1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+class Decoder(nn.Module):
+    def __init__(self, ch: int, out_ch: int, ch_mult: list[int], num_res_blocks: int, in_channels: int, resolution: int, z_channels: int):
+        super().__init__()
+        self.ch = ch
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.ffactor = 2 ** (self.num_resolutions - 1)
+        block_in = ch * ch_mult[self.num_resolutions - 1]
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in)
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch * ch_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out))
+                block_in = block_out
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in)
+                curr_res = curr_res * 2
+            self.up.insert(0, up)
+        self.norm_out = nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: Tensor) -> Tensor:
+        upscale_dtype = next(self.up.parameters()).dtype
+        h = self.conv_in(z)
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = h.to(upscale_dtype)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+        h = self.norm_out(h)
+        h = swish(h)
+        h = self.conv_out(h)
+        return h
+
+class DiagonalGaussian(nn.Module):
+    def __init__(self, sample: bool = True, chunk_dim: int = 1):
+        super().__init__()
+        self.sample = sample
+        self.chunk_dim = chunk_dim
+
+    def forward(self, z: Tensor) -> Tensor:
+        mean, logvar = torch.chunk(z, 2, dim=self.chunk_dim)
+        if self.sample:
+            std = torch.exp(0.5 * logvar)
+            return mean + std * torch.randn_like(mean)
+        else:
+            return mean
+
+class AutoEncoder(nn.Module):
+    def __init__(self, params: AutoEncoderParams, sample_z: bool = False):
+        super().__init__()
+        self.params = params
+        self.encoder = Encoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.decoder = Decoder(
+            resolution=params.resolution,
+            in_channels=params.in_channels,
+            ch=params.ch,
+            out_ch=params.out_ch,
+            ch_mult=params.ch_mult,
+            num_res_blocks=params.num_res_blocks,
+            z_channels=params.z_channels,
+        )
+        self.reg = DiagonalGaussian(sample=sample_z)
+        self.scale_factor = params.scale_factor
+        self.shift_factor = params.shift_factor
+
+    def encode(self, x: Tensor) -> Tensor:
+        z = self.reg(self.encoder(x))
+        z = self.scale_factor * (z - self.shift_factor)
+        return z
+
+    def decode(self, z: Tensor) -> Tensor:
+        z = z / self.scale_factor + self.shift_factor
+        return self.decoder(z)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.decode(self.encode(x))

flux/modules/conditioner.py

@@ -0,0 +1,34 @@
+# Flux Project: https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/conditioner.py
+from torch import Tensor, nn
+from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
+
+class HFEmbedder(nn.Module):
+    def __init__(self, version: str, max_length: int, **hf_kwargs):
+        super().__init__()
+        self.is_clip = version.startswith("openai")
+        self.max_length = max_length
+        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"
+        if self.is_clip:
+            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
+        else:
+            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(version, max_length=max_length)
+            self.hf_module: T5EncoderModel = T5EncoderModel.from_pretrained(version, **hf_kwargs)
+        self.hf_module = self.hf_module.eval().requires_grad_(False)
+
+    def forward(self, text: list[str]) -> Tensor:
+        batch_encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            return_length=False,
+            return_overflowing_tokens=False,
+            padding="max_length",
+            return_tensors="pt",
+        )
+        outputs = self.hf_module(
+            input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
+            attention_mask=None,
+            output_hidden_states=False,
+        )
+        return outputs[self.output_key].bfloat16()

flux/modules/image_embedders.py

@@ -0,0 +1,69 @@
+# Flux Project: https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/image_embedders.py
+import cv2
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from PIL import Image
+from safetensors.torch import load_file as load_sft
+from torch import nn
+from transformers import AutoModelForDepthEstimation, AutoProcessor, SiglipImageProcessor, SiglipVisionModel
+from flux.util import print_load_warning
+
+class DepthImageEncoder:
+    depth_model_name = "LiheYoung/depth-anything-large-hf"
+
+    def __init__(self, device):
+        self.device = device
+        self.depth_model = AutoModelForDepthEstimation.from_pretrained(self.depth_model_name).to(device)
+        self.processor = AutoProcessor.from_pretrained(self.depth_model_name)
+
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        hw = img.shape[-2:]
+        img = torch.clamp(img, -1.0, 1.0)
+        img_byte = ((img + 1.0) * 127.5).byte()
+        img = self.processor(img_byte, return_tensors="pt")["pixel_values"]
+        depth = self.depth_model(img.to(self.device)).predicted_depth
+        depth = repeat(depth, "b h w -> b 3 h w")
+        depth = torch.nn.functional.interpolate(depth, hw, mode="bicubic", antialias=True)
+        depth = depth / 127.5 - 1.0
+        return depth
+
+class CannyImageEncoder:
+    def __init__(self, device, min_t: int = 50, max_t: int = 200):
+        self.device = device
+        self.min_t = min_t
+        self.max_t = max_t
+
+    def __call__(self, img: torch.Tensor) -> torch.Tensor:
+        assert img.shape[0] == 1, "Only batch size 1 is supported"
+        img = rearrange(img[0], "c h w -> h w c")
+        img = torch.clamp(img, -1.0, 1.0)
+        img_np = ((img + 1.0) * 127.5).numpy().astype(np.uint8)
+        canny = cv2.Canny(img_np, self.min_t, self.max_t)
+        canny = torch.from_numpy(canny).float() / 127.5 - 1.0
+        canny = rearrange(canny, "h w -> 1 1 h w")
+        canny = repeat(canny, "b 1 ... -> b 3 ...")
+        return canny.to(self.device)
+
+class ReduxImageEncoder(nn.Module):
+    siglip_model_name = "google/siglip-so400m-patch14-384"
+
+    def __init__(self, device, redux_path: str, redux_dim: int = 1152, txt_in_features: int = 4096, dtype=torch.bfloat16):
+        super().__init__()
+        self.redux_dim = redux_dim
+        self.device = device if isinstance(device, torch.device) else torch.device(device)
+        self.dtype = dtype
+        with self.device:
+            self.redux_up = nn.Linear(redux_dim, txt_in_features * 3, dtype=dtype)
+            self.redux_down = nn.Linear(txt_in_features * 3, txt_in_features, dtype=dtype)
+            sd = load_sft(redux_path, device=str(device))
+            missing, unexpected = self.load_state_dict(sd, strict=False, assign=True)
+            print_load_warning(missing, unexpected)
+            self.siglip = SiglipVisionModel.from_pretrained(self.siglip_model_name).to(dtype=dtype)
+        self.normalize = SiglipImageProcessor.from_pretrained(self.siglip_model_name)
+
+    def __call__(self, x: Image.Image) -> torch.Tensor:
+        imgs = self.normalize.preprocess(images=[x], do_resize=True, return_tensors="pt", do_convert_rgb=True)
+        _encoded_x = self.siglip(**imgs.to(device=self.device, dtype=self.dtype)).last_hidden_state
+        projected_x = self.redux_down(nn.functional.silu(self.redux_up(_encoded_x)))
+        return projected_x

flux/modules/layers.py

@@ -0,0 +1,99 @@
+"""
+    Modified: https://github.com/datnguyen-tien204
+    Flux Project: https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
+"""
+import math
+import torch
+from torch import nn, Tensor
+from einops import rearrange
+from flux.utils.math import attention, rope
+
+
+class EmbedND(nn.Module):
+    def __init__(self, dim: int, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: Tensor) -> Tensor:
+        n_axes = ids.shape[-1]
+        emb = torch.cat(
+            [rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)],
+            dim=-3,
+        )
+        return emb.unsqueeze(1)
+
+
+def timestep_embedding(t: Tensor, dim, max_period=10000, time_factor: float = 1000.0):
+    t = time_factor * t
+    half = dim // 2
+    freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(t.device)
+    args = t[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    if torch.is_floating_point(t):
+        embedding = embedding.to(t)
+    return embedding
+
+
+class MLPEmbedder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int):
+        super().__init__()
+        self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
+        self.silu = nn.SiLU()
+        self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.out_layer(self.silu(self.in_layer(x)))
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: Tensor):
+        x_dtype = x.dtype
+        x = x.float()
+        rrms = torch.rsqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + 1e-6)
+
+        # --- DEBUG PRINT ---
+        if x.shape[-1] != self.scale.shape[0]:
+            print(f"!!! SHAPE MISMATCH ERROR !!!")
+            print(f"Input x shape: {x.shape}")
+            print(f"Scale shape: {self.scale.shape}")
+        # -------------------
+
+        return (x * rrms).to(dtype=x_dtype) * self.scale
+
+
+class QKNorm(torch.nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.query_norm = RMSNorm(dim)
+        self.key_norm = RMSNorm(dim)
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q.to(v), k.to(v)
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = False):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.norm = QKNorm(head_dim)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x: Tensor, pe: Tensor) -> Tensor:
+        qkv = self.qkv(x)
+        q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
+        q, k = self.norm(q, k, v)
+        x = attention(q, k, v, pe=pe)
+        x = self.proj(x)
+        return x

flux/modules/losses.py

@@ -0,0 +1,132 @@
+"""
+Perceptual and adversarial loss modules for enhanced training.
+Includes LPIPS, VGG perceptual loss, and PatchGAN discriminator.
+Author: https://github.com/datnguyen-tien204
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import vgg16, VGG16_Weights
+from torchvision import transforms, models
+from typing import List, Optional
+import lpips
+
+
+class CharbonnierLoss(nn.Module):
+    def __init__(self, eps=1e-4):
+        super().__init__()
+        self.eps = eps
+    def forward(self, x, y):
+        diff = x - y
+        return torch.mean(torch.sqrt(diff * diff + self.eps * self.eps))
+    
+
+class DayPreservationLoss(nn.Module):
+    def __init__(self): super().__init__()
+    def forward(self, pred_img, day_img, hour_val):
+        if isinstance(hour_val, torch.Tensor):
+            weights = []
+            for h in hour_val:
+                h = h.item() % 24.0
+                if 11.0 <= h <= 13.0: w = 1.0
+                elif 10.0 <= h < 11.0: w = (h - 10.0)
+                elif 13.0 < h <= 14.0: w = (14.0 - h)
+                else: w = 0.0
+                weights.append(w)
+            weight_tensor = torch.tensor(weights, device=pred_img.device, dtype=pred_img.dtype).view(-1, 1, 1, 1)
+            if weight_tensor.sum() > 0:
+                return (torch.abs(pred_img - day_img) * weight_tensor).mean()
+        return torch.tensor(0.0, device=pred_img.device)
+
+class PerceptualLoss(nn.Module):
+    def __init__(self, device):
+        super().__init__()
+        vgg = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features.to(device).eval()
+        for param in vgg.parameters(): param.requires_grad = False
+        self.slice1 = nn.Sequential(*list(vgg[:4])).to(device)
+        self.slice2 = nn.Sequential(*list(vgg[4:9])).to(device)
+        self.normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+
+    def forward(self, x, y):
+        x_norm = self.normalize((x + 1.0) / 2.0)
+        y_norm = self.normalize((y + 1.0) / 2.0)
+        x_f1 = self.slice1(x_norm)
+        y_f1 = self.slice1(y_norm)
+        x_f2 = self.slice2(x_f1)
+        y_f2 = self.slice2(y_f1)
+        return (F.l1_loss(x_f1, y_f1) * 0.5 + F.l1_loss(x_f2, y_f2) * 0.5) * 0.01
+
+class ColorMatchingLoss(nn.Module):
+    def __init__(self, device, use_perceptual=True):
+        super().__init__()
+        self.use_perceptual = use_perceptual
+        if use_perceptual: self.perceptual_loss = PerceptualLoss(device)
+    def forward(self, pred, target):
+        pred_01, target_01 = (pred + 1) / 2, (target + 1) / 2
+        mean_loss = F.mse_loss(pred_01.mean([2,3]), target_01.mean([2,3]))
+        std_loss = F.mse_loss(pred_01.std([2,3]), target_01.std([2,3]))
+
+        eps = 1e-3
+        p_c, t_c = pred_01.clamp(eps, 1), target_01.clamp(eps, 1)
+        def ratios(x): return (torch.log(x[:,0:1])-torch.log(x[:,1:2])).clamp(-3,3), (torch.log(x[:,0:1])-torch.log(x[:,2:3])).clamp(-3,3), (torch.log(x[:,1:2])-torch.log(x[:,2:3])).clamp(-3,3)
+        prg, prb, pgb = ratios(p_c)
+        trg, trb, tgb = ratios(t_c)
+        ratio_loss = (F.mse_loss(prg, trg) + F.mse_loss(prb, trb) + F.mse_loss(pgb, tgb)) / 3
+
+        loss = mean_loss * 1.0 + std_loss * 0.5 + ratio_loss * 0.3
+        if self.use_perceptual: loss += self.perceptual_loss(pred, target)
+        return loss
+
+class ContinuityLoss(nn.Module):
+    def __init__(self, max_rate=0.15):
+        super().__init__()
+        self.max_rate = max_rate
+
+    def forward(self, outputs, phis):
+        # outputs: [N, 3, H, W] - N can be batch size OR temporal buffer size
+        if outputs.shape[0] < 2: return torch.tensor(0.0, device=outputs.device)
+
+        # Sort by phi to ensure temporal order
+        idx = torch.argsort(phis)
+        out, phi = outputs[idx], phis[idx]
+
+        # Calculate Brightness
+        b = 0.299*out[:,0]+0.587*out[:,1]+0.114*out[:,2]
+
+        # Calculate diffs
+        diff_phi = (phi[1:] - phi[:-1])
+
+        # Filter out cases where phi is too close (avoid division by zero)
+        mask = diff_phi > 1e-4
+        if not mask.any(): return torch.tensor(0.0, device=outputs.device)
+
+        # Rate of change: |dB| / dPhi
+        dB = torch.abs(b[1:] - b[:-1])
+        dB = dB[mask]                         # [K, H, W]
+        dphi = diff_phi[mask].view(-1,1,1)    # [K,1,1]
+
+        rate = dB / dphi
+        return torch.mean(F.relu(rate - self.max_rate).clamp(0, 0.5)**2)
+
+class ColorContinuityLoss(nn.Module):
+    def __init__(self, max_change=0.10):
+        super().__init__()
+        self.max_change = max_change
+
+    def forward(self, outputs, phis):
+        if outputs.shape[0] < 2: return torch.tensor(0.0, device=outputs.device)
+
+        idx = torch.argsort(phis)
+        out, phi = outputs[idx], phis[idx]
+
+        diff_phi = (phi[1:] - phi[:-1])
+        mask = diff_phi > 1e-4
+        if not mask.any(): return torch.tensor(0.0, device=outputs.device)
+
+        # Rate of change: |dColor| / dPhi
+        rate = torch.abs(out[1:] - out[:-1])[mask] / diff_phi[mask].view(-1,1,1,1)
+
+        loss = 0
+        for c in range(3):
+            loss += torch.mean(F.relu(rate[:,c] - self.max_change).clamp(0, 0.3)**2)
+        return loss / 3

flux/modules/moe.py

@@ -0,0 +1,149 @@
+"""
+Mixture-of-Experts (MoE) module for time-based style control.
+Implements sparse expert selection based on time intervals.
+Author: https://github.com/datnguyen-tien204
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional, List
+import math
+
+class Expert(nn.Module):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
+        super().__init__()
+        self.network = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, output_dim)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.network(x)
+
+class TimeBasedGatingNetwork(nn.Module):
+    def __init__(self, num_experts: int, input_dim: int = 2, temperature: float = 0.2,
+                 custom_centers: Optional[List[float]] = None,
+                 custom_widths: Optional[List[float]] = None):
+        super().__init__()
+        self.num_experts = num_experts
+        self.temperature = temperature
+
+        if custom_centers is not None:
+            centers_rad = [c * (2 * math.pi / 24.0) for c in custom_centers]
+            self.register_buffer('expert_centers', torch.tensor(centers_rad))
+        else:
+            centers = [i * (2 * math.pi / num_experts) for i in range(num_experts)]
+            self.register_buffer('expert_centers', torch.tensor(centers))
+
+        if custom_widths is not None:
+            widths_rad = [(w * (2 * math.pi / 24.0)) / 2.0 for w in custom_widths]
+            init_log_widths = torch.log(torch.tensor(widths_rad) + 1e-5)
+            self.expert_log_widths = nn.Parameter(init_log_widths)
+        else:
+            self.expert_log_widths = nn.Parameter(torch.zeros(num_experts))
+
+        self.min_width = 0.1
+
+    def forward(self, phi: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        phi = phi.view(-1, 1)
+        centers = self.expert_centers.view(1, -1)
+
+        diff = torch.remainder(phi - centers + math.pi, 2 * math.pi) - math.pi
+
+        widths = F.softplus(self.expert_log_widths).view(1, -1) + self.min_width
+
+        logits = -diff ** 2 / (2 * widths ** 2)
+        weights = F.softmax(logits / self.temperature, dim=-1)
+
+        return weights, logits
+
+class MixtureOfExperts(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, num_experts, k=2, temperature=0.2, noise_std=0.1,
+                 custom_centers=None, custom_widths=None):
+        super().__init__()
+        self.num_experts = num_experts
+        self.k = k
+        self.noise_std = noise_std
+
+        self.experts = nn.ModuleList([
+            Expert(input_dim, hidden_dim, output_dim) for _ in range(num_experts)
+        ])
+
+        self.gate = TimeBasedGatingNetwork(
+            num_experts, input_dim=2, temperature=temperature,
+            custom_centers=custom_centers, custom_widths=custom_widths
+        )
+        self.register_buffer('load_balance_loss_coef', torch.tensor(0.01))
+
+    def _add_noise(self, logits):
+        if self.training and self.noise_std > 0:
+            noise = torch.randn_like(logits) * self.noise_std
+            return logits + noise
+        return logits
+
+    def _compute_load_balance_loss(self, weights):
+        mean_weights = weights.mean(dim=0)
+        target = torch.ones_like(mean_weights) / self.num_experts
+        loss = F.mse_loss(mean_weights, target)
+        return loss * self.load_balance_loss_coef
+
+    def forward(self, x, phi, return_aux=False):
+        is_sequence = x.dim() == 3
+        if is_sequence:
+            B, L, C = x.shape
+            x = x.reshape(B * L, C)
+            phi = phi.unsqueeze(1).expand(B, L).reshape(B * L)
+
+        weights, gate_logits = self.gate(phi)
+
+        if self.training:
+            gate_logits = self._add_noise(gate_logits)
+            weights = F.softmax(gate_logits / self.gate.temperature, dim=-1)
+
+        if self.k < self.num_experts:
+            top_k_weights, top_k_indices = torch.topk(weights, self.k, dim=-1)
+            top_k_weights = top_k_weights / top_k_weights.sum(dim=-1, keepdim=True)
+            sparse_weights = torch.zeros_like(weights)
+            sparse_weights.scatter_(1, top_k_indices, top_k_weights)
+            weights = sparse_weights
+
+        expert_outputs = torch.stack([e(x) for e in self.experts], dim=1)
+        output = torch.einsum('be,beo->bo', weights, expert_outputs)
+
+        if is_sequence:
+            output = output.reshape(B, L, -1)
+            weights = weights.reshape(B, L, -1).mean(dim=1)
+
+        aux = None
+        if return_aux:
+            aux = {'load_balance_loss': self._compute_load_balance_loss(weights)}
+
+        return output, aux
+
+class MoEStyleAdapter(nn.Module):
+    def __init__(self, style_dim, hidden_dim, num_experts, custom_centers=None, custom_widths=None):
+        super().__init__()
+        self.moe = MixtureOfExperts(
+            style_dim, hidden_dim, style_dim, num_experts, k=2, temperature=0.2, noise_std=0.1,
+            custom_centers=custom_centers, custom_widths=custom_widths
+        )
+    def forward(self, style, phi, return_aux=False):
+        adapted, aux = self.moe(style, phi, return_aux=return_aux)
+        return adapted + style, aux
+
+class MoEFeatureProcessor(nn.Module):
+    def __init__(self, feature_dim, num_experts, expand_ratio=4, custom_centers=None, custom_widths=None):
+        super().__init__()
+        self.moe = MixtureOfExperts(
+            feature_dim, feature_dim * expand_ratio, feature_dim, num_experts, k=2, temperature=0.2, noise_std=0.1,
+            custom_centers=custom_centers, custom_widths=custom_widths
+        )
+        self.norm = nn.LayerNorm(feature_dim)
+    def forward(self, x, phi, return_aux=False):
+        normed = self.norm(x)
+        processed, aux = self.moe(normed, phi, return_aux=return_aux)
+        return x + processed, aux

flux/modules/normalization.py

@@ -0,0 +1,263 @@
+"""
+Adaptive Normalization Modules for ChronFlux
+Includes AdaIN (Adaptive Instance Normalization) and AdaLN (Adaptive Layer Normalization)
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+NIGHT_SKY_INTENSITY = -0.95
+WARM_BIAS_RGB = [0.2, 0.0, -0.2]
+
+
+class AdaIN(nn.Module):
+    """
+    Adaptive Instance Normalization (AdaIN) for style transfer.
+    Used for generating pseudo-style targets during training with extensive augmentation.
+
+    Reference: Huang & Belongie, "Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization"
+    """
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+
+    def calc_mean_std(self, feat):
+        size = feat.size()
+        N, C = size[:2]
+        feat_view = feat.view(N, C, -1)
+
+        feat_std, feat_mean = torch.std_mean(feat_view, dim=2, unbiased=False)
+
+        feat_std = (feat_std + self.eps).view(N, C, 1, 1)
+        feat_mean = feat_mean.view(N, C, 1, 1)
+
+        return feat_mean, feat_std
+
+    def forward(self, content, style):
+        """
+        Apply AdaIN: transfers style statistics to content
+
+        Args:
+            content: Content features [B, C, H, W]
+            style: Style features [B, C, H, W]
+
+        Returns:
+            Stylized content with style's mean and std
+        """
+        size = content.size()
+        content_mean, content_std = self.calc_mean_std(content)
+        style_mean, style_std = self.calc_mean_std(style)
+
+        normalized = (content - content_mean) / content_std
+        stylized = normalized * style_std + style_mean
+
+        return stylized
+
+
+class AdaINTransfer(nn.Module):
+    """
+    Extended AdaIN for pseudo-style target generation with augmentation support.
+    Used specifically for training pipeline to create robust style targets.
+    """
+    def __init__(self, eps=1e-5):
+        super().__init__()
+        self.adain = AdaIN(eps=eps)
+
+        self.register_buffer(
+            'warm_bias_template',
+            torch.tensor(WARM_BIAS_RGB, dtype=torch.float32).view(1, 3, 1, 1)
+        )
+
+    def adaptive_night_transfer(self, content, style, sky_mask=None):
+        """
+        Transfer night style to content with optional sky masking
+
+        Args:
+            content: Day images [B, C, H, W]
+            style: Night images [B, C, H, W]
+            sky_mask: Binary mask for sky regions [B, 1, H, W]
+
+        Returns:
+            Night-styled content
+        """
+        target = self.adain(content, style)
+
+        if sky_mask is not None:
+            target = target * (1 - sky_mask) + (NIGHT_SKY_INTENSITY * sky_mask)
+
+        return target.clamp(-0.95, 0.95)
+
+    def adaptive_sunset_transfer(self, content, style_night, sky_mask=None):
+        """
+        Create sunset style by channel manipulation of night style
+
+        Args:
+            content: Day images [B, C, H, W]
+            style_night: Night images for base style [B, C, H, W]
+            sky_mask: Binary mask for sky regions [B, 1, H, W]
+
+        Returns:
+            Sunset-styled content with warm tones
+        """
+        content_mean, content_std = self.adain.calc_mean_std(content)
+        style_mean, style_std = self.adain.calc_mean_std(style_night)
+
+        mu_swap = style_mean.clone()
+        mu_swap[:, 0] = style_mean[:, 2]
+        mu_swap[:, 2] = style_mean[:, 0]
+
+        std_swap = style_std.clone()
+        std_swap[:, 0] = style_std[:, 2]
+        std_swap[:, 2] = style_std[:, 0]
+
+        B = content.shape[0]
+        warm_bias = self.warm_bias_template.to(content.dtype).expand(B, 3, 1, 1)
+
+        mu_final = mu_swap * 0.7 + warm_bias * 0.3
+        std_final = std_swap * 1.4
+
+        target = (content - content_mean) / content_std * std_final + mu_final
+
+        if sky_mask is not None:
+            sky_color = mu_final.view(B, 3, 1, 1)
+            target = target * (1 - sky_mask) + (target * 0.4 + sky_color * 0.6) * sky_mask
+
+        return target.clamp(-0.95, 0.95)
+
+
+class AdaLNModulation(nn.Module):
+    """
+    Adaptive Layer Normalization (AdaLN) Modulation.
+    Generates scale, shift, and gate parameters for layer-wise adaptive normalization.
+    Provides better control over style consistency and layer-wise adjustments.
+    """
+    def __init__(self, hidden_size, components=6):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(hidden_size, components * hidden_size, bias=True)
+        self.components = components
+
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+    def forward(self, vec):
+        """
+        Generate modulation parameters
+
+        Args:
+            vec: Conditioning vector [B, hidden_size]
+
+        Returns:
+            Tuple of (shift, scale, gate) parameters for each component
+        """
+        chunk = self.linear(self.silu(vec))
+        return chunk.chunk(self.components, dim=1)
+
+
+class AdaLNBlock(nn.Module):
+    """
+    Base AdaLN normalization block with scale, shift, and gate.
+    Can be used as a building block for various architectures.
+    """
+    def __init__(self, hidden_size, use_gate=True):
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.use_gate = use_gate
+
+    def forward(self, x, scale, shift, gate=None):
+        """
+        Apply AdaLN: normalize, then scale and shift, optionally gate
+
+        Args:
+            x: Input features [B, L, D]
+            scale: Scale parameter [B, D]
+            shift: Shift parameter [B, D]
+            gate: Optional gate parameter [B, D]
+
+        Returns:
+            Modulated features
+        """
+        x_norm = self.norm(x)
+        x_mod = x_norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+        if self.use_gate and gate is not None:
+            x_mod = gate.unsqueeze(1) * x_mod
+
+        return x_mod
+
+
+class AdaLNZero(nn.Module):
+    """
+    AdaLN-Zero: AdaLN with zero-initialized output scaling.
+    Particularly useful for residual connections and stable training.
+    """
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.scale = nn.Parameter(torch.zeros(hidden_size))
+        self.shift = nn.Parameter(torch.zeros(hidden_size))
+
+    def forward(self, x, modulation_scale, modulation_shift):
+        """
+        Apply AdaLN with zero-initialization
+
+        Args:
+            x: Input features [B, L, D]
+            modulation_scale: Additional scale from conditioning [B, D]
+            modulation_shift: Additional shift from conditioning [B, D]
+
+        Returns:
+            Modulated features
+        """
+        x_norm = self.norm(x)
+        x_mod = x_norm * (self.scale + modulation_scale.unsqueeze(1)) + (self.shift + modulation_shift.unsqueeze(1))
+        return x_mod
+
+
+class StyleAdaLN(nn.Module):
+    """
+    Style-aware AdaLN that combines style and conditioning vectors.
+    Particularly useful for style transfer with additional controls.
+    """
+    def __init__(self, hidden_size, style_dim):
+        super().__init__()
+        self.style_proj = nn.Linear(style_dim, hidden_size)
+        self.adaln = AdaLNModulation(hidden_size, components=2)
+        self.norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, style_vec, cond_vec=None):
+        """
+        Apply style-aware normalization
+
+        Args:
+            x: Input features [B, L, D]
+            style_vec: Style vector [B, style_dim]
+            cond_vec: Optional additional conditioning [B, D]
+
+        Returns:
+            Style-modulated features
+        """
+        style_emb = self.style_proj(style_vec)
+
+        if cond_vec is not None:
+            combined = style_emb + cond_vec
+        else:
+            combined = style_emb
+
+        scale, shift = self.adaln(combined)
+
+        x_norm = self.norm(x)
+        x_mod = x_norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+        return x_mod
+
+
+__all__ = [
+    'AdaIN',
+    'AdaINTransfer',
+    'AdaLNModulation',
+    'AdaLNBlock',
+    'AdaLNZero',
+    'StyleAdaLN',
+]

flux/modules/sampling.py

@@ -0,0 +1,379 @@
+"""
+State-of-the-art sampling techniques (2024-2025).
+Includes Flow Matching with Optimal Transport and Consistency Models.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Callable
+import math
+
+
+class FlowMatching:
+    """
+    Flow Matching with Optimal Transport for improved trajectory learning.
+    Based on recent advances in continuous normalizing flows.
+    """
+    def __init__(self, sigma_min: float = 0.001):
+        """
+        Args:
+            sigma_min: Minimum noise level for numerical stability
+        """
+        self.sigma_min = sigma_min
+    
+    def compute_optimal_transport_path(self, x0: torch.Tensor, x1: torch.Tensor, 
+                                      t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute optimal transport path from noise (x0) to data (x1).
+        
+        Args:
+            x0: Source distribution (noise) [B, ...]
+            x1: Target distribution (data) [B, ...]
+            t: Time steps [B] in range [0, 1]
+        Returns:
+            x_t: Interpolated samples
+            v_t: Vector field (velocity)
+        """
+        t = t.view(-1, *([1] * (x0.dim() - 1)))
+        
+        # Optimal transport interpolation (linear for Gaussian)
+        x_t = (1 - t) * x0 + t * x1
+        
+        # Velocity is the derivative: v = dx/dt = x1 - x0
+        v_t = x1 - x0
+        
+        return x_t, v_t
+    
+    def compute_conditional_flow(self, x0: torch.Tensor, x1: torch.Tensor,
+                                t: torch.Tensor, sigma: Optional[float] = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute conditional flow with added noise for robustness.
+        
+        Args:
+            x0: Noise [B, ...]
+            x1: Data [B, ...]
+            t: Time [B]
+            sigma: Optional noise level (defaults to sigma_min * (1-t))
+        Returns:
+            x_t: Noisy interpolation
+            v_t: Target velocity
+        """
+        if sigma is None:
+            sigma = self.sigma_min * (1 - t.view(-1, *([1] * (x0.dim() - 1))))
+        else:
+            sigma = sigma * torch.ones_like(t).view(-1, *([1] * (x0.dim() - 1)))
+        
+        # Conditional path
+        x_t, v_t = self.compute_optimal_transport_path(x0, x1, t)
+        
+        # Add Gaussian noise for conditional flow matching
+        eps = torch.randn_like(x_t)
+        x_t = x_t + sigma * eps
+        
+        return x_t, v_t
+    
+    def sample(self, model: Callable, x_init: torch.Tensor, 
+              num_steps: int = 50, **model_kwargs) -> torch.Tensor:
+        """
+        Sample using Euler method for ODE solving.
+        
+        Args:
+            model: Velocity prediction model
+            x_init: Initial noise [B, ...]
+            num_steps: Number of integration steps
+            **model_kwargs: Additional arguments for model
+        Returns:
+            x_final: Generated samples
+        """
+        x = x_init
+        dt = 1.0 / num_steps
+        
+        for step in range(num_steps):
+            t_current = torch.full((x.shape[0],), step / num_steps, 
+                                  device=x.device, dtype=x.dtype)
+            
+            # Predict velocity
+            with torch.no_grad():
+                v = model(x, timesteps=t_current, **model_kwargs)
+            
+            # Euler step
+            x = x + v * dt
+        
+        return x
+    
+    def sample_heun(self, model: Callable, x_init: torch.Tensor,
+                   num_steps: int = 50, **model_kwargs) -> torch.Tensor:
+        """
+        Higher-order sampling using Heun's method (2nd order).
+        More accurate but requires 2x model evaluations.
+        
+        Args:
+            model: Velocity prediction model
+            x_init: Initial noise
+            num_steps: Number of steps
+            **model_kwargs: Model arguments
+        Returns:
+            x_final: Generated samples
+        """
+        x = x_init
+        dt = 1.0 / num_steps
+        
+        for step in range(num_steps):
+            t_current = torch.full((x.shape[0],), step / num_steps,
+                                  device=x.device, dtype=x.dtype)
+            
+            # First stage
+            with torch.no_grad():
+                v1 = model(x, timesteps=t_current, **model_kwargs)
+            
+            x_temp = x + v1 * dt
+            
+            # Second stage
+            t_next = torch.full((x.shape[0],), (step + 1) / num_steps,
+                               device=x.device, dtype=x.dtype)
+            
+            with torch.no_grad():
+                v2 = model(x_temp, timesteps=t_next, **model_kwargs)
+            
+            # Heun's method: average of two velocities
+            x = x + (v1 + v2) * 0.5 * dt
+        
+        return x
+
+
+class ConsistencyModel:
+    """
+    Consistency Models for fast, single-step generation.
+    Based on "Consistency Models" (Song et al., 2023).
+    """
+    def __init__(self, sigma_min: float = 0.002, sigma_max: float = 80.0,
+                 rho: float = 7.0):
+        """
+        Args:
+            sigma_min: Minimum noise level
+            sigma_max: Maximum noise level
+            rho: Schedule parameter
+        """
+        self.sigma_min = sigma_min
+        self.sigma_max = sigma_max
+        self.rho = rho
+    
+    def get_schedule(self, num_steps: int, device: torch.device) -> torch.Tensor:
+        """
+        Get noise schedule for consistency training.
+        
+        Args:
+            num_steps: Number of discretization steps
+            device: Device for tensors
+        Returns:
+            sigmas: Noise levels [num_steps + 1]
+        """
+        rho_inv = 1.0 / self.rho
+        steps = torch.arange(num_steps + 1, device=device)
+        sigmas = (self.sigma_max ** rho_inv + 
+                 steps / num_steps * (self.sigma_min ** rho_inv - self.sigma_max ** rho_inv)) ** self.rho
+        return sigmas
+    
+    def consistency_function(self, model: Callable, x: torch.Tensor,
+                            sigma: torch.Tensor, **model_kwargs) -> torch.Tensor:
+        """
+        Apply consistency function: maps (x, sigma) to denoised output.
+        
+        Args:
+            model: Denoising model
+            x: Noisy input
+            sigma: Noise level
+            **model_kwargs: Model arguments
+        Returns:
+            denoised: Consistency function output
+        """
+        # Skip connection for sigma_min
+        c_skip = self.sigma_min ** 2 / (sigma ** 2 + self.sigma_min ** 2)
+        c_out = sigma * self.sigma_min / (sigma ** 2 + self.sigma_min ** 2).sqrt()
+        
+        # Predict denoised output
+        with torch.no_grad():
+            # Convert sigma to timestep for model
+            t = sigma / self.sigma_max
+            model_out = model(x, timesteps=t, **model_kwargs)
+        
+        # Apply consistency function
+        denoised = c_skip * x + c_out * model_out
+        
+        return denoised
+    
+    def sample_onestep(self, model: Callable, shape: tuple,
+                      device: torch.device, dtype: torch.dtype,
+                      **model_kwargs) -> torch.Tensor:
+        """
+        Single-step generation from noise to data.
+        
+        Args:
+            model: Consistency model
+            shape: Output shape
+            device: Device
+            dtype: Data type
+            **model_kwargs: Model arguments
+        Returns:
+            samples: Generated samples
+        """
+        # Sample from prior (high noise)
+        x = torch.randn(*shape, device=device, dtype=dtype) * self.sigma_max
+        sigma = torch.full((shape[0],), self.sigma_max, device=device, dtype=dtype)
+        
+        # One-step denoising
+        samples = self.consistency_function(model, x, sigma, **model_kwargs)
+        
+        return samples
+    
+    def sample_multistep(self, model: Callable, shape: tuple,
+                        num_steps: int, device: torch.device,
+                        dtype: torch.dtype, **model_kwargs) -> torch.Tensor:
+        """
+        Multi-step generation for better quality.
+        
+        Args:
+            model: Consistency model
+            shape: Output shape
+            num_steps: Number of steps
+            device: Device
+            dtype: Data type
+            **model_kwargs: Model arguments
+        Returns:
+            samples: Generated samples
+        """
+        # Get noise schedule
+        sigmas = self.get_schedule(num_steps, device)
+        
+        # Start from maximum noise
+        x = torch.randn(*shape, device=device, dtype=dtype) * self.sigma_max
+        
+        # Iterative refinement
+        for i in range(num_steps, 0, -1):
+            sigma_current = sigmas[i]
+            sigma_next = sigmas[i - 1]
+            
+            # Denoise to current level
+            x = self.consistency_function(
+                model, x, 
+                torch.full((shape[0],), sigma_current, device=device, dtype=dtype),
+                **model_kwargs
+            )
+            
+            # Add noise for next iteration (except last)
+            if i > 1:
+                noise = torch.randn_like(x)
+                x = x + noise * (sigma_next ** 2 - self.sigma_min ** 2).sqrt()
+        
+        return x
+
+
+class ClassifierFreeGuidance:
+    """
+    Enhanced Classifier-Free Guidance for better conditioning.
+    """
+    def __init__(self, guidance_scale: float = 7.5, min_scale: float = 1.0):
+        """
+        Args:
+            guidance_scale: Guidance strength
+            min_scale: Minimum guidance scale
+        """
+        self.guidance_scale = guidance_scale
+        self.min_scale = min_scale
+    
+    def apply_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
+                      guidance_scale: Optional[float] = None) -> torch.Tensor:
+        """
+        Apply classifier-free guidance to model outputs.
+        
+        Args:
+            cond_out: Conditional model output
+            uncond_out: Unconditional model output
+            guidance_scale: Optional override for guidance scale
+        Returns:
+            guided_out: Guided output
+        """
+        scale = guidance_scale if guidance_scale is not None else self.guidance_scale
+        scale = max(scale, self.min_scale)
+        
+        # CFG formula: out = uncond + scale * (cond - uncond)
+        guided_out = uncond_out + scale * (cond_out - uncond_out)
+        
+        return guided_out
+    
+    def apply_dynamic_guidance(self, cond_out: torch.Tensor, uncond_out: torch.Tensor,
+                              t: torch.Tensor, min_t: float = 0.0, 
+                              max_t: float = 1.0) -> torch.Tensor:
+        """
+        Apply time-dependent guidance that decreases over time.
+        
+        Args:
+            cond_out: Conditional output
+            uncond_out: Unconditional output
+            t: Time steps [B]
+            min_t: Time to start decreasing guidance
+            max_t: Time to reach minimum guidance
+        Returns:
+            guided_out: Dynamically guided output
+        """
+        # Linearly decrease guidance scale from max to min over time
+        t_normalized = (t - min_t) / (max_t - min_t)
+        t_normalized = t_normalized.clamp(0, 1)
+        
+        # Scale goes from guidance_scale to min_scale
+        scale = self.guidance_scale + (self.min_scale - self.guidance_scale) * t_normalized
+        scale = scale.view(-1, *([1] * (cond_out.dim() - 1)))
+        
+        guided_out = uncond_out + scale * (cond_out - uncond_out)
+        
+        return guided_out
+
+
+class AdaptiveSampler:
+    """
+    Adaptive sampling that combines multiple techniques for optimal quality/speed.
+    """
+    def __init__(self, use_flow_matching: bool = True,
+                 use_consistency: bool = False,
+                 use_cfg: bool = True):
+        self.use_flow_matching = use_flow_matching
+        self.use_consistency = use_consistency
+        self.use_cfg = use_cfg
+        
+        if use_flow_matching:
+            self.flow_matching = FlowMatching()
+        if use_consistency:
+            self.consistency_model = ConsistencyModel()
+        if use_cfg:
+            self.cfg = ClassifierFreeGuidance()
+    
+    def sample(self, model: Callable, x_init: torch.Tensor,
+              num_steps: int = 50, method: str = "heun",
+              cfg_scale: Optional[float] = None,
+              **model_kwargs) -> torch.Tensor:
+        """
+        Adaptive sampling with multiple techniques.
+        
+        Args:
+            model: Model for sampling
+            x_init: Initial noise
+            num_steps: Number of steps
+            method: Sampling method ("euler", "heun", "consistency")
+            cfg_scale: Optional CFG scale
+            **model_kwargs: Model arguments
+        Returns:
+            samples: Generated samples
+        """
+        if method == "consistency" and self.use_consistency:
+            return self.consistency_model.sample_multistep(
+                model, x_init.shape, num_steps,
+                x_init.device, x_init.dtype, **model_kwargs
+            )
+        elif method == "heun" and self.use_flow_matching:
+            return self.flow_matching.sample_heun(
+                model, x_init, num_steps, **model_kwargs
+            )
+        else:  # euler
+            return self.flow_matching.sample(
+                model, x_init, num_steps, **model_kwargs
+            )

flux/utils/__init__.py

@@ -0,0 +1,4 @@
+import os
+import logging
+import numpy as np
+from flux.utils import logger_utils

flux/utils/chronflux_datasets.py

@@ -0,0 +1,26 @@
+"""
+    Author: https://github.com/datnguyen-tien204
+"""
+
+import os
+import numpy as np
+from PIL import Image
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+class DayNightDataset(Dataset):
+    def __init__(self, day_dir, night_dir, size=384):
+        self.day_files = [os.path.join(day_dir, f) for f in os.listdir(day_dir) if f.endswith(('.jpg', '.png'))]
+        self.night_files = [os.path.join(night_dir, f) for f in os.listdir(night_dir) if f.endswith(('.jpg', '.png'))]
+        self.transform = transforms.Compose(
+            [transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
+
+    def __len__(self): return len(self.day_files)
+
+    def __getitem__(self, idx):
+        day_img = Image.open(self.day_files[idx]).convert('RGB')
+        night_img = Image.open(np.random.choice(self.night_files)).convert('RGB')
+        return self.transform(day_img), self.transform(night_img)

flux/utils/logger_utils.py

@@ -0,0 +1,19 @@
+class ConsoleLogger:
+    HEADER = '\033[95m'
+    BLUE = '\033[94m'
+    CYAN = '\033[96m'
+    GREEN = '\033[92m'
+    WARNING = '\033[93m'
+    FAIL = '\033[91m'
+    ENDC = '\033[0m'
+    BOLD = '\033[1m'
+    UNDERLINE = '\033[4m'
+
+    @staticmethod
+    def info(msg): print(f"{ConsoleLogger.GREEN}{ConsoleLogger.BOLD}[INFO]{ConsoleLogger.ENDC} {msg}")
+    @staticmethod
+    def step(msg): print(f"{ConsoleLogger.CYAN}[STEP]{ConsoleLogger.ENDC} {msg}")
+    @staticmethod
+    def warn(msg): print(f"{ConsoleLogger.WARNING}{ConsoleLogger.BOLD}[WARN]{ConsoleLogger.ENDC} {msg}")
+    @staticmethod
+    def section(msg): print(f"\n{ConsoleLogger.HEADER}{ConsoleLogger.BOLD}{'='*60}\n{msg}\n{'='*60}{ConsoleLogger.ENDC}")

flux/utils/math.py

@@ -0,0 +1,48 @@
+"""
+Author: https://github.com/datnguyen-tien204
+"""
+
+import torch
+from einops import rearrange
+from torch import Tensor
+from flash_attn import flash_attn_func
+
+
+def attention(q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+    # Ensure tensors are in bf16 or fp16 for FlashAttention =====
+    original_dtype = q.dtype
+    if q.dtype not in [torch.float16, torch.bfloat16]:
+        q = q.to(torch.bfloat16)
+        k = k.to(torch.bfloat16)
+        v = v.to(torch.bfloat16)
+    # ======================================================================
+
+    x = flash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)
+    x = rearrange(x, "B L H D -> B L (H D)")
+
+    # Convert back to original dtype if needed =====
+    if original_dtype not in [torch.float16, torch.bfloat16]:
+        x = x.to(original_dtype)
+    # ==========================================================
+
+    return x
+
+
+def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
+    assert dim % 2 == 0
+    scale = torch.arange(0, dim, 2, dtype=torch.float32, device=pos.device) / dim
+    omega = 1.0 / (theta ** scale)
+    out = torch.einsum("...n,d->...nd", pos, omega)
+    out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)
+    out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
+    return out.float()
+
+
+def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor) -> tuple[Tensor, Tensor]:
+    xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
+
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)

flux/utils/target_generation.py

@@ -0,0 +1,63 @@
+"""
+    Author: https://github.com/datnguyen-tien204
+"""
+import numpy as np
+import torch
+import math
+
+PI = math.pi
+
+def symmetric_brightness_curve(hour):
+    return 0.035 + (1.0 - 0.035) * (0.5 - 0.5 * np.cos(2 * PI * (hour / 24.0)))
+
+
+def rgb_curve_original(hour):
+    t = (hour % 24.0) / 24.0
+    sunset = np.exp(-((t - 0.75) / 0.08) ** 2)
+    sunrise = np.exp(-((t - 0.25) / 0.08) ** 2)
+    night = np.exp(-((t - 0.0) / 0.10) ** 2) + np.exp(-((t - 1.0) / 0.10) ** 2)
+    noon = np.exp(-((t - 0.5) / 0.12) ** 2)
+    return [1.0 + 0.4 * sunset + 0.22 * sunrise - 0.35 * night,
+            1.0 - 0.1 * sunset - 0.06 * night,
+            1.0 - 0.5 * sunset + 0.3 * night + 0.15 * noon]
+
+
+def create_target_final_v2(day_img, style_img, depth_map, hour_val, adain_module, style_stats):
+    device, dtype, B = day_img.device, day_img.dtype, day_img.shape[0]
+    img_original = (day_img + 1.0) / 2.0
+    target_list = []
+
+    for i in range(B):
+        h = hour_val[i].item()
+        bright = symmetric_brightness_curve(h)
+        rgb = rgb_curve_original(h)
+
+        tint = torch.tensor(rgb, device=device, dtype=dtype).view(1, 3, 1, 1)
+        tint = tint * (0.6 + 0.4 * style_stats["mean"][i].view(1, 3, 1, 1))
+
+        img_i = img_original[i:i + 1]
+        target_base = img_i * tint
+
+        if bright < 0.8:
+            adain_out = (adain_module(day_img[i:i + 1], style_img[i:i + 1]) + 1.0) / 2.0
+            night_int = (1.0 - bright) ** 0.8 * 0.85
+
+            d_i = depth_map[i:i + 1]
+            d_norm = (d_i - d_i.min()) / (d_i.max() - d_i.min() + 1e-6)
+            sky_mask = torch.sigmoid((0.15 - d_norm) * 10.0)
+
+            blend = night_int * (0.7 + 0.3 * sky_mask)
+            target_mixed = target_base * (1 - blend) + adain_out * blend
+        else:
+            target_mixed = target_base
+            sky_mask = torch.zeros_like(target_mixed[:, :1])
+
+        target_calc = target_mixed * bright
+        if bright < 0.20:
+            target_calc += img_i.mean(1, keepdim=True) * tint * 0.03
+            if bright < 0.8: target_calc *= (1.0 - sky_mask * 0.5)
+
+        noon_w = torch.tensor(np.exp(-(abs(h - 12.0) ** 2) / 8.0), device=device, dtype=dtype)
+        target_list.append(target_calc * (1 - noon_w) + img_i * noon_w)
+
+    return (torch.cat(target_list, dim=0) * 2.0 - 1.0).clamp(-1, 1)

flux/utils/util.py

@@ -0,0 +1,76 @@
+import os
+from dataclasses import dataclass
+from pathlib import Path
+import torch
+from einops import rearrange
+from PIL import Image
+from safetensors.torch import load_file as load_sft
+from huggingface_hub import hf_hub_download
+from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams
+
+CHECKPOINTS_DIR = Path("checkpoints")
+CHECKPOINTS_DIR.mkdir(exist_ok=True)
+
+
+@dataclass
+class VAEConfig:
+    repo_id: str
+    repo_ae: str
+    ae_params: AutoEncoderParams
+
+
+flux_vae_config = VAEConfig(
+    repo_id="black-forest-labs/FLUX.1-dev",
+    repo_ae="ae.safetensors",
+    ae_params=AutoEncoderParams(
+        resolution=64,
+        in_channels=3,
+        ch=128,
+        out_ch=3,
+        ch_mult=[1, 2, 4, 4],
+        num_res_blocks=2,
+        z_channels=16,
+        scale_factor=0.3611,
+        shift_factor=0.1159,
+    ),
+)
+
+
+def get_checkpoint_path(filename: str, repo_id: str) -> Path:
+    local_path = CHECKPOINTS_DIR / filename
+    if not local_path.exists():
+        print(f"Downloading {filename} from {repo_id}...")
+        try:
+            downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=CHECKPOINTS_DIR)
+            return Path(downloaded_path)
+        except Exception as e:
+            print(f"Failed to download: {e}")
+    return local_path
+
+
+def load_ae(device: str | torch.device = "cuda") -> AutoEncoder:
+    print("Init AE")
+    ckpt_path = get_checkpoint_path(flux_vae_config.repo_ae, flux_vae_config.repo_id)
+
+    if ckpt_path.exists():
+        print(f"Loading AE checkpoint: {ckpt_path}")
+        with torch.device("meta"):
+            ae = AutoEncoder(flux_vae_config.ae_params)
+        sd = load_sft(str(ckpt_path), device=str(device))
+        ae.load_state_dict(sd, strict=False, assign=True)
+    else:
+        raise RuntimeError("AE Checkpoint is missing! Model cannot train with random VAE.")
+
+    return ae
+
+
+def save_image(x: torch.Tensor, path: str):
+    x = x.clamp(-1, 1)
+    x = rearrange(x[0], "c h w -> h w c")
+    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
+    img.save(path, quality=95, subsampling=0)
+
+
+def print_load_warning(missing: list[str], unexpected: list[str]):
+    if len(missing) > 0 and len(unexpected) > 0:
+        print(f"Parameters missing: {len(missing)}. Unexpected: {len(unexpected)}")

infer_1step_style.py

@@ -0,0 +1,210 @@
+import os
+import sys
+import math
+import torch
+import torch.nn.functional as F
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+from torchvision import transforms
+
+sys.path.insert(0, "/workspace/ChronFLUX/ChronFlux-v1")
+
+from flux.utils.util import load_ae
+from flux.chronflux import ChronFlux
+from flux.modules.advanced_encoders import EncoderFactory, DepthEstimator
+from config import get_config
+
+PI = math.pi
+
+# ===== MODEL =====
+CHECKPOINT = "/workspace/ChronFLUX/ChronFlux-v1/distill_checkpoints/rift_1step_epoch_5.pt"
+
+# ===== DATA (FIXED MAPPING) =====
+DATA_PAIRS = [
+    {
+        "input":  "/workspace/ChronFLUX/ChronFlux-v1/EvaluationDS/ChronFLUX-Bench/SDSD/outdoor_png/clear",
+        "output": "/workspace/ChronFLUX/ChronFlux-v1/EvaluationDS/DS-BenchFinal/SDSD/outdoor_png/dark_gen",
+    },
+    {
+        "input":  "/workspace/ChronFLUX/ChronFlux-v1/EvaluationDS/ChronFLUX-Bench/SDSD/indoor_png/clear",
+        "output": "/workspace/ChronFLUX/ChronFlux-v1/EvaluationDS/DS-BenchFinal/SDSD/indoor_png/dark_gen",
+    },
+    {
+        "input":  "/workspace/ChronFLUX/ChronFlux-v1/EvaluationDS/ChronFLUX-Bench/SMID/clear",
+        "output": "/workspace/ChronFLUX/ChronFlux-v1/EvaluationDS/DS-BenchFinal/SMID/dark_gen",
+    },
+]
+
+# ===== STYLE =====
+FIXED_STYLE_PATH = "/workspace/ChronFLUX/rafael_ds_ver2/train/night/frame_000002.jpg"
+
+# ===== PARAMS =====
+HOURS = list(range(1, 25))
+
+IMAGE_SIZE = 384
+DEPTH_SIZE = 48
+DEVICE     = "cuda"
+DTYPE      = torch.bfloat16
+
+
+def load_model(ckpt_path, model_config, device, dtype):
+    model = ChronFlux(
+        in_channels=16,
+        hidden_size=1024,
+        num_double_layers=8,
+        num_single_layers=16,
+        num_heads=16,
+        style_dim=512,
+        use_moe=True,
+        use_depth_attn=True,
+        use_adain=True,
+        moe_params={"num_experts": model_config.moe.num_experts}
+    ).to(device).to(dtype)
+
+    state = torch.load(ckpt_path, map_location="cpu")
+    if isinstance(state, dict) and "model_state_dict" in state:
+        state = state["model_state_dict"]
+
+    model.load_state_dict(state, strict=True)
+    model.eval().requires_grad_(False)
+
+    print(f"[INFO] Loaded checkpoint: {ckpt_path}")
+    return model
+
+
+@torch.no_grad()
+def infer_1step(model, ae, content_lat, depth_small, style_vec, phi, device, dtype):
+    t = torch.tensor([0.0], device=device, dtype=dtype)
+    z = torch.randn_like(content_lat)
+
+    v, _ = model(z, content_lat, depth_small, style_vec, t, phi)
+    z_final = (z + v * 1.0).clamp(-8, 8)
+
+    return ae.decode(z_final).clamp(-1, 1)
+
+
+def save_tensor(tensor, path):
+    img = tensor.squeeze(0).cpu().float()
+    img = ((img + 1.0) / 2.0).clamp(0, 1)
+    img = (img.permute(1, 2, 0).numpy() * 255).astype(np.uint8)
+    Image.fromarray(img).save(path)
+
+
+def main():
+    model_config = get_config("q1_paper")
+
+    print("[INFO] Loading AutoEncoder...")
+    ae = load_ae(DEVICE).to(DEVICE).to(DTYPE).eval().requires_grad_(False)
+
+    print("[INFO] Loading Style Encoder (DINOv2)...")
+    style_enc = EncoderFactory.create_encoder(
+        "dinov2", DEVICE, 512, False, model_size="base"
+    ).to(DEVICE).to(DTYPE).eval().requires_grad_(False)
+
+    print("[INFO] Loading Depth Estimator...")
+    depth_est = DepthEstimator(DEVICE)
+
+    print("[INFO] Loading distilled model...")
+    model = load_model(CHECKPOINT, model_config, DEVICE, DTYPE)
+
+    transform = transforms.Compose([
+        transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+
+    print("[INFO] Encoding style image...")
+    style_t = transform(
+        Image.open(FIXED_STYLE_PATH).convert("RGB")
+    ).unsqueeze(0).to(DEVICE).to(DTYPE)
+
+    with torch.no_grad():
+        style_vec = style_enc(style_t, return_spatial=True)
+
+    print(f"[INFO] Style vector shape: {style_vec.shape}")
+
+    # ===== LOOP DATA PAIRS =====
+    for pair in DATA_PAIRS:
+        content_dir = pair["input"]
+        output_dir  = pair["output"]
+
+        os.makedirs(output_dir, exist_ok=True)
+
+        if not os.path.exists(content_dir):
+            print(f"[WARNING] Missing folder: {content_dir}")
+            continue
+
+        content_files = sorted([
+            f for f in os.listdir(content_dir)
+            if f.lower().endswith((".jpg", ".jpeg", ".png"))
+        ])
+
+        if not content_files:
+            print(f"[WARNING] No images in: {content_dir}")
+            continue
+
+        print(f"\n[DATASET]")
+        print(f"[INPUT ] {content_dir}")
+        print(f"[OUTPUT] {output_dir}")
+        print(f"[INFO  ] {len(content_files)} images")
+
+        for fname in tqdm(content_files, desc=os.path.basename(content_dir)):
+            stem = os.path.splitext(fname)[0]
+
+            try:
+                content_tensor = transform(
+                    Image.open(os.path.join(content_dir, fname)).convert("RGB")
+                ).unsqueeze(0).to(DEVICE).to(DTYPE)
+
+                with torch.no_grad():
+                    depth_full = depth_est.predict(
+                        content_tensor.float(), IMAGE_SIZE
+                    ).to(DTYPE)
+
+                    depth_small = F.interpolate(
+                        depth_full,
+                        (DEPTH_SIZE, DEPTH_SIZE),
+                        mode="nearest"
+                    )
+
+                    content_lat = ae.encode(content_tensor)
+
+                # ===== LOOP HOURS =====
+                for h in HOURS:
+                    phi = torch.tensor(
+                        [h * (2 * PI / 24.0)],
+                        device=DEVICE,
+                        dtype=DTYPE
+                    )
+
+                    out = infer_1step(
+                        model,
+                        ae,
+                        content_lat,
+                        depth_small,
+                        style_vec,
+                        phi,
+                        DEVICE,
+                        DTYPE
+                    )
+
+                    save_path = os.path.join(
+                        output_dir,
+                        f"{stem}_h{h:02d}.png"
+                    )
+
+                    save_tensor(out, save_path)
+
+            except Exception as e:
+                print(f"\n[ERROR] {fname}: {e}")
+                import traceback
+                traceback.print_exc()
+
+        print(f"[DONE] {content_dir}: {len(content_files)*len(HOURS)} files")
+
+    print("\n🎉 ALL DATASETS DONE")
+
+
+if __name__ == "__main__":
+    main()

infer_multistep.py

@@ -0,0 +1,300 @@
+import os
+import sys
+import argparse
+import math
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from torchvision.utils import save_image
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+from tqdm import tqdm
+
+sys.path.insert(0, "/workspace/RIFT-9226")
+
+from flux.utils.util import load_ae
+from flux.chronflux import ChronFlux
+from flux.modules.advanced_encoders import EncoderFactory, DepthEstimator
+from config import get_config
+
+PI = np.pi
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Helpers
+# ─────────────────────────────────────────────────────────────────────────────
+
+def add_label(tensor_img, text):
+    img = (tensor_img.clamp(-1, 1) + 1) / 2
+    img = img.squeeze(0).permute(1, 2, 0).cpu().float().numpy()
+    img = (img * 255).astype(np.uint8)
+    pil_img = Image.fromarray(img)
+    draw = ImageDraw.Draw(pil_img)
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 22)
+    except Exception:
+        font = ImageFont.load_default()
+    x, y = 6, 6
+    for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
+        draw.text((x + dx, y + dy), text, fill="black", font=font)
+    draw.text((x, y), text, fill="white", font=font)
+    out = np.array(pil_img).astype(np.float32) / 255.0
+    out = torch.from_numpy(out).permute(2, 0, 1).unsqueeze(0)
+    return out * 2 - 1
+
+
+def load_model(checkpoint_path, model_config, device, dtype):
+    model = ChronFlux(
+        in_channels=16,
+        hidden_size=1024,
+        num_double_layers=8,
+        num_single_layers=16,
+        num_heads=16,
+        style_dim=512,
+        use_moe=True,
+        use_depth_attn=True,
+        use_adain=True,
+        moe_params={"num_experts": model_config.moe.num_experts}
+    ).to(device).to(dtype)
+
+    state = torch.load(checkpoint_path, map_location="cpu")
+    if "model_state_dict" in state:
+        state = state["model_state_dict"]
+    model.load_state_dict(state, strict=True)
+    model.eval()
+    return model
+
+
+def get_cosine_schedule(num_steps):
+    return [1.0 - math.cos((i / num_steps) * math.pi / 2.0) for i in range(num_steps)]
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Core denoising loop  (shared by single-seed and multi-seed paths)
+# ─────────────────────────────────────────────────────────────────────────────
+
+def denoise(
+        model, z_init, content_lat, depth_small, style_vec,
+        phi, device, dtype, denoise_steps,
+        latent_clamp=20.0,          # clamp during ODE integration
+        latent_clamp_final=15.0,    # tighter clamp right before decode
+):
+    """
+    Heun 2nd-order rectified-flow ODE solver.
+
+    Returns the final latent (NOT yet decoded).
+    Caller is responsible for ae.decode(z).
+    """
+    t_schedule = get_cosine_schedule(denoise_steps)
+    z = z_init.clone()
+
+    for i in range(denoise_steps):
+        t_val = t_schedule[i]
+        dt    = (t_schedule[i + 1] - t_schedule[i]) if i < denoise_steps - 1 else (1.0 - t_schedule[i])
+
+        t_cur = torch.tensor([t_val], device=device, dtype=dtype)
+
+        # ── Euler predictor ─────────────────────────────────────────────────
+        v1, _ = model(z, content_lat, depth_small, style_vec, t_cur, phi)
+        z_mid  = (z + v1 * dt).clamp(-latent_clamp, latent_clamp)
+
+        # ── Heun corrector ───────────────────────────────────────────────────
+        t_next_val = t_schedule[i + 1] if i < denoise_steps - 1 else 1.0
+        t_next = torch.tensor([t_next_val], device=device, dtype=dtype)
+        v2, _ = model(z_mid, content_lat, depth_small, style_vec, t_next, phi)
+
+        z = z + (v1 + v2) * 0.5 * dt
+        z = z.clamp(-latent_clamp, latent_clamp)               # keep ODE stable
+
+    # ── Final pre-decode clamp ───────────────────────────────────────────────
+    z = z.clamp(-latent_clamp_final, latent_clamp_final)
+    return z
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  Per-image inference
+# ─────────────────────────────────────────────────────────────────────────────
+
+def run_inference(
+        model, ae, style_enc, depth_est,
+        content_path, style_path,
+        device, dtype,
+        denoise_steps, image_size, depth_size,
+        noise_alpha=0.92,      # [DEPRECATED] kept for signature compatibility
+        num_seeds=1,           # >1 → average over multiple trajectories
+        hours=None,            # list of hours to render; default = 0..24
+):
+    """
+    Returns a strip tensor:
+        [content | style | depth | hour_0 | hour_1 | ... | hour_24]
+    """
+    if hours is None:
+        hours = list(range(25))      # 0 h … 24 h
+
+    transform = transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+
+    content_tensor = transform(Image.open(content_path).convert("RGB")).unsqueeze(0).to(device).to(dtype)
+    style_tensor   = transform(Image.open(style_path).convert("RGB")).unsqueeze(0).to(device).to(dtype)
+
+    with torch.no_grad():
+        # ── Encode once ─────────────────────────────────────────────────────
+        depth_full  = depth_est.predict(content_tensor.float(), image_size).to(dtype)
+        depth_small = F.interpolate(depth_full, size=(depth_size, depth_size), mode="nearest")
+
+        content_lat = ae.encode(content_tensor)
+
+        style_vec = style_enc(style_tensor, return_spatial=True)
+
+        depth_vis = (depth_full.repeat(1, 3, 1, 1) - 0.5) / 0.5
+
+        generated = []
+        for h in hours:
+            phi_val = h * (2 * PI / 24.0)
+            phi     = torch.tensor([phi_val], device=device, dtype=dtype)
+
+            if num_seeds == 1:
+                # ── Single trajectory ────────────────────────────────────────
+                z_init = torch.randn_like(content_lat)
+
+                z_final = denoise(
+                    model, z_init, content_lat, depth_small, style_vec,
+                    phi, device, dtype, denoise_steps
+                )
+                out = ae.decode(z_final).clamp(-1, 1)
+
+            else:
+                decoded_sum = torch.zeros(1, 3, image_size, image_size, device=device, dtype=torch.float32)
+                for _ in range(num_seeds):
+                    z_init = torch.randn_like(content_lat)
+
+                    z_final = denoise(
+                        model, z_init, content_lat, depth_small, style_vec,
+                        phi, device, dtype, denoise_steps
+                    )
+                    decoded_sum += ae.decode(z_final).clamp(-1, 1).float()
+                out = (decoded_sum / num_seeds).clamp(-1, 1)
+
+            generated.append(add_label(out.float(), f"{h}h").to(device))
+
+    content_labeled = add_label(content_tensor.float(), "Content").to(device)
+    style_labeled   = add_label(style_tensor.float(),   "Style").to(device)
+    depth_labeled   = add_label(depth_vis.float(),      "Depth").to(device)
+
+    strip = torch.cat([content_labeled, style_labeled, depth_labeled] + generated, dim=3)
+    return strip
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+#  CLI entry-point
+# ─────────────────────────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="ChronFlux 24-hour inference")
+    parser.add_argument("--checkpoint",    type=str,   default="/workspace/RIFT-9226/checkpoints/rift_epoch_30.pt")
+    parser.add_argument("--content_dir",   type=str,   default="/workspace/RIFT/rafael_ds_ver2/img_test/cnt")
+    parser.add_argument("--style_dir",     type=str,   default="/workspace/RIFT/rafael_ds_ver2/img_test/sty")
+    parser.add_argument("--output_dir",    type=str,   default="/workspace/RIFT-9226/inference_outputs")
+    parser.add_argument("--denoise_steps", type=int,   default=50,
+                        help="More steps = better quality. 50 is a good default; try 75 for maximum quality.")
+    parser.add_argument("--image_size",    type=int,   default=384)
+    parser.add_argument("--depth_size",    type=int,   default=48)
+    parser.add_argument("--noise_alpha",   type=float, default=0.85,
+                        help="[DEPRECATED] No longer used. Model requires pure noise init.")
+    parser.add_argument("--num_seeds",     type=int,   default=1,
+                        help="Average over N random seeds (slow but reduces noise artifacts).")
+    parser.add_argument("--device",        type=str,   default="cuda")
+    parser.add_argument("--dtype",         type=str,   default="bfloat16")
+    args = parser.parse_args()
+
+    device = args.device
+    dtype  = torch.bfloat16 if args.dtype == "bfloat16" else torch.float32
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    print(f"[INFO] Loading config...")
+    model_config = get_config("q1_paper")
+
+    print(f"[INFO] Loading AutoEncoder...")
+    ae = load_ae(device).to(device).to(dtype).eval().requires_grad_(False)
+
+    print(f"[INFO] Loading Style Encoder (DINOv2)...")
+    style_enc = EncoderFactory.create_encoder(
+        "dinov2", device, 512, False, model_size="base"
+    ).to(device).to(dtype).eval().requires_grad_(False)
+
+    print(f"[INFO] Loading Depth Estimator...")
+    depth_est = DepthEstimator(device)
+
+    print(f"[INFO] Loading ChronFlux from {args.checkpoint}...")
+    model = load_model(args.checkpoint, model_config, device, dtype)
+    model.requires_grad_(False)
+
+    print(f"[INFO] denoise_steps={args.denoise_steps} | num_seeds={args.num_seeds} | Pure Noise Init: ENABLED")
+
+    content_files = sorted([
+        f for f in os.listdir(args.content_dir)
+        if f.lower().endswith((".jpg", ".jpeg", ".png"))
+    ])
+    style_files = sorted([
+        f for f in os.listdir(args.style_dir)
+        if f.lower().endswith((".jpg", ".jpeg", ".png"))
+    ])
+
+    if not content_files:
+        print("[ERROR] No content images found.")
+        return
+    if not style_files:
+        print("[ERROR] No style images found.")
+        return
+
+    print(f"[INFO] Found {len(content_files)} content images and {len(style_files)} style images.")
+
+    for content_file in tqdm(content_files, desc="Content images"):
+        content_path = os.path.join(args.content_dir, content_file)
+        content_stem = os.path.splitext(content_file)[0]
+
+        for style_file in style_files:
+            style_path = os.path.join(args.style_dir, style_file)
+            style_stem = os.path.splitext(style_file)[0]
+
+            out_name = f"{content_stem}__sty_{style_stem}.png"
+            out_path = os.path.join(args.output_dir, out_name)
+
+            if os.path.exists(out_path):
+                print(f"[SKIP] {out_name} already exists.")
+                continue
+
+            try:
+                strip = run_inference(
+                    model=model,
+                    ae=ae,
+                    style_enc=style_enc,
+                    depth_est=depth_est,
+                    content_path=content_path,
+                    style_path=style_path,
+                    device=device,
+                    dtype=dtype,
+                    denoise_steps=args.denoise_steps,
+                    image_size=args.image_size,
+                    depth_size=args.depth_size,
+                    noise_alpha=args.noise_alpha,
+                    num_seeds=args.num_seeds,
+                )
+
+                save_image(strip, out_path, normalize=True, value_range=(-1, 1))
+                print(f"[SAVED] {out_path}")
+
+            except Exception as e:
+                print(f"[ERROR] Failed on {content_file} x {style_file}: {e}")
+                import traceback
+                traceback.print_exc()
+
+    print(f"\n[DONE] All outputs saved to: {args.output_dir}")
+
+
+if __name__ == "__main__":
+    main()

infer_multistep_vid.py

@@ -0,0 +1,300 @@
+import os
+import sys
+import argparse
+import math
+import torch
+import torch.nn.functional as F
+import cv2
+import numpy as np
+from tqdm import tqdm
+from PIL import Image
+from torchvision import transforms
+
+try:
+    from moviepy import VideoFileClip
+    MOVIEPY_AVAILABLE = True
+except ImportError:
+    MOVIEPY_AVAILABLE = False
+
+sys.path.insert(0, "/workspace/ChronFLUX")
+
+from flux.utils.util import load_ae
+from flux.chronflux import ChronFlux
+from flux.modules.advanced_encoders import EncoderFactory, DepthEstimator
+from config import get_config
+
+PI = math.pi
+
+
+def load_model(checkpoint_path, model_config, device, dtype):
+    model = ChronFlux(
+        in_channels=16,
+        hidden_size=1024,
+        num_double_layers=8,
+        num_single_layers=16,
+        num_heads=16,
+        style_dim=512,
+        use_moe=True,
+        use_depth_attn=True,
+        use_adain=True,
+        moe_params={"num_experts": model_config.moe.num_experts}
+    ).to(device).to(dtype)
+
+    state = torch.load(checkpoint_path, map_location="cpu")
+    if "model_state_dict" in state:
+        state = state["model_state_dict"]
+    model.load_state_dict(state, strict=True)
+    model.eval()
+    model.requires_grad_(False)
+    return model
+
+
+def get_cosine_schedule(num_steps):
+    ts = []
+    for i in range(num_steps):
+        t_val = 1.0 - math.cos((i / num_steps) * math.pi / 2.0)
+        ts.append(t_val)
+    return ts
+
+
+def get_safe_resolution(width, height, max_height=720):
+    if height <= max_height and width % 16 == 0 and height % 16 == 0:
+        return width, height
+    ratio = max_height / height
+    new_width = int(width * ratio)
+    new_height = max_height
+    if new_width % 16 != 0:
+        new_width = new_width - (new_width % 16)
+    if new_height % 16 != 0:
+        new_height = new_height - (new_height % 16)
+    return new_width, new_height
+
+
+def tensor_to_numpy_bgr(tensor):
+    img = tensor.clamp(-1, 1)
+    img = (img + 1.0) / 2.0
+    img = img.squeeze(0).permute(1, 2, 0).cpu().float().numpy()
+    img = (img * 255).astype(np.uint8)
+    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    return img_bgr
+
+
+def inference_single_frame(
+        model, ae, depth_est,
+        frame_tensor,
+        style_vec_cached,
+        phi, device, dtype,
+        denoise_steps, depth_size,
+        base_noise
+):
+    with torch.no_grad():
+        image_size = frame_tensor.shape[-1]
+
+        depth_full = depth_est.predict(frame_tensor.float(), image_size).to(dtype)
+        depth_small = F.interpolate(depth_full, size=(depth_size, depth_size), mode="nearest")
+
+        # content_gray = transforms.functional.rgb_to_grayscale(frame_tensor, num_output_channels=3)
+        # content_lat = ae.encode(content_gray)
+        content_lat = ae.encode(frame_tensor)
+
+        style_vec = style_vec_cached
+
+        t_schedule = get_cosine_schedule(denoise_steps)
+
+        z = base_noise.clone()
+
+        for i in range(denoise_steps):
+            t_val = t_schedule[i]
+            if i < denoise_steps - 1:
+                dt = t_schedule[i + 1] - t_schedule[i]
+            else:
+                dt = 1.0 - t_schedule[i]
+
+            t = torch.tensor([t_val], device=device, dtype=dtype)
+
+            v1, _ = model(z, content_lat, depth_small, style_vec, t, phi)
+            z_mid = (z + v1 * dt).clamp(-8, 8)
+
+            if i < denoise_steps - 1:
+                t_next_val = t_schedule[i + 1]
+            else:
+                t_next_val = 1.0
+            t_next = torch.tensor([t_next_val], device=device, dtype=dtype)
+            v2, _ = model(z_mid, content_lat, depth_small, style_vec, t_next, phi)
+
+            z = z + (v1 + v2) * 0.5 * dt
+            z = z.clamp(-8, 8)
+
+        z = z.clamp(-6, 6)
+        out = ae.decode(z).clamp(-1, 1)
+        return out
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ChronFlux Video Inference - Day/Night Cycle (Fast)")
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--video_path", type=str, required=True)
+    parser.add_argument("--style_path", type=str, required=True, help="Ảnh style ban đêm")
+    parser.add_argument("--output_dir", type=str, default="/workspace/ChronFLUX/video_outputs")
+    parser.add_argument("--denoise_steps", type=int, default=50)
+    parser.add_argument("--image_size", type=int, default=384)
+    parser.add_argument("--depth_size", type=int, default=48)
+    parser.add_argument("--max_height", type=int, default=384)
+    parser.add_argument("--hour_start", type=float, default=0.0)
+    parser.add_argument("--hour_end", type=float, default=24.0)
+    parser.add_argument("--cycle", action="store_true")
+    parser.add_argument("--no_audio", action="store_true")
+    parser.add_argument("--device", type=str, default="cuda")
+    parser.add_argument("--dtype", type=str, default="bfloat16")
+    args = parser.parse_args()
+
+    device = args.device
+    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float32
+
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    print("[INFO] Loading config...")
+    model_config = get_config("q1_paper")
+
+    print("[INFO] Loading AutoEncoder...")
+    ae = load_ae(device).to(device).to(dtype).eval().requires_grad_(False)
+
+    print("[INFO] Loading Style Encoder (DINOv2)...")
+    style_enc = EncoderFactory.create_encoder(
+        "dinov2", device, 512, False, model_size="base"
+    ).to(device).to(dtype).eval().requires_grad_(False)
+
+    print("[INFO] Loading Depth Estimator...")
+    depth_est = DepthEstimator(device)
+
+    print(f"[INFO] Loading ChronFlux from {args.checkpoint}...")
+    model = load_model(args.checkpoint, model_config, device, dtype)
+
+    transform = transforms.Compose([
+        transforms.Resize((args.image_size, args.image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.5], [0.5])
+    ])
+
+    print("[INFO] Loading & encoding style image (cached for all frames)...")
+    style_pil = Image.open(args.style_path).convert("RGB")
+    style_tensor = transform(style_pil).unsqueeze(0).to(device).to(dtype)
+    with torch.no_grad():
+        style_vec_cached = style_enc(style_tensor, return_spatial=True)
+    print(f"[INFO] Style vector cached: {style_vec_cached.shape}")
+    # ─────────────────────────────────────────────────────────────────────────
+
+    print("[INFO] Opening video...")
+    cap = cv2.VideoCapture(args.video_path)
+    if not cap.isOpened():
+        raise IOError(f"Không thể mở video: {args.video_path}")
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    proc_w, proc_h = get_safe_resolution(orig_w, orig_h, args.max_height)
+    print(f"[INFO] Video: {orig_w}x{orig_h} → resize {proc_w}x{proc_h}, {fps:.2f} FPS, {total_frames} frames")
+    print(f"[INFO] Hour range: {args.hour_start}h → {args.hour_end}h, denoise_steps={args.denoise_steps} (Heun)")
+    print(f"[INFO] Speedup: style_enc skipped for {total_frames - 1} frames")
+
+    video_name = os.path.splitext(os.path.basename(args.video_path))[0]
+    temp_path = os.path.join(args.output_dir, f"{video_name}_temp.mp4")
+    output_path = os.path.join(args.output_dir, f"{video_name}_chronflux.mp4")
+
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(temp_path, fourcc, fps, (proc_w, proc_h))
+
+    frame_idx = 0
+    pbar = tqdm(total=total_frames, desc="Processing frames")
+
+    base_noise = None
+
+    while True:
+        ret, frame_bgr = cap.read()
+        if not ret:
+            break
+
+        frame_resized = cv2.resize(frame_bgr, (proc_w, proc_h), interpolation=cv2.INTER_AREA)
+        frame_rgb = cv2.cvtColor(frame_resized, cv2.COLOR_BGR2RGB)
+        frame_pil = Image.fromarray(frame_rgb)
+        frame_tensor_inf = transform(frame_pil).unsqueeze(0).to(device).to(dtype)
+
+        if base_noise is None:
+            with torch.no_grad():
+                content_gray = transforms.functional.rgb_to_grayscale(frame_tensor_inf, num_output_channels=3)
+                first_lat = ae.encode(content_gray)
+                base_noise = torch.randn_like(first_lat)
+            print(f"[INFO] Base noise initialized: {base_noise.shape}")
+
+        if args.cycle:
+            progress = (frame_idx % total_frames) / max(total_frames - 1, 1)
+            if (frame_idx // total_frames) % 2 == 0:
+                hour = args.hour_start + progress * (args.hour_end - args.hour_start)
+            else:
+                hour = args.hour_end - progress * (args.hour_end - args.hour_start)
+        else:
+            progress = frame_idx / max(total_frames - 1, 1)
+            hour = args.hour_start + progress * (args.hour_end - args.hour_start)
+
+        phi_val = hour * (2 * PI / 24.0)
+        phi = torch.tensor([phi_val], device=device, dtype=dtype)
+
+        out = inference_single_frame(
+            model=model,
+            ae=ae,
+            depth_est=depth_est,
+            frame_tensor=frame_tensor_inf,
+            style_vec_cached=style_vec_cached,
+            phi=phi,
+            device=device,
+            dtype=dtype,
+            denoise_steps=args.denoise_steps,
+            depth_size=args.depth_size,
+            base_noise=base_noise
+        )
+
+        out_bgr = tensor_to_numpy_bgr(out)
+        out_bgr_resized = cv2.resize(out_bgr, (proc_w, proc_h), interpolation=cv2.INTER_LINEAR)
+        writer.write(out_bgr_resized)
+
+        frame_idx += 1
+        pbar.update(1)
+        pbar.set_postfix({"hour": f"{hour:.1f}h", "phi": f"{phi_val:.3f}"})
+
+    cap.release()
+    writer.release()
+    pbar.close()
+    print(f"[INFO] Finished processing {frame_idx} frames.")
+
+    if not args.no_audio and MOVIEPY_AVAILABLE:
+        print("[INFO] Adding audio from original video...")
+        try:
+            original_clip = VideoFileClip(args.video_path)
+            generated_clip = VideoFileClip(temp_path)
+            if original_clip.audio is not None:
+                final_clip = generated_clip.set_audio(original_clip.audio)
+                final_clip.write_videofile(output_path, codec="libx264", audio_codec="aac", logger="bar")
+                original_clip.close()
+                generated_clip.close()
+                os.remove(temp_path)
+                print("[INFO] Audio added successfully.")
+            else:
+                original_clip.close()
+                generated_clip.close()
+                os.rename(temp_path, output_path)
+                print("[INFO] No audio track found in original video.")
+        except Exception as e:
+            print(f"[WARN] Could not add audio: {e}")
+            os.rename(temp_path, output_path)
+    else:
+        if not MOVIEPY_AVAILABLE and not args.no_audio:
+            print("[WARN] moviepy not installed, skipping audio.")
+        os.rename(temp_path, output_path)
+
+    print(f"\n[DONE] Output saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,15 @@
+torch>=2.6.0
+torchvision>=0.19.0
+numpy>=1.24.0
+pillow>=10.3.0
+tqdm
+einops
+safetensors
+transformers>=4.48.0
+accelerate
+opencv-python
+sentencepiece
+huggingface_hub
+timm>=0.9.0
+lpips
+flash-attn>=2.0.0