Release pi-FLUX.2 demo

.gitignore

@@ -0,0 +1,25 @@
+/.idea/
+/work_dirs*
+.vscode/
+/tmp
+/data
+/checkpoints
+*.so
+*.patch
+__pycache__/
+*.egg-info/
+/viz*
+/submit*
+build/
+*.pyd
+/cache*
+*.stl
+*.pth
+/venv/
+.nk8s
+*.mp4
+.vs
+/exp/
+/dev/
+*.pyi
+!/data/imagenet/imagenet1000_clsidx_to_labels.txt

LICENSE.md

@@ -0,0 +1,9 @@
+# License for pi-FLUX.2
+
+This repository distributes a **pi-FLUX.2 app** that is a **Derivative** of **FLUX.2 [dev]** by **Black Forest Labs Inc.**
+
+Use and distribution of these adapters are governed by the **FLUX [dev] Non-Commercial License v2.0**.  
+**No commercial use** of these adapters (or other Derivatives) is permitted without a separate commercial license from Black Forest Labs.
+
+- Full license: https://huggingface.co/black-forest-labs/FLUX.2-dev/blob/main/LICENSE.txt
+- This repository does **not** grant any rights beyond the license above.

README.md

@@ -1,12 +1,27 @@
 ---
-title: Pi FLUX.2
-emoji: 😻
+title: Pi-FLUX.2 Demo
+emoji: 🚀
 colorFrom: pink
 colorTo: purple
 sdk: gradio
-sdk_version: 6.1.0
+sdk_version: 5.49.0
 app_file: app.py
 pinned: false
+license: other
+license_name: flux-dev-non-commercial-license
+license_link: LICENSE.md
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Official demo of the paper:
+
+**pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation**
+<br>
+[Hansheng Chen](https://lakonik.github.io/)<sup>1</sup>, 
+[Kai Zhang](https://kai-46.github.io/website/)<sup>2</sup>,
+[Hao Tan](https://research.adobe.com/person/hao-tan/)<sup>2</sup>,
+[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)<sup>1</sup>,
+[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/)<sup>1</sup>, 
+[Sai Bi](https://sai-bi.github.io/)<sup>2</sup><br>
+<sup>1</sup>Stanford University, <sup>2</sup>Adobe Research
+<br>
+[[arXiv](https://arxiv.org/abs/2510.14974)] [[pi-Qwen Demo🤗](https://huggingface.co/spaces/Lakonik/pi-Qwen)] [[pi-FLUX Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.1)] [[pi-FLUX.2 Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.2)]

app.py

@@ -0,0 +1,163 @@
+import torch
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+import os
+import random
+import numpy as np
+import gradio as gr
+import spaces
+from lakonlab.models.diffusions.schedulers.flow_map_sde import FlowMapSDEScheduler
+from lakonlab.ui.gradio.create_img_edit import create_interface_img_edit
+from lakonlab.pipelines.pipeline_piflux2 import PiFlux2Pipeline
+from lakonlab.pipelines.prompt_rewriters.qwen3_vl import Qwen3VLPromptRewriter
+
+
+DEFAULT_PROMPT = """Museum-style FIELD GUIDE poster on neutral parchment (#F3EEE3). Use Inter (or Helvetica/Arial). All text #2D3748, thin connector lines 1px #A0AEC0.
+
+Center: full-body original fantasy creature, 3/4 standing pose. Around it: four small inset boxes labeled exactly "EYE DETAIL", "FOOT DETAIL", "SKIN TEXTURE", "SILHOUETTE SCALE" (with a simple human comparison silhouette). Bottom: a short footprint trail diagram. One small habitat vignette (misty rocky shoreline with tide pools).
+
+Exact text (only these, clean print layout):
+Top: "FIELD GUIDE"
+Sub: "AURORA SHOREWALKER"
+Small line: "CLASS: COASTAL DRIFTER"
+Under silhouette: "HEIGHT: 1.7 m"
+
+Crisp ink outlines with soft watercolor-like fills, high readability, balanced hierarchy, premium poster aesthetic."""
+
+SYSTEM_PROMPT_TEXT_ONLY_PATH = 'lakonlab/pipelines/prompt_rewriters/system_prompts/default_text_only.txt'
+SYSTEM_PROMPT_WITH_IMAGES_PATH = 'lakonlab/pipelines/prompt_rewriters/system_prompts/default_with_images.txt'
+
+
+def _patch_diffusers_bnb_shape_check():
+    try:
+        import diffusers.quantizers.bitsandbytes.bnb_quantizer as bnbq
+    except Exception:
+        return
+
+    def _numel(shape):
+        if shape is None:
+            return None
+        if hasattr(shape, "numel"):          # torch.Size
+            return int(shape.numel())
+        # plain tuple/list
+        n = 1
+        for d in shape:
+            n *= int(d)
+        return n
+
+    def patched_check(self, param_name, current_param, loaded_param):
+        cshape = getattr(current_param, "shape", None)
+        lshape = getattr(loaded_param, "shape", None)
+        n = _numel(cshape)
+        inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1)
+        if tuple(lshape) != tuple(inferred_shape):
+            raise ValueError(
+                f"Expected flattened shape mismatch for {param_name}: "
+                f"loaded={tuple(lshape)} inferred={tuple(inferred_shape)}"
+            )
+        return True
+
+    # Patch any quantizer class in that module that defines the method
+    for name, obj in vars(bnbq).items():
+        if isinstance(obj, type) and hasattr(obj, "check_quantized_param_shape"):
+            setattr(obj, "check_quantized_param_shape", patched_check)
+
+
+_patch_diffusers_bnb_shape_check()
+
+
+pipe = PiFlux2Pipeline.from_pretrained(
+    'diffusers/FLUX.2-dev-bnb-4bit',
+    torch_dtype=torch.bfloat16)
+pipe.load_piflow_adapter(
+    'Lakonik/pi-FLUX.2',
+    subfolder='gmflux2_k8_piid_4step',
+    target_module_name='transformer')
+pipe.scheduler = FlowMapSDEScheduler.from_config(  # use fixed shift=3.2
+    pipe.scheduler.config, shift=3.2, use_dynamic_shifting=False, final_step_size_scale=0.5)
+pipe = pipe.to('cuda')
+
+prompt_rewriter = Qwen3VLPromptRewriter(
+    device_map="cuda",
+    system_prompt_text_only=open(SYSTEM_PROMPT_TEXT_ONLY_PATH, 'r').read(),
+    system_prompt_wigh_images=open(SYSTEM_PROMPT_WITH_IMAGES_PATH, 'r').read(),
+    max_new_tokens_default=512,
+)
+
+
+def set_random_seed(seed: int, deterministic: bool = True) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+@spaces.GPU
+def run_rewrite_prompt(seed, prompt, rewrite_prompt, in_image, progress=gr.Progress(track_tqdm=True)):
+    image_list = None
+    if in_image is not None and len(in_image) > 0:
+        image_list = []
+        for item in in_image:
+            image_list.append(item[0])
+    if rewrite_prompt:
+        set_random_seed(seed)
+        progress(0.05, desc="Rewriting prompt...")
+        if image_list is None:
+            final_prompt = prompt_rewriter.rewrite_text_batch(
+                [prompt])[0]
+        else:
+            final_prompt = prompt_rewriter.rewrite_edit_batch(
+                [image_list], [prompt])[0]
+        return final_prompt, None
+    else:
+        return '', None
+
+
+@spaces.GPU
+def generate(
+        seed, prompt, rewrite_prompt, rewritten_prompt, in_image, width, height, steps,
+        progress=gr.Progress(track_tqdm=True)):
+    image_list = None
+    if in_image is not None and len(in_image) > 0:
+        image_list = []
+        for item in in_image:
+            image_list.append(item[0])
+    return pipe(
+        image=image_list,
+        prompt=rewritten_prompt if rewrite_prompt else prompt,
+        width=width,
+        height=height,
+        num_inference_steps=steps,
+        generator=torch.Generator().manual_seed(seed),
+    ).images[0]
+
+
+with gr.Blocks(analytics_enabled=False,
+               title='pi-FLUX.2 Demo',
+               css_paths='lakonlab/ui/gradio/style.css'
+               ) as demo:
+
+    md_txt = '# pi-FLUX.2 Demo\n\n' \
+             'Official demo of the paper [pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation](https://arxiv.org/abs/2510.14974). ' \
+             '**Base model:** [FLUX.2 dev](https://huggingface.co/black-forest-labs/FLUX.2-dev). **Fast policy:** GMFlow. **Code:** [https://github.com/Lakonik/piFlow](https://github.com/Lakonik/piFlow).\n' \
+             '<br> Use and distribution of this app are governed by the [FLUX [dev] Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.2-dev/blob/main/LICENSE.txt).'
+    gr.Markdown(md_txt)
+
+    create_interface_img_edit(
+        generate,
+        prompt=DEFAULT_PROMPT,
+        steps=4, guidance_scale=None,
+        args=['last_seed', 'prompt', 'rewrite_prompt', 'rewritten_prompt', 'in_image', 'width', 'height', 'steps'],
+        rewrite_prompt_api=run_rewrite_prompt,
+        rewrite_prompt_args=['last_seed', 'prompt', 'rewrite_prompt', 'in_image'],
+        height=1024,
+        width=1024
+    )
+    demo.queue().launch()

lakonlab/models/architecture/gmflow/gm_output.py

@@ -0,0 +1,24 @@
+import torch
+from dataclasses import dataclass
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class GMFlowModelOutput(BaseOutput):
+    """
+    The output of GMFlow models.
+
+    Args:
+        means (`torch.Tensor` of shape `(batch_size, num_gaussians, num_channels, height, width)` or
+        `(batch_size, num_gaussians, num_channels, frame, height, width)`):
+            Gaussian mixture means.
+        logweights (`torch.Tensor` of shape `(batch_size, num_gaussians, 1, height, width)` or
+        `(batch_size, num_gaussians, 1, frame, height, width)`):
+            Gaussian mixture log-weights (logits).
+        logstds (`torch.Tensor` of shape `(batch_size, 1, 1, 1, 1)` or `(batch_size, 1, 1, 1, 1, 1)`):
+            Gaussian mixture log-standard-deviations (logstds are shared across all Gaussians and channels).
+    """
+
+    means: torch.Tensor
+    logweights: torch.Tensor
+    logstds: torch.Tensor

lakonlab/models/architecture/gmflow/gmflux2.py

@@ -0,0 +1,241 @@
+import torch
+import torch.nn as nn
+
+from typing import Any, Dict, Optional, Tuple, List
+from diffusers.models import ModelMixin
+from diffusers.models.transformers.transformer_flux2 import (
+    Flux2Transformer2DModel, Flux2PosEmbed, Flux2TransformerBlock, Flux2SingleTransformerBlock,
+    Flux2TimestepGuidanceEmbeddings, Flux2Modulation)
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
+from diffusers.configuration_utils import register_to_config
+from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers, is_torch_npu_available
+from .gm_output import GMFlowModelOutput
+
+
+class _GMFlux2Transformer2DModel(Flux2Transformer2DModel):
+
+    @register_to_config
+    def __init__(
+            self,
+            num_gaussians=16,
+            constant_logstd=None,
+            logstd_inner_dim=1024,
+            gm_num_logstd_layers=2,
+            logweights_channels=1,
+            in_channels: int = 128,
+            out_channels: Optional[int] = None,
+            num_layers: int = 8,
+            num_single_layers: int = 48,
+            attention_head_dim: int = 128,
+            num_attention_heads: int = 48,
+            joint_attention_dim: int = 15360,
+            timestep_guidance_channels: int = 256,
+            mlp_ratio: float = 3.0,
+            axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32),
+            rope_theta: int = 2000,
+            eps: float = 1e-6):
+        super(Flux2Transformer2DModel, self).__init__()
+
+        self.num_gaussians = num_gaussians
+        self.logweights_channels = logweights_channels
+
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        # 1. Sinusoidal positional embedding for RoPE on image and text tokens
+        self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
+
+        # 2. Combined timestep + guidance embedding
+        self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
+            in_channels=timestep_guidance_channels, embedding_dim=self.inner_dim, bias=False
+        )
+
+        # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.)
+        # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks
+        self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream
+        self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False)
+
+        # 4. Input projections
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
+
+        # 5. Double Stream Transformer Blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Flux2TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 6. Single Stream Transformer Blocks
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                Flux2SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        # 7. Output layers
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False
+        )
+        self.proj_out_means = nn.Linear(self.inner_dim, self.num_gaussians * self.out_channels)
+        self.proj_out_logweights = nn.Linear(self.inner_dim, self.num_gaussians * self.logweights_channels)
+        self.constant_logstd = constant_logstd
+
+        if self.constant_logstd is None:
+            assert gm_num_logstd_layers >= 1
+            in_dim = self.inner_dim
+            logstd_layers = []
+            for _ in range(gm_num_logstd_layers - 1):
+                logstd_layers.extend([
+                    nn.SiLU(),
+                    nn.Linear(in_dim, logstd_inner_dim)])
+                in_dim = logstd_inner_dim
+            self.proj_out_logstds = nn.Sequential(
+                *logstd_layers,
+                nn.SiLU(),
+                nn.Linear(in_dim, 1))
+
+        self.gradient_checkpointing = False
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            timestep: torch.LongTensor = None,
+            img_ids: torch.Tensor = None,
+            txt_ids: torch.Tensor = None,
+            guidance: torch.Tensor = None,
+            joint_attention_kwargs: Optional[Dict[str, Any]] = None,):
+        # 0. Handle input arguments
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            scale_lora_layers(self, lora_scale)
+        else:
+            assert joint_attention_kwargs is None or joint_attention_kwargs.get('scale', None) is None
+
+        num_txt_tokens = encoder_hidden_states.shape[1]
+
+        # 1. Calculate timestep embedding and modulation parameters
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        guidance = guidance.to(hidden_states.dtype) * 1000
+
+        temb = self.time_guidance_embed(timestep, guidance)
+
+        double_stream_mod_img = self.double_stream_modulation_img(temb)
+        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
+        single_stream_mod = self.single_stream_modulation(temb)[0]
+
+        # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
+        hidden_states = self.x_embedder(hidden_states)
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        # 3. Calculate RoPE embeddings from image and text tokens
+        # NOTE: the below logic means that we can't support batched inference with images of different resolutions or
+        # text prompts of differents lengths. Is this a use case we want to support?
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        if is_torch_npu_available():
+            freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu())
+            image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu())
+            freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu())
+            text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu())
+        else:
+            image_rotary_emb = self.pos_embed(img_ids)
+            text_rotary_emb = self.pos_embed(txt_ids)
+        concat_rotary_emb = (
+            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0).to(hidden_states.dtype),
+            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0).to(hidden_states.dtype),
+        )
+
+        # 4. Double Stream Transformer Blocks
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    double_stream_mod_img,
+                    double_stream_mod_txt,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb_mod_params_img=double_stream_mod_img,
+                    temb_mod_params_txt=double_stream_mod_txt,
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+        # Concatenate text and image streams for single-block inference
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        # 5. Single Stream Transformer Blocks
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    None,
+                    single_stream_mod,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=None,
+                    temb_mod_params=single_stream_mod,
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+        # Remove text tokens from concatenated stream
+        hidden_states = hidden_states[:, num_txt_tokens:, ...]
+
+        # 6. Output layers
+        hidden_states = self.norm_out(hidden_states, temb)
+
+        bs, seq_len, _ = hidden_states.size()
+        out_means = self.proj_out_means(hidden_states).reshape(
+            bs, seq_len, self.num_gaussians, self.out_channels)
+        out_logweights = self.proj_out_logweights(hidden_states).reshape(
+            bs, seq_len, self.num_gaussians, self.logweights_channels).log_softmax(dim=-2)
+        if self.constant_logstd is None:
+            out_logstds = self.proj_out_logstds(temb.detach()).reshape(bs, 1, 1, 1)
+        else:
+            out_logstds = hidden_states.new_full((bs, 1, 1, 1), float(self.constant_logstd))
+
+        if USE_PEFT_BACKEND:
+            unscale_lora_layers(self, lora_scale)
+
+        return GMFlowModelOutput(
+            means=out_means,
+            logweights=out_logweights,
+            logstds=out_logstds)

