Add gradio app

.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

README.md

@@ -1,13 +1,25 @@
 ---
-title: Pi Qwen
-emoji: 🏃
-colorFrom: red
-colorTo: blue
+title: Pi-Qwen Demo
+emoji: 🚀
+colorFrom: yellow
+colorTo: pink
 sdk: gradio
-sdk_version: 5.49.1
+sdk_version: 4.18.0
 app_file: app.py
 pinned: false
 license: apache-2.0
 ---
 
-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]()] [[pi-Qwen Demo🤗](https://huggingface.co/spaces/Lakonik/pi-Qwen)] [[pi-FLUX Demo🤗](https://huggingface.co/spaces/Lakonik/pi-FLUX.1)] 

app.py

@@ -0,0 +1,60 @@
+import torch
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+import os
+import gradio as gr
+import spaces
+from diffusers import FlowMatchEulerDiscreteScheduler
+from lakonlab.ui.gradio.create_text_to_img import create_interface_text_to_img
+from lakonlab.pipelines.piqwen_pipeline import PiQwenImagePipeline
+
+from huggingface_hub import login
+login(token=os.getenv('HF_TOKEN'))
+
+
+DEFAULT_PROMPT = ('Photo of a coffee shop entrance featuring a chalkboard sign reading "π-Qwen Coffee 😊 $2 per cup," '
+                  'with a neon light beside it displaying "π-通义千问". Next to it hangs a poster showing a beautiful '
+                  'Chinese woman, and beneath the poster is written "e≈2.71828-18284-59045-23536-02874-71352".')
+
+
+pipe = PiQwenImagePipeline.from_pretrained(
+    'Qwen/Qwen-Image',
+    torch_dtype=torch.bfloat16)
+pipe.load_piflow_adapter(
+    'Lakonik/pi-Qwen-Image',
+    subfolder='gmqwen_k8_piid_4step',
+    target_module_name='transformer')
+pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(  # use fixed shift=3.2
+    pipe.scheduler.config, shift=3.2, shift_terminal=None, use_dynamic_shifting=False)
+pipe = pipe.to('cuda')
+
+
+@spaces.GPU
+def generate(seed, prompt, width, height, steps):
+    return pipe(
+        prompt=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-Qwen Demo',
+               css='lakonlab/ui/gradio/style.css'
+               ) as demo:
+
+    md_txt = '# pi-Qwen Demo\n\n' \
+             'Official demo of the paper [pi-Flow: Policy-Based Few-Step Generation via Imitation Distillation](). ' \
+             '**Base model:** [Qwen-Image](https://huggingface.co/Qwen/Qwen-Image). **Fast policy:** GMFlow. **Code:** [https://github.com/Lakonik/piFlow](https://github.com/Lakonik/piFlow).'
+    gr.Markdown(md_txt)
+
+    create_interface_text_to_img(
+        generate,
+        prompt=DEFAULT_PROMPT,
+        steps=4, guidance_scale=None,
+        args=['last_seed', 'prompt', 'width', 'height', 'steps'])
+    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/gmflux.py

@@ -0,0 +1,225 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+from typing import Any, Dict, Optional, Tuple
+from diffusers.models import ModelMixin
+from diffusers.models.transformers.transformer_flux import (
+    FluxTransformer2DModel, FluxPosEmbed, FluxTransformerBlock, FluxSingleTransformerBlock)
+from diffusers.models.embeddings import (
+    CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings)
+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
+from .gm_output import GMFlowModelOutput
+
+
+class _GMFluxTransformer2DModel(FluxTransformer2DModel):
+
+    @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 = 64,
+            out_channels: Optional[int] = None,
+            num_layers: int = 19,
+            num_single_layers: int = 38,
+            attention_head_dim: int = 128,
+            num_attention_heads: int = 24,
+            joint_attention_dim: int = 4096,
+            pooled_projection_dim: int = 768,
+            guidance_embeds: bool = False,
+            axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)):
+        super(FluxTransformer2DModel, 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
+
+        self.pos_embed = FluxPosEmbed(theta=10000, axes_dim=axes_dims_rope)
+
+        text_time_guidance_cls = (
+            CombinedTimestepGuidanceTextProjEmbeddings if guidance_embeds else CombinedTimestepTextProjEmbeddings
+        )
+        self.time_text_embed = text_time_guidance_cls(
+            embedding_dim=self.inner_dim, pooled_projection_dim=pooled_projection_dim
+        )
+
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim)
+        self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        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,
+            pooled_projections: torch.Tensor = None,
+            timestep: torch.Tensor = None,
+            img_ids: torch.Tensor = None,
+            txt_ids: torch.Tensor = None,
+            guidance: torch.Tensor = None,
+            joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+            controlnet_block_samples=None,
+            controlnet_single_block_samples=None,
+            controlnet_blocks_repeat: bool = False):
+        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
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype) * 1000
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+
+        temb = (
+            self.time_text_embed(timestep, pooled_projections)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, pooled_projections)
+        )
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        ids = torch.cat((txt_ids, img_ids), dim=0)
+        image_rotary_emb = self.pos_embed(ids)
+        image_rotary_emb = tuple([x.to(hidden_states.dtype) for x in image_rotary_emb])
+
+        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
+            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
+            ip_hidden_states = self.encoder_hid_proj(ip_adapter_image_embeds)
+            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})
+
+        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,
+                    temb,
+                    image_rotary_emb,
+                    joint_attention_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                # For Xlabs ControlNet.
+                if controlnet_blocks_repeat:
+                    hidden_states = (
+                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
+                    )
+                else:
+                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+        for index_block, block in enumerate(self.single_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,
+                    temb,
+                    image_rotary_emb,
+                    joint_attention_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        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/architecture/gmflow/gmqwen.py

@@ -0,0 +1,149 @@
+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_qwenimage import (
+    QwenImageTransformer2DModel, QwenEmbedRope, QwenImageTransformerBlock, QwenTimestepProjEmbeddings)
+from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle, RMSNorm
+from diffusers.configuration_utils import register_to_config
+from diffusers.utils import USE_PEFT_BACKEND, scale_lora_layers, unscale_lora_layers
+from .gm_output import GMFlowModelOutput
+
+
+class _GMQwenImageTransformer2DModel(QwenImageTransformer2DModel):
+
+    @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 = 64,
+            out_channels: Optional[int] = None,
+            num_layers: int = 60,
+            attention_head_dim: int = 128,
+            num_attention_heads: int = 24,
+            joint_attention_dim: int = 3584,
+            axes_dims_rope: Tuple[int, int, int] = (16, 56, 56)):
+        super(QwenImageTransformer2DModel, 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
+
+        self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
+
+        self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+
+        self.img_in = nn.Linear(in_channels, self.inner_dim)
+        self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                QwenImageTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        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,
+            encoder_hidden_states_mask: torch.Tensor = None,
+            timestep: torch.LongTensor = None,
+            img_shapes: Optional[List[Tuple[int, int, int]]] = None,
+            txt_seq_lens: Optional[List[int]] = None,
+            attention_kwargs: Optional[Dict[str, Any]] = None):
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            scale_lora_layers(self, lora_scale)
+        else:
+            assert attention_kwargs is None or attention_kwargs.get('scale', None) is None
+
+        hidden_states = self.img_in(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype)
+        encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        temb = self.time_text_embed(timestep, hidden_states)
+
+        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+
+        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,
+                    encoder_hidden_states_mask,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_states_mask=encoder_hidden_states_mask,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=attention_kwargs,
+                )
+
+        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 u(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,108 @@
+# 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.
+        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,
+            eps: float = 1e-4):
+        self.x_t_src = x_t_src
+        self.ndim = x_t_src.dim()
+        self.shift = shift
+        self.eps = eps
+
+        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 u(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)
+        x_0 = self._interpolate(
+            self.denoising_output_x_0, (raw_t - self.raw_t_dst) / self.segment_size)
+        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.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,175 @@
+# Copyright (c) 2025 Hansheng Chen
+
+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):
+    alpha_t_src = 1 - sigma_t_src
+    alpha_t = 1 - sigma_t
+
+    sigma_t_src_sq = sigma_t_src.square()
+    sigma_t_sq = sigma_t.square()
+
+    # compute gaussian params
+    denom = (alpha_t.square() * sigma_t_src_sq - alpha_t_src.square() * sigma_t_sq).clamp(min=eps)  # ζ
+    g_mean = (alpha_t * sigma_t_src_sq * x_t - alpha_t_src * sigma_t_sq * x_t_src) / denom  # ν / ζ
+    g_var = sigma_t_sq * sigma_t_src_sq / denom
+
+    # gm_mul_iso_gaussian
+    g_mean = g_mean.unsqueeze(gm_dim)  # (bs, *, 1, out_channels, h, w)
+    g_var = g_var.unsqueeze(gm_dim)  # (bs, *, 1, 1, 1, 1)
+
+    gm_diffs = gm_means - g_mean  # (bs, *, num_gaussians, out_channels, h, w)
+    norm_factor = (g_var + gm_vars).clamp(min=eps)
+
+    out_means = (g_var * gm_means + gm_vars * g_mean) / norm_factor
+    # (bs, *, num_gaussians, 1, h, w)
+    logweights_delta = gm_diffs.square().sum(dim=channel_dim, keepdim=True) * (-0.5 / norm_factor)
+    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-4):
+        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 u(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/pipelines/piflow_loader.py

@@ -0,0 +1,275 @@
+# 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 lakonlab.models.architecture.gmflow.gmflux import _GMFluxTransformer2DModel
+from lakonlab.models.architecture.gmflow.gmqwen import _GMQwenImageTransformer2DModel
+
+
+LOCAL_CLASS_MAPPING = {
+    "GMFluxTransformer2DModel": _GMFluxTransformer2DModel,
+    "GMQwenImageTransformer2DModel": _GMQwenImageTransformer2DModel,
+}
+
+_SET_ADAPTER_SCALE_FN_MAPPING.update(
+    _GMFluxTransformer2DModel=lambda model_cls, weights: weights,
+    _GMQwenImageTransformer2DModel=lambda model_cls, weights: weights,
+)
+
+logger = logging.get_logger(__name__)
+
+
+class PiFlowLoaderMixin:
+
+    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,
+            )
+
+        # 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)
+
+        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)
+
+        # 4. Load model weights
+
+        if model_file is not None:
+            base_state_dict = base_module.state_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:
+                    base_state_dict[k.removeprefix(f"{target_module_name}.")] = adapter_state_dict[k]
+
+            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_state_dict(
+                    base_state_dict, strict=False, assign=True)
+                piflow_module.load_lora_adapter(
+                    lora_state_dict, prefix=None, adapter_name=adapter_name)
+
+                setattr(self, target_module_name, piflow_module)
+
+        else:
+            adapter_name = None
+
+        if adapter_name is None:
+            logger.warning(
+                f"No LoRA weights were found in {pretrained_model_name_or_path}."
+            )
+
+        return adapter_name