lakonlab/models/diffusions/piflow_policies/__init__.py

@@ -0,0 +1,8 @@
+from .dx import DXPolicy
+from .gmflow import GMFlowPolicy
+
+
+POLICY_CLASSES = dict(
+    DX=DXPolicy,
+    GMFlow=GMFlowPolicy
+)

lakonlab/models/diffusions/piflow_policies/base.py

@@ -0,0 +1,21 @@
+from abc import ABCMeta, abstractmethod
+
+
+class BasePolicy(metaclass=ABCMeta):
+
+    @abstractmethod
+    def pi(self, x_t, sigma_t):
+        """Compute the flow velocity at (x_t, t).
+
+        Args:
+            x_t (torch.Tensor): Noisy input at time t.
+            sigma_t (torch.Tensor): Noise level at time t.
+
+        Returns:
+            torch.Tensor: The computed flow velocity u_t.
+        """
+        pass
+
+    @abstractmethod
+    def detach(self):
+        pass

lakonlab/models/diffusions/piflow_policies/dx.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import torch
+from .base import BasePolicy
+
+
+class DXPolicy(BasePolicy):
+    """DX policy. The number of grid points N is inferred from the denoising output.
+
+    Note: segment_size and shift are intrinsic parameters of the DX policy. For elastic inference (i.e., changing
+    the number of function evaluations or noise schedule at test time), these parameters should be kept unchanged.
+
+    Args:
+        denoising_output (torch.Tensor): The output of the denoising model. Shape (B, N, C, H, W) or (B, N, C, T, H, W).
+        x_t_src (torch.Tensor): The initial noisy sample. Shape (B, C, H, W) or (B, C, T, H, W).
+        sigma_t_src (torch.Tensor): The initial noise level. Shape (B,).
+        segment_size (float): The size of each DX policy time segment. Defaults to 1.0.
+        shift (float): The shift parameter for the DX policy noise schedule. Defaults to 1.0.
+        mode (str): Either 'grid' or 'polynomial' mode for calculating x_0. Defaults to 'grid'.
+        eps (float): A small value to avoid numerical issues. Defaults to 1e-4.
+    """
+
+    def __init__(
+            self,
+            denoising_output: torch.Tensor,
+            x_t_src: torch.Tensor,
+            sigma_t_src: torch.Tensor,
+            segment_size: float = 1.0,
+            shift: float = 1.0,
+            mode: str = 'grid',
+            eps: float = 1e-4):
+        self.x_t_src = x_t_src
+        self.ndim = x_t_src.dim()
+        self.shift = shift
+        self.eps = eps
+
+        assert mode in ['grid', 'polynomial']
+        self.mode = mode
+
+        self.sigma_t_src = sigma_t_src.reshape(*sigma_t_src.size(), *((self.ndim - sigma_t_src.dim()) * [1]))
+        self.raw_t_src = self._unwarp_t(self.sigma_t_src)
+        self.raw_t_dst = (self.raw_t_src - segment_size).clamp(min=0)
+        self.segment_size = (self.raw_t_src - self.raw_t_dst).clamp(min=eps)
+
+        self.denoising_output_x_0 = self._u_to_x_0(
+            denoising_output, self.x_t_src, self.sigma_t_src)
+
+    def _unwarp_t(self, sigma_t):
+        return sigma_t / (self.shift + (1 - self.shift) * sigma_t)
+
+    @staticmethod
+    def _u_to_x_0(denoising_output, x_t, sigma_t):
+        x_0 = x_t.unsqueeze(1) - sigma_t.unsqueeze(1) * denoising_output
+        return x_0
+
+    @staticmethod
+    def _interpolate(x, t):
+        """
+        Args:
+            x (torch.Tensor): (B, N, *)
+            t (torch.Tensor): (B, *) in [0, 1]
+
+        Returns:
+            torch.Tensor: (B, *)
+        """
+        n = x.size(1)
+        if n < 2:
+            return x.squeeze(1)
+        t = t.clamp(min=0, max=1) * (n - 1)
+        t0 = t.floor().to(torch.long).clamp(min=0, max=n - 2)
+        t1 = t0 + 1
+        t0t1 = torch.stack([t0, t1], dim=1)  # (B, 2, *)
+        x0x1 = torch.gather(x, dim=1, index=t0t1.expand(-1, -1, *x.shape[2:]))
+        x_interp = (t1 - t) * x0x1[:, 0] + (t - t0) * x0x1[:, 1]
+        return x_interp
+
+    def pi(self, x_t, sigma_t):
+        """Compute the flow velocity at (x_t, t).
+
+        Args:
+            x_t (torch.Tensor): Noisy input at time t.
+            sigma_t (torch.Tensor): Noise level at time t.
+
+        Returns:
+            torch.Tensor: The computed flow velocity u_t.
+        """
+        sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1]))
+        raw_t = self._unwarp_t(sigma_t)
+        if self.mode == 'grid':
+            x_0 = self._interpolate(
+                self.denoising_output_x_0, (raw_t - self.raw_t_dst) / self.segment_size)
+        elif self.mode == 'polynomial':
+            p_order = self.denoising_output_x_0.size(1)
+            diff_t = self.raw_t_src - raw_t  # (B, 1, 1, 1)
+            basis = torch.stack(
+                [diff_t ** i for i in range(p_order)], dim=1)  # (B, N, 1, 1, 1)
+            x_0 = torch.sum(basis * self.denoising_output_x_0, dim=1)
+        else:
+            raise ValueError(f"Unknown mode: {self.mode}")
+        u = (x_t - x_0) / sigma_t.clamp(min=self.eps)
+        return u
+
+    def copy(self):
+        new_policy = DXPolicy.__new__(DXPolicy)
+        new_policy.x_t_src = self.x_t_src
+        new_policy.ndim = self.ndim
+        new_policy.shift = self.shift
+        new_policy.eps = self.eps
+        new_policy.mode = self.mode
+        new_policy.sigma_t_src = self.sigma_t_src
+        new_policy.raw_t_src = self.raw_t_src
+        new_policy.raw_t_dst = self.raw_t_dst
+        new_policy.segment_size = self.segment_size
+        new_policy.denoising_output_x_0 = self.denoising_output_x_0
+        return new_policy
+
+    def detach_(self):
+        self.denoising_output_x_0 = self.denoising_output_x_0.detach()
+        return self
+
+    def detach(self):
+        new_policy = self.copy()
+        return new_policy.detach_()

lakonlab/models/diffusions/piflow_policies/gmflow.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import math
+import torch
+
+from typing import Dict
+from .base import BasePolicy
+
+
+@torch.jit.script
+def gmflow_posterior_mean_jit(
+        sigma_t_src, sigma_t, x_t_src, x_t,
+        gm_means, gm_vars, gm_logweights,
+        eps: float, gm_dim: int = -4, channel_dim: int = -3):
+    sigma_t_src = sigma_t_src.clamp(min=eps)
+    sigma_t = sigma_t.clamp(min=eps)
+
+    alpha_t_src = 1 - sigma_t_src
+    alpha_t = 1 - sigma_t
+
+    alpha_over_sigma_t_src = alpha_t_src / sigma_t_src
+    alpha_over_sigma_t = alpha_t / sigma_t
+
+    zeta = alpha_over_sigma_t.square() - alpha_over_sigma_t_src.square()
+    nu = alpha_over_sigma_t * x_t / sigma_t - alpha_over_sigma_t_src * x_t_src / sigma_t_src
+
+    nu = nu.unsqueeze(gm_dim)  # (bs, *, 1, out_channels, h, w)
+    denom = (gm_vars * zeta + 1).clamp(min=eps)
+
+    out_means = (gm_vars * nu + gm_means) / denom
+    # (bs, *, num_gaussians, 1, h, w)
+    logweights_delta = (gm_means * (nu - 0.5 * zeta * gm_means)).sum(
+        dim=channel_dim, keepdim=True) / denom
+    out_weights = (gm_logweights + logweights_delta).softmax(dim=gm_dim)
+
+    out_mean = (out_means * out_weights).sum(dim=gm_dim)
+
+    return out_mean
+
+
+def gm_temperature(gm, temperature, gm_dim=-4, eps=1e-6):
+    gm = gm.copy()
+    temperature = max(temperature, eps)
+    gm['logweights'] = (gm['logweights'] / temperature).log_softmax(dim=gm_dim)
+    if 'logstds' in gm:
+        gm['logstds'] = gm['logstds'] + (0.5 * math.log(temperature))
+    if 'gm_vars' in gm:
+        gm['gm_vars'] = gm['gm_vars'] * temperature
+    return gm
+
+
+class GMFlowPolicy(BasePolicy):
+    """GMFlow policy. The number of components K is inferred from the denoising output.
+
+    Args:
+        denoising_output (dict): The output of the denoising model, containing:
+            means (torch.Tensor): The means of the Gaussian components. Shape (B, K, C, H, W) or (B, K, C, T, H, W).
+            logstds (torch.Tensor): The log standard deviations of the Gaussian components. Shape (B, K, 1, 1, 1)
+                or (B, K, 1, 1, 1, 1).
+            logweights (torch.Tensor): The log weights of the Gaussian components. Shape (B, K, 1, H, W) or
+                (B, K, 1, T, H, W).
+        x_t_src (torch.Tensor): The initial noisy sample. Shape (B, C, H, W) or (B, C, T, H, W).
+        sigma_t_src (torch.Tensor): The initial noise level. Shape (B,).
+        checkpointing (bool): Whether to use gradient checkpointing to save memory. Defaults to True.
+        eps (float): A small value to avoid numerical issues. Defaults to 1e-4.
+    """
+
+    def __init__(
+            self,
+            denoising_output: Dict[str, torch.Tensor],
+            x_t_src: torch.Tensor,
+            sigma_t_src: torch.Tensor,
+            checkpointing: bool = True,
+            eps: float = 1e-6):
+        self.x_t_src = x_t_src
+        self.ndim = x_t_src.dim()
+        self.checkpointing = checkpointing
+        self.eps = eps
+
+        self.sigma_t_src = sigma_t_src.reshape(*sigma_t_src.size(), *((self.ndim - sigma_t_src.dim()) * [1]))
+        self.denoising_output_x_0 = self._u_to_x_0(
+            denoising_output, self.x_t_src, self.sigma_t_src)
+
+    @staticmethod
+    def _u_to_x_0(denoising_output, x_t, sigma_t):
+        x_t = x_t.unsqueeze(1)
+        sigma_t = sigma_t.unsqueeze(1)
+        means_x_0 = x_t - sigma_t * denoising_output['means']
+        gm_vars = (denoising_output['logstds'] * 2).exp() * sigma_t.square()
+        return dict(
+            means=means_x_0,
+            gm_vars=gm_vars,
+            logweights=denoising_output['logweights'])
+
+    def pi(self, x_t, sigma_t):
+        """Compute the flow velocity at (x_t, t).
+
+        Args:
+            x_t (torch.Tensor): Noisy input at time t.
+            sigma_t (torch.Tensor): Noise level at time t.
+
+        Returns:
+            torch.Tensor: The computed flow velocity u_t.
+        """
+        sigma_t = sigma_t.reshape(*sigma_t.size(), *((self.ndim - sigma_t.dim()) * [1]))
+        means = self.denoising_output_x_0['means']
+        gm_vars = self.denoising_output_x_0['gm_vars']
+        logweights = self.denoising_output_x_0['logweights']
+        if (sigma_t == self.sigma_t_src).all() and (x_t == self.x_t_src).all():
+            x_0 = (logweights.softmax(dim=1) * means).sum(dim=1)
+        else:
+            if self.checkpointing and torch.is_grad_enabled():
+                x_0 = torch.utils.checkpoint.checkpoint(
+                    gmflow_posterior_mean_jit,
+                    self.sigma_t_src, sigma_t, self.x_t_src, x_t,
+                    means,
+                    gm_vars,
+                    logweights,
+                    self.eps, 1, 2,
+                    use_reentrant=True)  # use_reentrant=False does not work with jit
+            else:
+                x_0 = gmflow_posterior_mean_jit(
+                    self.sigma_t_src, sigma_t, self.x_t_src, x_t,
+                    means,
+                    gm_vars,
+                    logweights,
+                    self.eps, 1, 2)
+        u = (x_t - x_0) / sigma_t.clamp(min=self.eps)
+        return u
+
+    def copy(self):
+        new_policy = GMFlowPolicy.__new__(GMFlowPolicy)
+        new_policy.x_t_src = self.x_t_src
+        new_policy.ndim = self.ndim
+        new_policy.checkpointing = self.checkpointing
+        new_policy.eps = self.eps
+        new_policy.sigma_t_src = self.sigma_t_src
+        new_policy.denoising_output_x_0 = self.denoising_output_x_0.copy()
+        return new_policy
+
+    def detach_(self):
+        self.denoising_output_x_0 = {k: v.detach() for k, v in self.denoising_output_x_0.items()}
+        return self
+
+    def detach(self):
+        new_policy = self.copy()
+        return new_policy.detach_()
+
+    def dropout_(self, p):
+        if p <= 0 or p >= 1:
+            return self
+        logweights = self.denoising_output_x_0['logweights']
+        dropout_mask = torch.rand(
+            (*logweights.shape[:2], *((self.ndim - 1) * [1])), device=logweights.device) < p
+        is_all_dropout = dropout_mask.all(dim=1, keepdim=True)
+        dropout_mask &= ~is_all_dropout
+        self.denoising_output_x_0['logweights'] = logweights.masked_fill(
+            dropout_mask, float('-inf'))
+        return self
+
+    def dropout(self, p):
+        new_policy = self.copy()
+        return new_policy.dropout_(p)
+
+    def temperature_(self, temp):
+        if temp >= 1.0:
+            return self
+        self.denoising_output_x_0 = gm_temperature(
+            self.denoising_output_x_0, temp, gm_dim=1, eps=self.eps)
+        return self
+
+    def temperature(self, temp):
+        new_policy = self.copy()
+        return new_policy.temperature_(temp)