lakonlab/pipelines/piflux_pipeline.py

@@ -0,0 +1,491 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import numpy as np
+import torch
+
+from typing import Dict, List, Optional, Union, Any, Callable
+from functools import partial
+from transformers import (
+    CLIPImageProcessor,
+    CLIPTextModel,
+    CLIPTokenizer,
+    CLIPVisionModelWithProjection,
+    T5EncoderModel,
+    T5TokenizerFast,
+)
+from diffusers.utils import is_torch_xla_available
+from diffusers.image_processor import PipelineImageInput
+from diffusers.models import AutoencoderKL, FluxTransformer2DModel
+from diffusers.pipelines.flux.pipeline_flux import (
+    FluxPipeline, calculate_shift, FluxPipelineOutput, retrieve_timesteps)
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from lakonlab.models.diffusions.piflow_policies import POLICY_CLASSES
+from .piflow_loader import PiFlowLoaderMixin
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+def retrieve_raw_timesteps(
+    num_inference_steps: int,
+    total_substeps: int,
+    final_step_size_scale: float
+):
+    r"""
+    Retrieve the raw times and the number of substeps for each inference step.
+
+    Args:
+        num_inference_steps (`int`):
+            Number of inference steps.
+        total_substeps (`int`):
+            Total number of substeps (e.g., 128).
+        final_step_size_scale (`float`):
+            Scale for the final step size (e.g., 0.5).
+
+    Returns:
+        `Tuple[List[float], List[int], int]`: A tuple where the first element is the raw timestep schedule, the second
+        element is the number of substeps for each inference step, and the third element is the rounded total number of
+        substeps.
+    """
+    base_segment_size = 1 / (num_inference_steps - 1 + final_step_size_scale)
+    raw_timesteps = []
+    num_inference_substeps = []
+    _raw_t = 1.0
+    for i in range(num_inference_steps):
+        if i < num_inference_steps - 1:
+            segment_size = base_segment_size
+        else:
+            segment_size = base_segment_size * final_step_size_scale
+        _num_inference_substeps = max(round(segment_size * total_substeps), 1)
+        num_inference_substeps.append(_num_inference_substeps)
+        raw_timesteps.extend(np.linspace(
+            _raw_t, _raw_t - segment_size, _num_inference_substeps, endpoint=False).clip(min=0.0).tolist())
+        _raw_t = _raw_t - segment_size
+    total_substeps = sum(num_inference_substeps)
+    return raw_timesteps, num_inference_substeps, total_substeps
+
+
+class PiFluxPipeline(FluxPipeline, PiFlowLoaderMixin):
+    r"""
+    The policy-based Flux pipeline for text-to-image generation.
+
+    Reference: Todo: add paper link
+
+    Args:
+        transformer ([`FluxTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`CLIPTextModel`]):
+            [CLIP](https://huggingface.co/docs/transformers/model_doc/clip#transformers.CLIPTextModel), specifically
+            the [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14) variant.
+        text_encoder_2 ([`T5EncoderModel`]):
+            [T5](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5EncoderModel), specifically
+            the [google/t5-v1_1-xxl](https://huggingface.co/google/t5-v1_1-xxl) variant.
+        tokenizer (`CLIPTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        tokenizer_2 (`T5TokenizerFast`):
+            Second Tokenizer of class
+            [T5TokenizerFast](https://huggingface.co/docs/transformers/en/model_doc/t5#transformers.T5TokenizerFast).
+        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.
+    """
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        tokenizer: CLIPTokenizer,
+        text_encoder_2: T5EncoderModel,
+        tokenizer_2: T5TokenizerFast,
+        transformer: FluxTransformer2DModel,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+        policy_type: str = 'GMFlow',
+        policy_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__(
+            scheduler,
+            vae,
+            text_encoder,
+            tokenizer,
+            text_encoder_2,
+            tokenizer_2,
+            transformer,
+            image_encoder,
+            feature_extractor
+        )
+        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
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width, patch_size=1, target_patch_size=2):
+        scale = target_patch_size // patch_size
+        latents = latents.view(
+            batch_size,
+            num_channels_latents * patch_size * patch_size,
+            height // target_patch_size, scale, width // target_patch_size, scale)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(
+            batch_size,
+            (height // target_patch_size) * (width // target_patch_size),
+            num_channels_latents * target_patch_size * target_patch_size)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor, patch_size=2, target_patch_size=1):
+        batch_size, num_patches, channels = latents.shape
+        scale = patch_size // target_patch_size
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = (int(height) // (vae_scale_factor * patch_size))
+        width = (int(width) // (vae_scale_factor * patch_size))
+
+        latents = latents.view(
+            batch_size, height, width, channels // (scale * scale), scale, scale)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (scale * scale), height * scale, width * scale)
+
+        return latents
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 4,
+        total_substeps: int = 128,
+        final_step_size_scale: float = 0.5,
+        temperature: Union[float, str] = 'auto',
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_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,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            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.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead.
+            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.
+            total_substeps (`int`, *optional*, defaults to 128):
+                The total number of substeps for policy-based flow integration.
+            final_step_size_scale (`float`, *optional*, defaults to 0.5):
+                The scale for the final step size.
+            temperature (`float` or `"auto"`, *optional*, defaults to `"auto"`):
+                The tmperature parameter for the flow policy.
+            guidance_scale (`float`, *optional*, defaults to 3.5):
+                Embedded guiddance scale is enabled by setting `guidance_scale` > 1. Higher `guidance_scale` encourages
+                a model to generate images more aligned with `prompt` at the expense of lower image quality.
+
+                Guidance-distilled models approximates true classifer-free guidance for `guidance_scale` > 1. Refer to
+                the [paper](https://huggingface.co/papers/2210.03142) to learn more.
+            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.FloatTensor`, *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.FloatTensor`, *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.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` 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.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_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`.
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_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 prompt embeddings
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        raw_timesteps, num_inference_substeps, total_substeps = retrieve_raw_timesteps(
+            num_inference_steps, total_substeps, final_step_size_scale)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, _ = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=raw_timesteps,
+            mu=mu,
+        )
+        assert len(timesteps) == total_substeps
+        self._num_timesteps = total_substeps
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        image_embeds = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+
+        # 6. Denoising loop
+        self.scheduler.set_begin_index(0)
+        timestep_id = 0
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i in range(num_inference_steps):
+                if self.interrupt:
+                    continue
+
+                t_src = timesteps[timestep_id]
+                sigma_t_src = t_src / self.scheduler.config.num_train_timesteps
+                is_final_step = i == (num_inference_steps - 1)
+
+                self._current_timestep = t_src
+                if image_embeds is not None:
+                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+
+                with self.transformer.cache_context("cond"):
+                    denoising_output = self.transformer(
+                        hidden_states=latents.to(dtype=self.transformer.dtype),
+                        timestep=t_src.expand(latents.shape[0]) / 1000,
+                        guidance=guidance,
+                        pooled_projections=pooled_prompt_embeds,
+                        encoder_hidden_states=prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                    )
+
+                # unpack and create policy
+                latents = self._unpack_latents(
+                    latents, height, width, self.vae_scale_factor, target_patch_size=1)
+                if self.policy_type == 'GMFlow':
+                    denoising_output = self._unpack_gm(
+                        denoising_output, height, width, num_channels_latents, gm_patch_size=1)
+                    denoising_output = {k: v.to(torch.float32) for k, v in denoising_output.items()}
+                    policy = self.policy_class(
+                        denoising_output, latents, sigma_t_src)
+                    if not is_final_step:
+                        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(
+                        denoising_output, height, width, self.vae_scale_factor, target_patch_size=1)
+                    denoising_output = denoising_output.reshape(latents.size(0), -1, *latents.shape[1:])
+                    denoising_output = denoising_output.to(torch.float32)
+                    policy = self.policy_class(
+                        denoising_output, latents, sigma_t_src)
+                else:
+                    raise ValueError(f'Unknown policy type: {self.policy_type}.')
+
+                # compute the previous noisy sample x_t -> x_t-1
+                for _ in range(num_inference_substeps[i]):
+                    t = timesteps[timestep_id]
+                    sigma_t = t / self.scheduler.config.num_train_timesteps
+                    u = policy.u(latents, sigma_t)
+                    latents = self.scheduler.step(u, t, latents, return_dict=False)[0]
+                    timestep_id += 1
+
+                # repack
+                latents = self._pack_latents(
+                    latents, latents.size(0), num_channels_latents,
+                    2 * (int(height) // (self.vae_scale_factor * 2)),
+                    2 * (int(width) // (self.vae_scale_factor * 2)),
+                    patch_size=1)
+
+                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:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            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 FluxPipelineOutput(images=image)

lakonlab/pipelines/piqwen_pipeline.py

@@ -0,0 +1,429 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import numpy as np
+import torch
+
+from typing import Dict, List, Optional, Union, Any, Callable
+from functools import partial
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer
+from diffusers.utils import is_torch_xla_available
+from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
+from diffusers.pipelines.qwenimage.pipeline_qwenimage import (
+    QwenImagePipeline, calculate_shift, retrieve_timesteps, QwenImagePipelineOutput)
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from lakonlab.models.diffusions.piflow_policies import POLICY_CLASSES
+from .piflow_loader import PiFlowLoaderMixin
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+def retrieve_raw_timesteps(
+    num_inference_steps: int,
+    total_substeps: int,
+    final_step_size_scale: float
+):
+    r"""
+    Retrieve the raw times and the number of substeps for each inference step.
+
+    Args:
+        num_inference_steps (`int`):
+            Number of inference steps.
+        total_substeps (`int`):
+            Total number of substeps (e.g., 128).
+        final_step_size_scale (`float`):
+            Scale for the final step size (e.g., 0.5).
+
+    Returns:
+        `Tuple[List[float], List[int], int]`: A tuple where the first element is the raw timestep schedule, the second
+        element is the number of substeps for each inference step, and the third element is the rounded total number of
+        substeps.
+    """
+    base_segment_size = 1 / (num_inference_steps - 1 + final_step_size_scale)
+    raw_timesteps = []
+    num_inference_substeps = []
+    _raw_t = 1.0
+    for i in range(num_inference_steps):
+        if i < num_inference_steps - 1:
+            segment_size = base_segment_size
+        else:
+            segment_size = base_segment_size * final_step_size_scale
+        _num_inference_substeps = max(round(segment_size * total_substeps), 1)
+        num_inference_substeps.append(_num_inference_substeps)
+        raw_timesteps.extend(np.linspace(
+            _raw_t, _raw_t - segment_size, _num_inference_substeps, endpoint=False).clip(min=0.0).tolist())
+        _raw_t = _raw_t - segment_size
+    total_substeps = sum(num_inference_substeps)
+    return raw_timesteps, num_inference_substeps, total_substeps
+
+
+class PiQwenImagePipeline(QwenImagePipeline, PiFlowLoaderMixin):
+    r"""
+    The policy-based QwenImage pipeline for text-to-image generation.
+
+    Reference: Todo: add paper link
+
+    Args:
+        transformer ([`QwenImageTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`Qwen2.5-VL-7B-Instruct`]):
+            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the
+            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant.
+        tokenizer (`QwenTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+        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.
+    """
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLQwenImage,
+        text_encoder: Qwen2_5_VLForConditionalGeneration,
+        tokenizer: Qwen2Tokenizer,
+        transformer: QwenImageTransformer2DModel,
+        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
+
+    @staticmethod
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width, patch_size=1, target_patch_size=2):
+        scale = target_patch_size // patch_size
+        latents = latents.view(
+            batch_size,
+            num_channels_latents * patch_size * patch_size,
+            height // target_patch_size, scale, width // target_patch_size, scale)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(
+            batch_size,
+            (height // target_patch_size) * (width // target_patch_size),
+            num_channels_latents * target_patch_size * target_patch_size)
+
+        return latents
+
+    @staticmethod
+    def _unpack_latents(latents, height, width, vae_scale_factor, patch_size=2, target_patch_size=1):
+        batch_size, num_patches, channels = latents.shape
+        scale = patch_size // target_patch_size
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = (int(height) // (vae_scale_factor * patch_size))
+        width = (int(width) // (vae_scale_factor * patch_size))
+
+        latents = latents.view(
+            batch_size, height, width, channels // (scale * scale), scale, scale)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (scale * scale), height * scale, width * scale)
+
+        return latents
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 4,
+        total_substeps: int = 128,
+        final_step_size_scale: float = 0.5,
+        temperature: Union[float, str] = 'auto',
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_embeds_mask: 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,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            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.
+            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.
+            total_substeps (`int`, *optional*, defaults to 128):
+                The total number of substeps for policy-based flow integration.
+            final_step_size_scale (`float`, *optional*, defaults to 0.5):
+                The scale for the final step size.
+            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.FloatTensor`, *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.FloatTensor`, *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.flux.FluxPipelineOutput`] 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`.
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        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 prompt embeddings
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+        img_shapes = [[(1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2)]] * batch_size
+
+        # 5. Prepare timesteps
+        raw_timesteps, num_inference_substeps, total_substeps = retrieve_raw_timesteps(
+            num_inference_steps, total_substeps, final_step_size_scale)
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, _ = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=raw_timesteps,
+            mu=mu,
+        )
+        assert len(timesteps) == total_substeps
+        self._num_timesteps = total_substeps
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+
+        # 6. Denoising loop
+        self.scheduler.set_begin_index(0)
+        timestep_id = 0
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i in range(num_inference_steps):
+                if self.interrupt:
+                    continue
+
+                t_src = timesteps[timestep_id]
+                sigma_t_src = t_src / self.scheduler.config.num_train_timesteps
+                is_final_step = i == (num_inference_steps - 1)
+
+                self._current_timestep = t_src
+
+                with self.transformer.cache_context("cond"):
+                    denoising_output = self.transformer(
+                        hidden_states=latents.to(dtype=self.transformer.dtype),
+                        timestep=t_src.expand(latents.shape[0]) / 1000,
+                        encoder_hidden_states_mask=prompt_embeds_mask,
+                        encoder_hidden_states=prompt_embeds,
+                        img_shapes=img_shapes,
+                        txt_seq_lens=txt_seq_lens,
+                        attention_kwargs=self.attention_kwargs,
+                    )
+
+                # unpack and create policy
+                latents = self._unpack_latents(
+                    latents, height, width, self.vae_scale_factor, target_patch_size=1)
+                if self.policy_type == 'GMFlow':
+                    denoising_output = self._unpack_gm(
+                        denoising_output, height, width, num_channels_latents, gm_patch_size=1)
+                    denoising_output = {k: v.to(torch.float32) for k, v in denoising_output.items()}
+                    policy = self.policy_class(
+                        denoising_output, latents, sigma_t_src)
+                    if not is_final_step:
+                        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(
+                        denoising_output, height, width, self.vae_scale_factor, target_patch_size=1)
+                    denoising_output = denoising_output.reshape(latents.size(0), -1, *latents.shape[1:])
+                    denoising_output = denoising_output.to(torch.float32)
+                    policy = self.policy_class(
+                        denoising_output, latents, sigma_t_src)
+                else:
+                    raise ValueError(f'Unknown policy type: {self.policy_type}.')
+
+                # compute the previous noisy sample x_t -> x_t-1
+                for _ in range(num_inference_substeps[i]):
+                    t = timesteps[timestep_id]
+                    sigma_t = t / self.scheduler.config.num_train_timesteps
+                    u = policy.u(latents, sigma_t)
+                    latents = self.scheduler.step(u, t, latents, return_dict=False)[0]
+                    timestep_id += 1
+
+                # repack
+                latents = self._pack_latents(
+                    latents, latents.size(0), num_channels_latents,
+                    2 * (int(height) // (self.vae_scale_factor * 2)),
+                    2 * (int(width) // (self.vae_scale_factor * 2)),
+                    patch_size=1)
+
+                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:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)[:, :, None]
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents * latents_std + latents_mean
+            image = self.vae.decode(latents.to(self.vae.dtype), return_dict=False)[0][:, :, 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 QwenImagePipelineOutput(images=image)

lakonlab/ui/gradio/create_text_to_img.py

@@ -0,0 +1,53 @@
+import gradio as gr
+from .shared_opts import create_base_opts, create_generate_bar, set_seed, create_prompt_opts
+
+
+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'])
+        create_prompt_opts(var_dict, create_negative_prompt=create_negative_prompt, prompt=prompt)
+        with gr.Column(variant='compact', elem_classes=['custom-spacing']):
+            with gr.Row(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)
+        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,64 @@
+import random
+import gradio as gr
+
+
+def create_prompt_opts(var_dict, create_negative_prompt=True, prompt='', negatove_prompt=''):
+    var_dict['prompt'] = gr.Textbox(
+        prompt, label='Prompt', show_label=False, lines=2, placeholder='Prompt', container=False, interactive=True)
+    if create_negative_prompt:
+        var_dict['negative_prompt'] = gr.Textbox(
+            negatove_prompt, label='Negative prompt', show_label=False, lines=2,
+            placeholder='Negative prompt', container=False, interactive=True)
+
+
+def create_generate_bar(var_dict, text='Generate', variant='primary', seed=-1):
+    with gr.Row(equal_height=False):
+        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'])
+        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_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'])
+        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,59 @@
+.force-hide-container {
+    margin: 0;
+    box-shadow: none;
+    --block-border-width: 0;
+    background: transparent;
+    padding: 0;
+    overflow: visible;
+}
+
+.svelte-sfqy0y {
+    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;
+}
+
+.unequal-height {
+    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 .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;
+}

requirements.txt

@@ -0,0 +1,8 @@
+numpy==1.26.4
+torch==2.6.0
+diffusers==0.35.1
+peft==0.17.0
+sentencepiece
+accelerate
+transformers==4.54.1
+gradio==4.18.0