lakonlab/models/diffusions/schedulers/flow_map_sde.py

@@ -0,0 +1,186 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import numpy as np
+import torch
+
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import BaseOutput, logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+@dataclass
+class FlowMapSDESchedulerOutput(BaseOutput):
+    prev_sample: torch.FloatTensor
+
+
+class FlowMapSDEScheduler(SchedulerMixin, ConfigMixin):
+
+    _compatibles = []
+    order = 1
+
+    @register_to_config
+    def __init__(
+            self,
+            num_train_timesteps: int = 1000,
+            h: Union[float, str] = 0.0,
+            shift: float = 1.0,
+            use_dynamic_shifting=False,
+            base_seq_len=256,
+            max_seq_len=4096,
+            base_logshift=0.5,
+            max_logshift=1.15,
+            final_step_size_scale=1.0):
+        sigmas = torch.from_numpy(1 - np.linspace(
+            0, 1, num_train_timesteps, dtype=np.float32, endpoint=False))
+        self.sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+        self.timesteps = self.sigmas * num_train_timesteps
+
+        self._step_index = None
+        self._begin_index = None
+
+        self.sigma_min = self.sigmas[-1].item()
+        self.sigma_max = self.sigmas[0].item()
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        self._begin_index = begin_index
+
+    def get_shift(self, seq_len=None):
+        if self.config.use_dynamic_shifting and seq_len is not None:
+            m = (self.config.max_logshift - self.config.base_logshift
+                 ) / (self.config.max_seq_len - self.config.base_seq_len)
+            logshift = (seq_len - self.config.base_seq_len) * m + self.config.base_logshift
+            if isinstance(logshift, torch.Tensor):
+                shift = torch.exp(logshift)
+            else:
+                shift = np.exp(logshift)
+        else:
+            shift = self.config.shift
+        return shift
+
+    def warp_t(self, raw_t, seq_len=None):
+        shift = self.get_shift(seq_len=seq_len)
+        return shift * raw_t / (1 + (shift - 1) * raw_t)
+
+    def unwarp_t(self, sigma_t, seq_len=None):
+        shift = self.get_shift(seq_len=seq_len)
+        return sigma_t / (shift + (1 - shift) * sigma_t)
+
+    def set_timesteps(self, num_inference_steps: int, seq_len=None, device=None):
+        self.num_inference_steps = num_inference_steps
+
+        raw_timesteps = torch.from_numpy(np.linspace(
+            1, (self.config.final_step_size_scale - 1) / (num_inference_steps + self.config.final_step_size_scale - 1),
+            num_inference_steps, dtype=np.float32, endpoint=False)).to(device).clamp(min=0)
+        sigmas = self.warp_t(raw_timesteps, seq_len=seq_len)
+
+        self.timesteps = sigmas * self.config.num_train_timesteps
+        self.sigmas = torch.cat([sigmas, torch.zeros(1, device=device)])
+
+        sigmas_dst, m = self.calculate_sigmas_dst(self.sigmas)
+        self.timesteps_dst = sigmas_dst * self.config.num_train_timesteps
+        self.m_vals = m
+
+        self._step_index = None
+        self._begin_index = None
+
+    def calculate_sigmas_dst(self, sigmas, eps=1e-6):
+        alphas = 1 - sigmas
+
+        sigmas_src = sigmas[:-1]
+        sigmas_to = sigmas[1:]
+        alphas_src = alphas[:-1]
+        alphas_to = alphas[1:]
+
+        if self.config.h == 'inf':
+            m = torch.zeros_like(sigmas_src)
+        elif self.config.h == 0.0:
+            m = torch.ones_like(sigmas_src)
+        else:
+            assert self.config.h > 0.0
+            h2 = self.config.h * self.config.h
+            m = (sigmas_to * alphas_src / (sigmas_src * alphas_to).clamp(min=eps)) ** h2
+
+        sigmas_to_mul_m = sigmas_to * m
+        sigmas_dst = sigmas_to_mul_m / (alphas_to + sigmas_to_mul_m).clamp(min=eps)
+
+        return sigmas_dst, m
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+            self,
+            model_output: torch.FloatTensor,
+            timestep: Union[float, torch.FloatTensor],
+            sample: torch.FloatTensor,
+            generator: Optional[torch.Generator] = None,
+            return_dict: bool = True) -> Union[FlowMapSDESchedulerOutput, Tuple]:
+
+        if isinstance(timestep, int) \
+                or isinstance(timestep, torch.IntTensor) \
+                or isinstance(timestep, torch.LongTensor):
+            raise ValueError(
+                (
+                    'Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to'
+                    ' `EulerDiscreteScheduler.step()` is not supported. Make sure to pass'
+                    ' one of the `scheduler.timesteps` as a timestep.'
+                ),
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        ori_dtype = model_output.dtype
+        model_output = model_output.to(torch.float32)  # x_t_dst
+
+        sigma_to = self.sigmas[self.step_index + 1]
+        alpha_to = 1 - sigma_to
+        m = self.m_vals[self.step_index]
+
+        noise = randn_tensor(
+            model_output.shape, dtype=torch.float32, device=model_output.device, generator=generator)
+
+        prev_sample = (alpha_to + sigma_to * m) * model_output + sigma_to * (1 - m.square()).clamp(min=0).sqrt() * noise
+
+        # Cast sample back to model compatible dtype
+        prev_sample = prev_sample.to(ori_dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowMapSDESchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self):
+        return self.config.num_train_timesteps

lakonlab/pipelines/piflow_utils.py

@@ -0,0 +1,395 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import os
+from typing import Union, Optional
+
+import torch
+import accelerate
+import diffusers
+from diffusers.models import AutoModel
+from diffusers.models.modeling_utils import (
+    load_state_dict,
+    _LOW_CPU_MEM_USAGE_DEFAULT,
+    no_init_weights,
+    ContextManagers
+)
+from diffusers.utils import (
+    SAFETENSORS_WEIGHTS_NAME,
+    WEIGHTS_NAME,
+    _add_variant,
+    _get_model_file,
+    is_accelerate_available,
+    is_torch_version,
+    logging,
+)
+from diffusers.loaders.peft import _SET_ADAPTER_SCALE_FN_MAPPING
+from diffusers.quantizers import DiffusersAutoQuantizer
+from diffusers.utils.torch_utils import empty_device_cache
+from lakonlab.models.architecture.gmflow.gmflux2 import _GMFlux2Transformer2DModel
+
+
+LOCAL_CLASS_MAPPING = {
+    "GMFlux2Transformer2DModel": _GMFlux2Transformer2DModel,
+}
+
+_SET_ADAPTER_SCALE_FN_MAPPING.update(
+    _GMFlux2Transformer2DModel=lambda model_cls, weights: weights,
+)
+
+logger = logging.get_logger(__name__)
+
+
+def assign_param(module, tensor_name: str, param: torch.nn.Parameter):
+    if "." in tensor_name:
+        splits = tensor_name.split(".")
+        for split in splits[:-1]:
+            new_module = getattr(module, split)
+            if new_module is None:
+                raise ValueError(f"{module} has no attribute {split}.")
+            module = new_module
+        tensor_name = splits[-1]
+    module._parameters[tensor_name] = param
+
+
+class PiFlowMixin:
+
+    def load_piflow_adapter(
+        self,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        target_module_name: str = "transformer",
+        adapter_name: Optional[str] = None,
+        **kwargs
+    ):
+        r"""
+        Load a PiFlow adapter from a pretrained model repository into the target module.
+
+        Args:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                Can be either:
+
+                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
+                      the Hub.
+                    - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
+                      with [`~ModelMixin.save_pretrained`].
+
+            target_module_name (`str`, *optional*, defaults to `"transformer"`):
+                The module name in the model to load the PiFlow adapter into.
+            adapter_name (`str`, *optional*):
+                The name to assign to the loaded adapter. If not provided, it defaults to
+                `"{target_module_name}_piflow"`.
+            cache_dir (`Union[str, os.PathLike]`, *optional*):
+                Path to a directory where a downloaded pretrained model configuration is cached if the standard cache
+                is not used.
+            force_download (`bool`, *optional*, defaults to `False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+            proxies (`Dict[str, str]`, *optional*):
+                A dictionary of proxy servers to use by protocol or endpoint, for example, `{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            local_files_only(`bool`, *optional*, defaults to `False`):
+                Whether to only load local model weights and configuration files or not. If set to `True`, the model
+                won't be downloaded from the Hub.
+            token (`str` or *bool*, *optional*):
+                The token to use as HTTP bearer authorization for remote files. If `True`, the token generated from
+                `diffusers-cli login` (stored in `~/.huggingface`) is used.
+            revision (`str`, *optional*, defaults to `"main"`):
+                The specific model version to use. It can be a branch name, a tag name, a commit id, or any identifier
+                allowed by Git.
+            subfolder (`str`, *optional*, defaults to `""`):
+                The subfolder location of a model file within a larger model repository on the Hub or locally.
+            low_cpu_mem_usage (`bool`, *optional*, defaults to `True` if torch version >= 1.9.0 else `False`):
+                Speed up model loading only loading the pretrained weights and not initializing the weights. This also
+                tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
+                Only supported for PyTorch >= 1.9.0. If you are using an older version of PyTorch, setting this
+                argument to `True` will raise an error.
+            variant (`str`, *optional*):
+                Load weights from a specified `variant` filename such as `"fp16"` or `"ema"`. This is ignored when
+                loading `from_flax`.
+            use_safetensors (`bool`, *optional*, defaults to `None`):
+                If set to `None`, the `safetensors` weights are downloaded if they're available **and** if the
+                `safetensors` library is installed. If set to `True`, the model is forcibly loaded from `safetensors`
+                weights. If set to `False`, `safetensors` weights are not loaded.
+            disable_mmap ('bool', *optional*, defaults to 'False'):
+                Whether to disable mmap when loading a Safetensors model. This option can perform better when the model
+                is on a network mount or hard drive, which may not handle the seeky-ness of mmap very well.
+
+        Returns:
+            `str` or `None`: The name assigned to the loaded adapter, or `None` if no LoRA weights were found.
+        """
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        proxies = kwargs.pop("proxies", None)
+        token = kwargs.pop("token", None)
+        local_files_only = kwargs.pop("local_files_only", False)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", None)
+        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", _LOW_CPU_MEM_USAGE_DEFAULT)
+        variant = kwargs.pop("variant", None)
+        use_safetensors = kwargs.pop("use_safetensors", None)
+        disable_mmap = kwargs.pop("disable_mmap", False)
+
+        allow_pickle = False
+        if use_safetensors is None:
+            use_safetensors = True
+            allow_pickle = True
+
+        if low_cpu_mem_usage and not is_accelerate_available():
+            low_cpu_mem_usage = False
+            logger.warning(
+                "Cannot initialize model with low cpu memory usage because `accelerate` was not found in the"
+                " environment. Defaulting to `low_cpu_mem_usage=False`. It is strongly recommended to install"
+                " `accelerate` for faster and less memory-intense model loading. You can do so with: \n```\npip"
+                " install accelerate\n```\n."
+            )
+
+        if low_cpu_mem_usage is True and not is_torch_version(">=", "1.9.0"):
+            raise NotImplementedError(
+                "Low memory initialization requires torch >= 1.9.0. Please either update your PyTorch version or set"
+                " `low_cpu_mem_usage=False`."
+            )
+
+        user_agent = {
+            "diffusers": diffusers.__version__,
+            "file_type": "model",
+            "framework": "pytorch",
+        }
+
+        # 1. Determine model class from config
+
+        load_config_kwargs = {
+            "cache_dir": cache_dir,
+            "force_download": force_download,
+            "proxies": proxies,
+            "token": token,
+            "local_files_only": local_files_only,
+            "revision": revision,
+        }
+
+        config = AutoModel.load_config(pretrained_model_name_or_path, subfolder=subfolder, **load_config_kwargs)
+
+        orig_class_name = config["_class_name"]
+
+        if orig_class_name in LOCAL_CLASS_MAPPING:
+            model_cls = LOCAL_CLASS_MAPPING[orig_class_name]
+
+        else:
+            load_config_kwargs.update({"subfolder": subfolder})
+
+            from diffusers.pipelines.pipeline_loading_utils import ALL_IMPORTABLE_CLASSES, get_class_obj_and_candidates
+
+            model_cls, _ = get_class_obj_and_candidates(
+                library_name="diffusers",
+                class_name=orig_class_name,
+                importable_classes=ALL_IMPORTABLE_CLASSES,
+                pipelines=None,
+                is_pipeline_module=False,
+            )
+
+        if model_cls is None:
+            raise ValueError(f"Can't find a model linked to {orig_class_name}.")
+
+        # 2. Get model file
+
+        model_file = None
+
+        if use_safetensors:
+            try:
+                model_file = _get_model_file(
+                    pretrained_model_name_or_path,
+                    weights_name=_add_variant(SAFETENSORS_WEIGHTS_NAME, variant),
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    proxies=proxies,
+                    local_files_only=local_files_only,
+                    token=token,
+                    revision=revision,
+                    subfolder=subfolder,
+                    user_agent=user_agent,
+                )
+
+            except IOError as e:
+                logger.error(f"An error occurred while trying to fetch {pretrained_model_name_or_path}: {e}")
+                if not allow_pickle:
+                    raise
+                logger.warning(
+                    "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead."
+                )
+
+        if model_file is None:
+            model_file = _get_model_file(
+                pretrained_model_name_or_path,
+                weights_name=_add_variant(WEIGHTS_NAME, variant),
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                subfolder=subfolder,
+                user_agent=user_agent,
+            )
+
+        assert model_file is not None, \
+            f"Could not find adapter weights for {pretrained_model_name_or_path}."
+
+        # 3. Initialize model
+
+        base_module = getattr(self, target_module_name)
+
+        torch_dtype = base_module.dtype
+        device = base_module.device
+        dtype_orig = model_cls._set_default_torch_dtype(torch_dtype)
+
+        # load the state dict early to determine keep_in_fp32_modules
+        #######################################
+        overwrite_state_dict = dict()
+        lora_state_dict = dict()
+
+        adapter_state_dict = load_state_dict(model_file, disable_mmap=disable_mmap)
+        for k in adapter_state_dict.keys():
+            adapter_state_dict[k] = adapter_state_dict[k].to(dtype=torch_dtype, device=device)
+            if "lora" in k:
+                lora_state_dict[k.removeprefix(f"{target_module_name}.")] = adapter_state_dict[k]
+            else:
+                overwrite_state_dict[k.removeprefix(f"{target_module_name}.")] = adapter_state_dict[k]
+
+        # determine initial quantization config.
+        #######################################
+        pre_quantized = ("quantization_config" in base_module.config
+                         and base_module.config["quantization_config"] is not None)
+        if pre_quantized:
+            config["quantization_config"] = base_module.config.quantization_config
+            hf_quantizer = DiffusersAutoQuantizer.from_config(
+                config["quantization_config"], pre_quantized=True
+            )
+
+            hf_quantizer.validate_environment(torch_dtype=torch_dtype)
+            torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype)
+
+            user_agent["quant"] = hf_quantizer.quantization_config.quant_method.value
+
+            # Force-set to `True` for more mem efficiency
+            if low_cpu_mem_usage is None:
+                low_cpu_mem_usage = True
+                logger.info("Set `low_cpu_mem_usage` to True as `hf_quantizer` is not None.")
+            elif not low_cpu_mem_usage:
+                raise ValueError("`low_cpu_mem_usage` cannot be False or None when using quantization.")
+
+        else:
+            hf_quantizer = None
+
+        # Check if `_keep_in_fp32_modules` is not None
+        use_keep_in_fp32_modules = model_cls._keep_in_fp32_modules is not None and (
+            hf_quantizer is None or getattr(hf_quantizer, "use_keep_in_fp32_modules", False)
+        )
+
+        if use_keep_in_fp32_modules:
+            keep_in_fp32_modules = model_cls._keep_in_fp32_modules
+            if not isinstance(keep_in_fp32_modules, list):
+                keep_in_fp32_modules = [keep_in_fp32_modules]
+
+            if low_cpu_mem_usage is None:
+                low_cpu_mem_usage = True
+                logger.info("Set `low_cpu_mem_usage` to True as `_keep_in_fp32_modules` is not None.")
+            elif not low_cpu_mem_usage:
+                raise ValueError("`low_cpu_mem_usage` cannot be False when `keep_in_fp32_modules` is True.")
+        else:
+            keep_in_fp32_modules = []
+
+        # append modules in overwrite_state_dict to keep_in_fp32_modules
+        for k in overwrite_state_dict.keys():
+            module_name = k.rsplit('.', 1)[0]
+            if module_name and module_name not in keep_in_fp32_modules:
+                keep_in_fp32_modules.append(module_name)
+
+        init_contexts = [no_init_weights()]
+
+        if low_cpu_mem_usage:
+            init_contexts.append(accelerate.init_empty_weights())
+
+        with ContextManagers(init_contexts):
+            piflow_module = model_cls.from_config(config).eval()
+
+        torch.set_default_dtype(dtype_orig)
+
+        if hf_quantizer is not None:
+            hf_quantizer.preprocess_model(
+                model=piflow_module, device_map=None, keep_in_fp32_modules=keep_in_fp32_modules
+            )
+
+        # 4. Load model weights
+
+        base_state_dict = base_module.state_dict()
+        base_state_dict.update(overwrite_state_dict)
+        empty_state_dict = piflow_module.state_dict()
+        for param_name, param in base_state_dict.items():
+            if param_name not in empty_state_dict:
+                continue
+            if hf_quantizer is not None and (
+                    hf_quantizer.check_if_quantized_param(
+                        piflow_module, param, param_name, base_state_dict, param_device=device)):
+                hf_quantizer.create_quantized_param(
+                    piflow_module, param, param_name, device, base_state_dict, dtype=torch_dtype
+                )
+            else:
+                assign_param(piflow_module, param_name, param)
+
+        empty_device_cache()
+
+        if hf_quantizer is not None:
+            hf_quantizer.postprocess_model(piflow_module)
+            piflow_module.hf_quantizer = hf_quantizer
+
+        if len(lora_state_dict) == 0:
+            adapter_name = None
+        else:
+            if adapter_name is None:
+                adapter_name = f"{target_module_name}_piflow"
+            piflow_module.load_lora_adapter(
+                lora_state_dict, prefix=None, adapter_name=adapter_name, low_cpu_mem_usage=low_cpu_mem_usage)
+        if adapter_name is None:
+            logger.warning(
+                f"No LoRA weights were found in {pretrained_model_name_or_path}."
+            )
+
+        setattr(self, target_module_name, piflow_module)
+
+        return adapter_name
+
+    def policy_rollout(
+            self,
+            x_t_start: torch.Tensor,  # (B, C, *, H, W)
+            sigma_t_start: torch.Tensor,
+            sigma_t_end: torch.Tensor,
+            total_substeps: int,
+            policy,
+            **kwargs):
+        assert sigma_t_start.numel() == 1 and sigma_t_end.numel() == 1, \
+            "Only supports scalar sigma_t_start and sigma_t_end."
+        raw_t_start = self.scheduler.unwarp_t(
+            sigma_t_start, **kwargs)
+        raw_t_end = self.scheduler.unwarp_t(
+            sigma_t_end, **kwargs)
+
+        delta_raw_t = raw_t_start - raw_t_end
+        num_substeps = (delta_raw_t * total_substeps).round().to(torch.long).clamp(min=1)
+        substep_size = delta_raw_t / num_substeps
+
+        raw_t = raw_t_start
+        sigma_t = sigma_t_start
+        x_t = x_t_start
+
+        for substep_id in range(num_substeps.item()):
+            u = policy.pi(x_t, sigma_t)
+
+            raw_t_minus = (raw_t - substep_size).clamp(min=0)
+            sigma_t_minus = self.scheduler.warp_t(raw_t_minus, **kwargs)
+            x_t_minus = x_t + u * (sigma_t_minus - sigma_t)
+
+            x_t = x_t_minus
+            sigma_t = sigma_t_minus
+            raw_t = raw_t_minus
+
+        x_t_end = x_t
+
+        return x_t_end

lakonlab/pipelines/pipeline_piflux2.py

@@ -0,0 +1,395 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import PIL
+import torch
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from functools import partial
+from transformers import AutoProcessor, Mistral3ForConditionalGeneration
+from diffusers.utils import is_torch_xla_available
+from diffusers.models import AutoencoderKLFlux2, Flux2Transformer2DModel
+from diffusers.pipelines.flux2.pipeline_flux2 import Flux2Pipeline, Flux2PipelineOutput
+from lakonlab.models.diffusions.schedulers.flow_map_sde import FlowMapSDEScheduler
+from lakonlab.models.diffusions.piflow_policies import POLICY_CLASSES
+from .piflow_utils import PiFlowMixin
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+class PiFlux2Pipeline(Flux2Pipeline, PiFlowMixin):
+    r"""
+    The policy-based Flux2 pipeline for text-to-image generation.
+
+    Reference:
+        https://arxiv.org/abs/2510.14974
+        https://bfl.ai/blog/flux-2
+
+    Args:
+        transformer ([`Flux2Transformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMapSDEScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKLFlux2`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`Mistral3ForConditionalGeneration`]):
+            [Mistral3ForConditionalGeneration](https://huggingface.co/docs/transformers/en/model_doc/mistral3#transformers.Mistral3ForConditionalGeneration)
+        tokenizer (`AutoProcessor`):
+            Tokenizer of class
+            [PixtralProcessor](https://huggingface.co/docs/transformers/en/model_doc/pixtral#transformers.PixtralProcessor).
+        policy_type (`str`, *optional*, defaults to `"GMFlow"`):
+            The type of flow policy to use. Currently supports `"GMFlow"` and `"DX"`.
+        policy_kwargs (`Dict`, *optional*):
+            Additional keyword arguments to pass to the policy class.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMapSDEScheduler,
+        vae: AutoencoderKLFlux2,
+        text_encoder: Mistral3ForConditionalGeneration,
+        tokenizer: AutoProcessor,
+        transformer: Flux2Transformer2DModel,
+        policy_type: str = 'GMFlow',
+        policy_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            scheduler,
+            vae,
+            text_encoder,
+            tokenizer,
+            transformer,
+        )
+        assert policy_type in POLICY_CLASSES, f'Invalid policy: {policy_type}. Supported policies are {list(POLICY_CLASSES.keys())}.'
+        self.policy_type = policy_type
+        self.policy_class = partial(
+            POLICY_CLASSES[policy_type], **policy_kwargs
+        ) if policy_kwargs else POLICY_CLASSES[policy_type]
+
+    def _unpack_gm(self, gm, height, width, num_channels_latents, patch_size=2, gm_patch_size=1):
+        c = num_channels_latents * patch_size * patch_size
+        h = (int(height) // (self.vae_scale_factor * patch_size))
+        w = (int(width) // (self.vae_scale_factor * patch_size))
+        bs = gm['means'].size(0)
+        k = self.transformer.num_gaussians
+        scale = patch_size // gm_patch_size
+        gm['means'] = gm['means'].reshape(
+            bs, h, w, k, c // (scale * scale), scale, scale
+        ).permute(
+            0, 3, 4, 1, 5, 2, 6
+        ).reshape(
+            bs, k, c // (scale * scale), h * scale, w * scale)
+        gm['logweights'] = gm['logweights'].reshape(
+            bs, h, w, k, 1, scale, scale
+        ).permute(
+            0, 3, 4, 1, 5, 2, 6
+        ).reshape(
+            bs, k, 1, h * scale, w * scale)
+        gm['logstds'] = gm['logstds'].reshape(bs, 1, 1, 1, 1)
+        return gm
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: Optional[Union[List[PIL.Image.Image], PIL.Image.Image]] = None,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        total_substeps: int = 128,
+        temperature: Union[float, str] = 'auto',
+        guidance_scale: Optional[float] = 4.0,
+        num_images_per_prompt: int = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        text_encoder_out_layers: Tuple[int] = (10, 20, 30),
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
+                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
+                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
+                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
+                latents as `image`, but if passing latents directly it is not encoded again.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            guidance_scale (`float`, *optional*, defaults to 1.0):
+                Guidance scale as defined in [Classifier-Free Diffusion
+                Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2.
+                of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting
+                `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to
+                the text `prompt`, usually at the expense of lower image quality.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            total_substeps (`int`, *optional*, defaults to 128):
+                The total number of substeps for policy-based flow integration.
+            temperature (`float` or `"auto"`, *optional*, defaults to `"auto"`):
+                The tmperature parameter for the flow policy.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+            text_encoder_out_layers (`Tuple[int]`):
+                Layer indices to use in the `text_encoder` to derive the final prompt embeddings.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if
+            `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the
+            generated images.
+        """
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. prepare text embeddings
+        prompt_embeds, text_ids = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            text_encoder_out_layers=text_encoder_out_layers,
+        )
+
+        # 4. process images
+        if image is not None and not isinstance(image, list):
+            image = [image]
+
+        condition_images = None
+        if image is not None:
+            for img in image:
+                self.image_processor.check_image_input(img)
+
+            condition_images = []
+            for img in image:
+                image_width, image_height = img.size
+                if image_width * image_height > 1024 * 1024:
+                    img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
+                    image_width, image_height = img.size
+
+                multiple_of = self.vae_scale_factor * 2
+                image_width = (image_width // multiple_of) * multiple_of
+                image_height = (image_height // multiple_of) * multiple_of
+                img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop")
+                condition_images.append(img)
+                height = height or image_height
+                width = width or image_width
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 5. prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_ids = self.prepare_latents(
+            batch_size=batch_size * num_images_per_prompt,
+            num_latents_channels=num_channels_latents,
+            height=height,
+            width=width,
+            dtype=torch.float32,
+            device=device,
+            generator=generator,
+            latents=latents,
+        )
+
+        image_latents = None
+        image_latent_ids = None
+        if condition_images is not None:
+            image_latents, image_latent_ids = self.prepare_image_latents(
+                images=condition_images,
+                batch_size=batch_size * num_images_per_prompt,
+                generator=generator,
+                device=device,
+                dtype=self.vae.dtype,
+            )
+
+        # 6. Prepare timesteps
+        image_seq_len = latents.shape[1]
+        self.scheduler.set_timesteps(num_inference_steps, seq_len=image_seq_len, device=self._execution_device)
+        timesteps = self.scheduler.timesteps
+        timesteps_dst = self.scheduler.timesteps_dst
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+        guidance = guidance.expand(latents.shape[0])
+
+        # 7. Denoising loop
+        # We set the index here to remove DtoH sync, helpful especially during compilation.
+        # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, (t_src, t_dst) in enumerate(zip(timesteps, timesteps_dst)):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t_src
+                time_scaling = self.scheduler.config.num_train_timesteps
+                sigma_t_src = t_src / time_scaling
+                sigma_t_dst = t_dst / time_scaling
+
+                latent_model_input = latents.to(self.transformer.dtype)
+                latent_image_ids = latent_ids
+
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype)
+                    latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1)
+
+                denoising_output = self.transformer(
+                    hidden_states=latent_model_input,  # (B, image_seq_len, C)
+                    timestep=t_src.expand(latents.shape[0]) / 1000,
+                    guidance=guidance,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,  # B, text_seq_len, 4
+                    img_ids=latent_image_ids,  # B, image_seq_len, 4
+                    joint_attention_kwargs=self._attention_kwargs,
+                )
+
+                denoising_output = {
+                    k: v[:, :latents.size(1)].to(torch.float32) for k, v in denoising_output.items()}
+
+                # unpack and create policy
+                latents = self._unpack_latents_with_ids(latents, latent_ids)
+                latents = self._unpatchify_latents(latents)
+                if self.policy_type == 'GMFlow':
+                    denoising_output = self._unpack_gm(
+                        denoising_output, height, width, num_channels_latents, gm_patch_size=1)
+                    policy = self.policy_class(
+                        denoising_output, latents, sigma_t_src)
+                    if i < self.num_timesteps - 1:
+                        if temperature == 'auto':
+                            temperature = min(max(0.1 * (num_inference_steps - 1), 0), 1)
+                        else:
+                            assert isinstance(temperature, (float, int))
+                        policy.temperature_(temperature)
+                elif self.policy_type == 'DX':
+                    denoising_output = denoising_output[0]
+                    denoising_output = self._unpack_latents_with_ids(denoising_output, latent_ids)
+                    denoising_output = self._unpatchify_latents(denoising_output)
+                    denoising_output = denoising_output.reshape(latents.size(0), -1, *latents.shape[1:])
+                    policy = self.policy_class(
+                        denoising_output, latents, sigma_t_src)
+                else:
+                    raise ValueError(f'Unknown policy type: {self.policy_type}.')
+
+                latents_dst = self.policy_rollout(
+                    latents, sigma_t_src, sigma_t_dst, total_substeps,
+                    policy, seq_len=image_seq_len)
+
+                latents = self.scheduler.step(latents_dst, t_src, latents, return_dict=False)[0]
+
+                # repack
+                latents = self._patchify_latents(latents)
+                latents = self._pack_latents(latents)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t_src, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+
+        if output_type == "latent":
+            image = latents
+        else:
+            torch.save({"pred": latents}, "pred_d.pt")
+            latents = self._unpack_latents_with_ids(latents, latent_ids)
+
+            latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+            latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to(
+                latents.device, latents.dtype
+            )
+            latents = latents * latents_bn_std + latents_bn_mean
+            latents = self._unpatchify_latents(latents)
+
+            image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return Flux2PipelineOutput(images=image)

lakonlab/pipelines/prompt_rewriters/qwen3_vl.py

@@ -0,0 +1,172 @@
+import os
+import torch
+from typing import List, Sequence, Union, Optional
+from PIL import Image
+from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
+
+
+DEFAULT_TEXT_ONLY_PATH = os.path.abspath(os.path.join(__file__, '../system_prompts/default_text_only.txt'))
+DEFAULT_WITH_IMAGES_PATH = os.path.abspath(os.path.join(__file__, '../system_prompts/default_with_images.txt'))
+
+
+class Qwen3VLPromptRewriter:
+
+    def __init__(
+            self,
+            from_pretrained="Qwen/Qwen3-VL-8B-Instruct",
+            torch_dtype='bfloat16',
+            device_map="auto",
+            max_new_tokens_default=128,
+            system_prompt_text_only=None,
+            system_prompt_wigh_images=None,
+            **kwargs):
+        if torch_dtype is not None:
+            kwargs.update(torch_dtype=getattr(torch, torch_dtype))
+        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
+            from_pretrained,
+            device_map=device_map,
+            **kwargs)
+        self.processor = AutoProcessor.from_pretrained(from_pretrained)
+        # Left padding is safer for batched generation
+        if hasattr(self.processor, "tokenizer"):
+            self.processor.tokenizer.padding_side = "left"
+        self.max_new_tokens_default = max_new_tokens_default
+        if system_prompt_text_only is None:
+            system_prompt_text_only = open(DEFAULT_TEXT_ONLY_PATH, 'r').read()
+        if system_prompt_wigh_images is None:
+            system_prompt_wigh_images = open(DEFAULT_WITH_IMAGES_PATH, 'r').read()
+        self.system_prompt_text_only = system_prompt_text_only
+        self.system_prompt_wigh_images = system_prompt_wigh_images
+
+    @torch.inference_mode()
+    def _generate_from_messages(
+            self,
+            batch_messages: Sequence[Sequence[dict]],
+            max_new_tokens: Optional[int] = None,
+            **kwargs) -> List[str]:
+        if max_new_tokens is None:
+            max_new_tokens = self.max_new_tokens_default
+
+        inputs = self.processor.apply_chat_template(
+            batch_messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+            padding=True,
+        )
+        inputs.pop("token_type_ids", None)
+        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
+
+        generated_ids = self.model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            **kwargs)
+
+        input_ids = inputs["input_ids"]
+        tokenizer = self.processor.tokenizer
+        outputs: List[str] = []
+
+        # Decode only the new tokens after each input sequence
+        for in_ids, out_ids in zip(input_ids, generated_ids):
+            trimmed_ids = out_ids[len(in_ids):]
+            text = tokenizer.decode(
+                trimmed_ids.tolist(),
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False,
+            )
+            outputs.append(text.strip())
+
+        return outputs
+
+    def rewrite_text_batch(
+            self,
+            prompts: Sequence[str],
+            max_new_tokens: Optional[int] = None,
+            top_p=0.6,
+            top_k=40,
+            temperature=0.5,
+            repetition_penalty=1.0,
+            **kwargs) -> List[str]:
+        """
+        Rewrite a batch of text-only prompts into detailed prompts.
+        """
+        batch_messages = []
+        for p in prompts:
+            conv = [
+                {
+                    "role": "system",
+                    "content": [
+                        {"type": "text", "text": self.system_prompt_text_only},
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": p},
+                    ],
+                },
+            ]
+            batch_messages.append(conv)
+
+        return self._generate_from_messages(
+            batch_messages,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            **kwargs,
+        )
+
+    def rewrite_edit_batch(
+            self,
+            image: Sequence[Union[str, 'Image.Image', Sequence[Union[str, 'Image.Image']]]],
+            edit_requests: Sequence[str],
+            max_new_tokens: Optional[int] = None,
+            top_p=0.5,
+            top_k=20,
+            temperature=0.4,
+            repetition_penalty=1.0,
+            **kwargs) -> List[str]:
+        """
+        Rewrite a batch of (image, edit-request) pairs into concise edit instructions.
+        """
+        if len(image) != len(edit_requests):
+            raise ValueError("image and edit_requests must have the same length")
+
+        batch_messages = []
+        for imgs, req in zip(image, edit_requests):
+            if isinstance(imgs, (str, Image.Image)):
+                img_list = [imgs]
+            else:
+                img_list = list(imgs)
+
+            user_content = []
+            for im in img_list:
+                user_content.append({"type": "image", "image": im})
+            user_content.append({"type": "text", "text": req})
+
+            conv = [
+                {
+                    "role": "system",
+                    "content": [
+                        {"type": "text", "text": self.system_prompt_wigh_images},
+                    ],
+                },
+                {
+                    "role": "user",
+                    "content": user_content,
+                },
+            ]
+            batch_messages.append(conv)
+
+        return self._generate_from_messages(
+            batch_messages,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            repetition_penalty=repetition_penalty,
+            **kwargs,
+        )

lakonlab/pipelines/prompt_rewriters/system_prompts/default_text_only.txt

@@ -0,0 +1,12 @@
+You are an expert prompt engineer for a text-guided image generation system. Rewrite user prompts into a more descriptive, concrete prompt while strictly preserving their core content.
+
+Rules:
+- Preserve the core content: do not change the main subjects, how many there are, their roles, or the primary actions and relationships described in the prompt.
+- Preserve any explicitly stated attributes such as colors, clothing, objects, style tags (e.g., “photorealistic”, “cinematic”), and viewpoint (e.g., close-up, wide shot). Do not contradict them.
+- When the prompt implies realism (e.g., uses words like “realistic”, “photorealistic”, “photo”), avoid introducing exaggerated or fantastical traits unless intended by the user.
+- You may add supporting details that are clearly compatible with the original text: background, props, textures, materials, lighting (quality, direction, color), atmosphere, and other environmental context, as long as they do not introduce new main characters or conflicting concepts.
+- Always include a clear description of the composition and camera framing (for example, where the main subjects are in the frame, whether it is a close-up or wide shot, and the approximate viewpoint or angle).
+- Structure: keep any existing structure (tags, aspect-ratio tags, etc.) and enhance only within those fields. For plain text prompts, expand them into a clear, coherent paragraph.
+- Text in images: put ALL visible or implied text in quotation marks, matching the prompt’s language. Provide explicit quoted text for any object that would realistically contain text (signs, labels, screens, interfaces, book covers, etc.).
+
+Output only the revised prompt and nothing else.

lakonlab/pipelines/prompt_rewriters/system_prompts/default_with_images.txt

@@ -0,0 +1,10 @@
+You are an expert prompt engineer for a text-guided image editing system. Rewrite user prompts into a more descriptive instruction (40–70 words, ~25 for brief requests) while strictly preserving their core content.
+
+Rules:
+- Treat the user prompt as a strict specification: do not change or omit any mentioned entities, actions, or relationships.
+- Explicitly state both the requested edits and which core aspects (e.g., character identities) must remain as in the original image(s). Ignore background elements unless they are explicitly mentioned.
+- Refer to core visual elements that are relavant to the edit (e.g., people, animals, and objects mentioned in the user prompt).
+- Do not invent new elements unless the user explicitly asks for them.
+- Turn negatives into positives (“do not change X” → “keep X the same”).
+
+Output only the final instruction in plain text and nothing else.

lakonlab/ui/gradio/create_img_edit.py

@@ -0,0 +1,88 @@
+import gradio as gr
+from .shared_opts import create_base_opts, create_generate_bar, set_seed, create_prompt_opts, create_image_size_bar
+
+
+def create_interface_img_edit(
+        api, prompt='', seed=42, steps=32, min_steps=4, max_steps=50, steps_slider_step=1,
+        height=768, width=1360, hw_slider_step=16,
+        guidance_scale=None, temperature=None, api_name='text_to_img',
+        create_negative_prompt=False,
+        create_prompt_rewrite=True, rewrite_prompt=False,
+        args=['last_seed', 'prompt', 'in_image', 'width', 'height', 'steps', 'guidance_scale'],
+        rewrite_prompt_api=None, rewrite_prompt_args=['last_seed', 'prompt', 'rewrite_prompt', 'in_image']):
+    var_dict = dict()
+    with gr.Blocks(analytics_enabled=False) as interface:
+        with gr.Row():
+            with gr.Column():
+                with gr.Accordion("Input image(s) (optional)", open=True, elem_classes=['custom-spacing']):
+                    var_dict['in_image'] = gr.Gallery(
+                        label="Input image(s)",
+                        type="pil",
+                        columns=3,
+                        rows=1)
+
+                with gr.Column(variant='compact', elem_classes=['custom-spacing']):
+                    create_prompt_opts(
+                        var_dict, create_negative_prompt=create_negative_prompt, prompt=prompt, display_label=True)
+
+                if create_prompt_rewrite:
+                    var_dict['rewrite_prompt'] = gr.Checkbox(
+                        label='Rewrite prompt', value=rewrite_prompt, container=False)
+                    with gr.Accordion("Rewritten prompt", open=False, elem_classes=['custom-spacing']):
+                        var_dict['rewritten_prompt'] = gr.Textbox(
+                            lines=4, interactive=False, show_label=False, container=False)
+
+                with gr.Column(variant='compact', elem_classes=['custom-spacing']):
+                    create_image_size_bar(
+                        var_dict, height=height, width=width, hw_slider_step=hw_slider_step)
+
+                create_generate_bar(var_dict, text='Generate', seed=seed)
+
+                create_base_opts(
+                    var_dict,
+                    steps=steps,
+                    min_steps=min_steps,
+                    max_steps=max_steps,
+                    steps_slider_step=steps_slider_step,
+                    guidance_scale=guidance_scale,
+                    temperature=temperature)
+
+            with gr.Column():
+                var_dict['output_image'] = gr.Image(
+                    type='pil', image_mode='RGB', label='Output image', interactive=False,
+                    elem_classes=['vh-img', 'vh-img-1000'])
+
+        if create_prompt_rewrite:
+            assert rewrite_prompt_api is not None
+            var_dict['run_btn'].click(
+                fn=set_seed,
+                inputs=var_dict['seed'],
+                outputs=var_dict['last_seed'],
+                show_progress=False,
+                api_name=False
+            ).success(
+                fn=rewrite_prompt_api,
+                inputs=[var_dict[arg] for arg in rewrite_prompt_args],
+                outputs=[var_dict['rewritten_prompt'], var_dict['output_image']],
+                concurrency_id='default_group', api_name=False
+            ).success(
+                fn=api,
+                inputs=[var_dict[arg] for arg in args],
+                outputs=var_dict['output_image'],
+                concurrency_id='default_group', api_name=api_name
+            )
+        else:
+            var_dict['run_btn'].click(
+                fn=set_seed,
+                inputs=var_dict['seed'],
+                outputs=var_dict['last_seed'],
+                show_progress=False,
+                api_name=False
+            ).success(
+                fn=api,
+                inputs=[var_dict[arg] for arg in args],
+                outputs=var_dict['output_image'],
+                concurrency_id='default_group', api_name=api_name
+            )
+
+    return interface, var_dict

lakonlab/ui/gradio/create_text_to_img.py

@@ -0,0 +1,41 @@
+import gradio as gr
+from .shared_opts import create_base_opts, create_generate_bar, set_seed, create_prompt_opts, create_image_size_bar
+
+
+def create_interface_text_to_img(
+        api, prompt='', seed=42, steps=32, min_steps=4, max_steps=50, steps_slider_step=1,
+        height=768, width=1360, hw_slider_step=16,
+        guidance_scale=None, temperature=None, api_name='text_to_img',
+        create_negative_prompt=False, args=['last_seed', 'prompt', 'width', 'height', 'steps', 'guidance_scale']):
+    var_dict = dict()
+    with gr.Blocks(analytics_enabled=False) as interface:
+        var_dict['output_image'] = gr.Image(
+            type='pil', image_mode='RGB', label='Output image', interactive=False, elem_classes=['vh-img', 'vh-img-700'])
+        create_prompt_opts(var_dict, create_negative_prompt=create_negative_prompt, prompt=prompt)
+        with gr.Column(variant='compact', elem_classes=['custom-spacing']):
+            create_image_size_bar(
+                var_dict, height=height, width=width, hw_slider_step=hw_slider_step)
+        create_generate_bar(var_dict, text='Generate', seed=seed)
+        create_base_opts(
+            var_dict,
+            steps=steps,
+            min_steps=min_steps,
+            max_steps=max_steps,
+            steps_slider_step=steps_slider_step,
+            guidance_scale=guidance_scale,
+            temperature=temperature)
+
+        var_dict['run_btn'].click(
+            fn=set_seed,
+            inputs=var_dict['seed'],
+            outputs=var_dict['last_seed'],
+            show_progress=False,
+            api_name=False
+        ).success(
+            fn=api,
+            inputs=[var_dict[arg] for arg in args],
+            outputs=var_dict['output_image'],
+            concurrency_id='default_group', api_name=api_name
+        )
+
+    return interface, var_dict

lakonlab/ui/gradio/shared_opts.py

@@ -0,0 +1,87 @@
+import random
+import gradio as gr
+
+
+def create_prompt_opts(
+        var_dict, create_negative_prompt=True, prompt='', negatove_prompt='', display_label=False):
+    if display_label:
+        kwargs = dict(show_label=True, container=True, elem_classes=['force-hide-container'])
+    else:
+        kwargs = dict(show_label=False, container=False)
+    var_dict['prompt'] = gr.Textbox(
+        prompt, label='Prompt', lines=2, placeholder='Prompt', interactive=True, **kwargs)
+    if create_negative_prompt:
+        var_dict['negative_prompt'] = gr.Textbox(
+            negatove_prompt, label='Negative prompt', lines=2,
+            placeholder='Negative prompt', interactive=True, **kwargs)
+
+
+def create_generate_bar(var_dict, text='Generate', variant='primary', seed=-1):
+    with gr.Row(equal_height=False, elem_classes=['generate-bar']):
+        var_dict['run_btn'] = gr.Button(text, variant=variant, scale=2)
+        var_dict['seed'] = gr.Number(
+            label='Seed', value=seed, min_width=100, precision=0, minimum=-1, maximum=2 ** 31,
+            elem_classes=['force-hide-container', 'seed-input'])
+        var_dict['random_seed'] = gr.Button('\U0001f3b2\ufe0f', elem_classes=['tool'])
+        var_dict['reuse_seed'] = gr.Button('\u267b\ufe0f', elem_classes=['tool'])
+        with gr.Column(visible=False):
+            var_dict['last_seed'] = gr.Number(value=seed, label='Last seed')
+    var_dict['reuse_seed'].click(
+        fn=lambda x: x,
+        inputs=var_dict['last_seed'],
+        outputs=var_dict['seed'],
+        show_progress=False,
+        api_name=False)
+    var_dict['random_seed'].click(
+        fn=lambda: -1,
+        outputs=var_dict['seed'],
+        show_progress=False,
+        api_name=False)
+
+
+def create_image_size_bar(var_dict, height=768, width=1360, hw_slider_step=16):
+    with gr.Row(equal_height=True, variant='compact', elem_classes=['force-hide-container']):
+        var_dict['width'] = gr.Slider(
+            label='Width', minimum=64, maximum=2048, step=hw_slider_step, value=width,
+            elem_classes=['force-hide-container'])
+        var_dict['switch_hw'] = gr.Button('\U000021C6', elem_classes=['tool'])
+        var_dict['height'] = gr.Slider(
+            label='Height', minimum=64, maximum=2048, step=hw_slider_step, value=height,
+            elem_classes=['force-hide-container'])
+        var_dict['switch_hw'].click(
+            fn=lambda w, h: (h, w),
+            inputs=[var_dict['width'], var_dict['height']],
+            outputs=[var_dict['width'], var_dict['height']],
+            show_progress=False,
+            api_name=False)
+
+
+def create_base_opts(var_dict,
+                     steps=24,
+                     min_steps=4,
+                     max_steps=50,
+                     steps_slider_step=1,
+                     guidance_scale=None,
+                     temperature=None,
+                     render=True):
+    with gr.Column(variant='compact', elem_classes=['custom-spacing'], render=render) as base_opts:
+        with gr.Row(variant='compact', elem_classes=['force-hide-container']):
+            var_dict['steps'] = gr.Slider(
+                min_steps, max_steps, value=steps, step=steps_slider_step, label='Sampling steps',
+                elem_classes=['force-hide-container'])
+        if guidance_scale is not None or temperature is not None:
+            with gr.Row(variant='compact', elem_classes=['force-hide-container']):
+                if guidance_scale is not None:
+                    var_dict['guidance_scale'] = gr.Slider(
+                        0.0, 30.0, value=guidance_scale, step=0.5, label='Guidance scale',
+                        elem_classes=['force-hide-container'])
+                if temperature is not None:
+                    var_dict['temperature'] = gr.Slider(
+                        0.0, 1.0, value=temperature, step=0.01, label='Temperature',
+                        elem_classes=['force-hide-container'])
+    return base_opts
+
+
+def set_seed(seed):
+    seed = random.randint(0, 2**31) if seed == -1 else seed
+    return seed

lakonlab/ui/gradio/style.css

@@ -0,0 +1,73 @@
+.force-hide-container {
+    margin: 0;
+    box-shadow: none;
+    --block-border-width: 0;
+    background: transparent;
+    padding: 0;
+    overflow: visible;
+}
+
+.svelte-1vd8eap {
+    display: flex;
+    flex-direction: inherit;
+    flex-wrap: wrap;
+    gap: 0;
+    box-shadow: none;
+    border: 0;
+    border-radius: 0;
+    background: transparent;
+    overflow-y: hidden;
+}
+
+.custom-spacing {
+    padding: 10px;
+    gap: 20px;
+    flex-grow: 0 !important;
+}
+
+.generate-bar {
+    align-items: flex-end;
+}
+
+.tool{
+    max-width: 40px;
+    min-width: 40px !important;
+}
+
+/* Center the component and allow it to use the full row width */
+.vh-img {
+  display: grid;
+  justify-items: center;
+}
+
+/* Container should size to the image, but never exceed the row width */
+.vh-img .image-container {
+  inline-size: fit-content !important;   /* prefers image’s natural width */
+  max-inline-size: 100% !important;      /* ...but clamps to available width */
+  margin-inline: auto;
+  overflow: hidden;                       /* avoid odd overflow on iOS */
+}
+
+/* Image scales by BOTH constraints: height cap and row width */
+.vh-img-700 .image-container img {
+  max-block-size: 700px !important;      /* fixed max height cap */
+  max-inline-size: 100%;                  /* never wider than container */
+  inline-size: auto;                      /* keep aspect ratio */
+  block-size: auto;
+  object-fit: contain;
+  display: block;
+}
+
+.vh-img-1000 .image-container img {
+  max-block-size: 1000px !important;      /* fixed max height cap */
+  max-inline-size: 100%;                  /* never wider than container */
+  inline-size: auto;                      /* keep aspect ratio */
+  block-size: auto;
+  object-fit: contain;
+  display: block;
+}
+
+.gradio-container .seed-input input {   /* remove the border and use box-shadow instead for height alignment */
+    border: none !important;
+    box-shadow: inset 0 0 0 var(--input-border-width) var(--input-border-color) !important;
+}

requirements.txt

@@ -0,0 +1,10 @@
+numpy==1.26.4
+torch==2.6.0
+torchvision==0.21.0
+diffusers==0.36.0
+peft==0.17.0
+sentencepiece
+accelerate
+transformers==4.57.3
+gradio==5.49.0
+bitsandbytes>=0.46.1