Add AsymFLUX.2-klein Space demo

.gitignore

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

LICENSE.md

@@ -0,0 +1,11 @@
+# License for AsymFLUX.2-klein
+
+This repository distributes an **AsymFLUX.2-klein app** that is a **Derivative** of
+**FLUX.2 klein Base 9B** by **Black Forest Labs Inc.**
+
+Use and distribution of this app are governed by the **FLUX Non-Commercial License**.
+No commercial use of this app or its derivatives is permitted without a separate
+commercial license from Black Forest Labs.
+
+- Full license: https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B/blob/main/LICENSE.md
+- This repository does not grant any rights beyond the license above.

README.md

@@ -1,14 +1,26 @@
 ---
-title: AsymFLUX.2 Klein
-emoji: 📈
-colorFrom: green
-colorTo: yellow
+title: AsymFLUX.2-klein Demo
+emoji: 🚀
+colorFrom: blue
+colorTo: purple
 sdk: gradio
-sdk_version: 6.14.0
-python_version: '3.12'
+sdk_version: 5.49.0
 app_file: app.py
 pinned: false
-short_description: High-fidelity pixel-space text-to-image generation
+license: other
+license_name: flux-non-commercial-license
+license_link: LICENSE.md
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Official demo of the paper:
+
+**Asymmetric Flow Models**
+<br>
+[Hansheng Chen](https://lakonik.github.io/),
+Jan Ackermann,
+Minseo Kim,
+[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/),
+[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)<br>
+Stanford University
+<br>
+[[Paper](https://hanshengchen.com/asymflow/static/paper.pdf)] [[Project](https://hanshengchen.com/asymflow)] [[Code](https://github.com/Lakonik/piFlow)]

app.py

@@ -0,0 +1,139 @@
+import math
+import os
+import random
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+
+from lakonlab.models.architectures import OklabColorEncoder
+from lakonlab.models.diffusions.schedulers import FlowAdapterScheduler
+from lakonlab.pipelines.pipeline_pixelflux2_klein import PixelFlux2KleinPipeline
+from lakonlab.pipelines.prompt_rewriters.qwen3_vl import Qwen3VLPromptRewriter
+from lakonlab.ui.gradio.create_text_to_img import create_interface_text_to_img
+
+
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+
+
+DEFAULT_PROMPT = (
+    'Restored color photo from the 1900s. A middle-aged man with cybernetic metal hands is sitting on an old wooden '
+    'chair and reading the newspaper. The newspaper has the prominent headline "AsymFLOW RELEASED" in large bold font. '
+    'Close-up shot focusing on the newspaper.'
+)
+DEFAULT_NEG_PROMPT = 'Low quality, worst quality, blurry, deformed, bad anatomy, unclear text'
+
+SYSTEM_PROMPT_TEXT_ONLY_PATH = 'lakonlab/pipelines/prompt_rewriters/system_prompts/default_text_only.txt'
+
+
+def set_random_seed(seed: int, deterministic: bool = True) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    if deterministic:
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+pipe = PixelFlux2KleinPipeline.from_pretrained(
+    'black-forest-labs/FLUX.2-klein-base-9B',
+    vae=OklabColorEncoder(
+        use_affine_norm=True,
+        mean=(0.56, 0.0, 0.01),
+        std=0.16),
+    scheduler=FlowAdapterScheduler(
+        shift=17.0,
+        use_dynamic_shifting=True,
+        base_seq_len=1024 ** 2,
+        max_seq_len=2048 ** 2,
+        base_logshift=math.log(17.0),
+        max_logshift=math.log(34.0),
+        dynamic_shifting_type='sqrt',
+        base_scheduler='UniPCMultistep'),
+    torch_dtype=torch.bfloat16)
+pipe.load_lakonlab_adapter(
+    'Lakonik/AsymFLUX.2-klein-9B',
+    target_module_name='transformer')
+pipe = pipe.to('cuda')
+
+prompt_rewriter = Qwen3VLPromptRewriter(
+    device_map='cuda',
+    system_prompt_text_only=open(SYSTEM_PROMPT_TEXT_ONLY_PATH, 'r').read(),
+    max_new_tokens_default=512,
+)
+
+
+@spaces.GPU
+def run_rewrite_prompt_gpu(seed, prompt, progress):
+    set_random_seed(seed)
+    progress(0.05, desc='Rewriting prompt...')
+    return prompt_rewriter.rewrite_text_batch([prompt])[0]
+
+
+def run_rewrite_prompt(seed, prompt, rewrite_prompt, progress=gr.Progress(track_tqdm=True)):
+    if rewrite_prompt:
+        return run_rewrite_prompt_gpu(seed, prompt, progress), None
+    return '', None
+
+
+@spaces.GPU
+def generate(
+        seed, prompt, negative_prompt, rewrite_prompt, rewritten_prompt, width, height, steps, guidance_scale,
+        progress=gr.Progress(track_tqdm=True)):
+    return pipe(
+        prompt=rewritten_prompt if rewrite_prompt else prompt,
+        negative_prompt=negative_prompt,
+        width=width,
+        height=height,
+        num_inference_steps=steps,
+        guidance_scale=guidance_scale,
+        generator=torch.Generator().manual_seed(seed),
+    ).images[0]
+
+
+with gr.Blocks(
+        analytics_enabled=False,
+        title='AsymFLUX.2-klein Demo',
+        css_paths='lakonlab/ui/gradio/style.css') as demo:
+    gr.Markdown(
+        '# AsymFLUX.2-klein Demo\n\n'
+        'Pixel-space text-to-image generation demo of the paper '
+        '[Asymmetric Flow Models](https://hanshengchen.com/asymflow/static/paper.pdf). '
+        '**Base model:** [FLUX.2 klein Base 9B](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B). '
+        '**Code:** [https://github.com/Lakonik/piFlow](https://github.com/Lakonik/piFlow).\n'
+        '<br> Use and distribution of this app are governed by the '
+        '[FLUX Non-Commercial License](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B/blob/main/LICENSE.md).'
+    )
+
+    create_interface_text_to_img(
+        generate,
+        prompt=DEFAULT_PROMPT,
+        negative_prompt=DEFAULT_NEG_PROMPT,
+        steps=38,
+        min_steps=4,
+        max_steps=50,
+        guidance_scale=4.0,
+        height=1280,
+        width=960,
+        create_negative_prompt=True,
+        create_prompt_rewrite=True,
+        args=[
+            'last_seed',
+            'prompt',
+            'negative_prompt',
+            'rewrite_prompt',
+            'rewritten_prompt',
+            'width',
+            'height',
+            'steps',
+            'guidance_scale',
+        ],
+        rewrite_prompt_api=run_rewrite_prompt,
+        rewrite_prompt_args=['last_seed', 'prompt', 'rewrite_prompt'])
+
+demo.queue().launch()

lakonlab/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.2.0'

lakonlab/models/__init__.py

@@ -0,0 +1 @@
+

lakonlab/models/architectures/__init__.py

@@ -0,0 +1,3 @@
+from .autoencoders import OklabColorEncoder, RGBColorEncoder
+
+__all__ = ['OklabColorEncoder', 'RGBColorEncoder']

lakonlab/models/architectures/asymflow/__init__.py

@@ -0,0 +1,3 @@
+from .asymflux2 import _AsymFlux2Transformer2DModel
+
+__all__ = ['_AsymFlux2Transformer2DModel']

lakonlab/models/architectures/asymflow/asymflux2.py

@@ -0,0 +1,309 @@
+from typing import Any, Optional, List
+
+import torch
+import torch.nn as nn
+
+from diffusers.models.transformers.transformer_flux2 import (
+    Flux2Transformer2DModel, Flux2PosEmbed, Flux2TransformerBlock, Flux2SingleTransformerBlock,
+    Flux2TimestepGuidanceEmbeddings, Flux2Modulation)
+from diffusers.models.normalization import AdaLayerNormContinuous
+from diffusers.configuration_utils import register_to_config
+from diffusers.utils import apply_lora_scale
+
+from .common import AsymFlowMixin
+
+
+class _AsymFlux2Transformer2DModel(AsymFlowMixin, Flux2Transformer2DModel):
+
+    @register_to_config
+    def __init__(
+            self,
+            patch_size=16,
+            in_channels: int = 3,
+            base_rank: int = 128,
+            num_layers: int = 8,
+            num_single_layers: int = 48,
+            attention_head_dim: int = 128,
+            num_attention_heads: int = 48,
+            joint_attention_dim: int = 15360,
+            timestep_guidance_channels: int = 256,
+            mlp_ratio: float = 3.0,
+            axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
+            rope_theta: int = 2000,
+            eps: float = 1e-6,
+            sigma_min: float = 1e-4,
+            num_timesteps=1,
+            guidance_embeds: bool = True):
+        super(Flux2Transformer2DModel, self).__init__()
+
+        self.patch_size = patch_size
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        # 1. Sinusoidal positional embedding for RoPE on image and text tokens
+        self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
+
+        # 2. Combined timestep + guidance embedding
+        self.time_guidance_embed = Flux2TimestepGuidanceEmbeddings(
+            in_channels=timestep_guidance_channels,
+            embedding_dim=self.inner_dim,
+            bias=False,
+            guidance_embeds=guidance_embeds,
+        )
+
+        # 3. Modulation (double stream and single stream blocks share modulation parameters, resp.)
+        # Two sets of shift/scale/gate modulation parameters for the double stream attn and FF sub-blocks
+        self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+        # Only one set of modulation parameters as the attn and FF sub-blocks are run in parallel for single stream
+        self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False)
+
+        # 4. Input projections
+        self.x_embedder = nn.Linear(in_channels * (patch_size ** 2), self.inner_dim, bias=False)
+        self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
+
+        # 5. Double Stream Transformer Blocks
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Flux2TransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 6. Single Stream Transformer Blocks
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                Flux2SingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                    mlp_ratio=mlp_ratio,
+                    eps=eps,
+                    bias=False,
+                )
+                for _ in range(num_single_layers)
+            ]
+        )
+
+        # 7. Output layers
+        self.norm_out = AdaLayerNormContinuous(
+            self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False
+        )
+        self.proj_out = nn.Linear(
+            self.inner_dim, self.out_channels * (patch_size ** 2), bias=False
+        )
+
+        # 8. AsymFlow attributes and buffers
+        self.base_rank = base_rank
+        self.sigma_min = sigma_min
+        self.num_timesteps = num_timesteps
+        self.init_asymflow_buffers(self.in_channels * (patch_size ** 2), self.base_rank)
+
+        self.gradient_checkpointing = False
+
+    @staticmethod
+    def patchify(latents, patch_size, pack_channels=True):
+        bs, c, h, w = latents.size()
+        latents = latents.reshape(
+            bs, c, h // patch_size, patch_size, w // patch_size, patch_size
+        ).permute(
+            0, 1, 3, 5, 2, 4
+        )
+        if pack_channels:
+            latents = latents.reshape(
+                bs, c * patch_size * patch_size, h // patch_size, w // patch_size)
+        else:
+            latents = latents.reshape(
+                bs, c, patch_size * patch_size, h // patch_size, w // patch_size)
+        return latents
+
+    @staticmethod
+    def unpatchify(latents, patch_size, packed_channels=True):
+        if packed_channels:
+            bs, c, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c // (patch_size * patch_size), patch_size, patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c // (patch_size * patch_size), h * patch_size, w * patch_size)
+        else:
+            bs, c, _, h, w = latents.size()
+            latents = latents.reshape(
+                bs, c, patch_size, patch_size, h, w
+            ).permute(
+                0, 1, 4, 2, 5, 3
+            ).reshape(
+                bs, c, h * patch_size, w * patch_size)
+        return latents
+
+    @staticmethod
+    def pack(latents):
+        bs, c, h, w = latents.shape
+        latents = latents.reshape(bs, c, h * w).permute(0, 2, 1)
+        return latents
+
+    @staticmethod
+    def unpack(latents, h, w):
+        bs, _, c = latents.shape
+        latents = latents.permute(0, 2, 1).reshape(bs, c, h, w)
+        return latents
+
+    @staticmethod
+    def _prepare_latent_ids(latents):
+        batch_size, _, height, width = latents.shape
+
+        t = torch.arange(1)
+        h = torch.arange(height)
+        w = torch.arange(width)
+        l = torch.arange(1)
+
+        latent_ids = torch.cartesian_prod(t, h, w, l)
+        latent_ids = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+        return latent_ids.to(device=latents.device)
+
+    @staticmethod
+    def _prepare_condition_latent_ids(
+            image_latents: List[torch.Tensor],
+            scale: int = 10):
+        if not isinstance(image_latents, list):
+            raise ValueError(f"Expected `image_latents` to be a list, got {type(image_latents)}.")
+
+        t_coords = [scale + scale * t for t in torch.arange(0, len(image_latents))]
+        t_coords = [t.view(-1) for t in t_coords]
+
+        image_latent_ids = []
+        for x, t in zip(image_latents, t_coords):
+            _, _, h, w = x.shape
+            x_ids = torch.cartesian_prod(t, torch.arange(h), torch.arange(w), torch.arange(1))
+            image_latent_ids.append(x_ids)
+
+        image_latent_ids = torch.cat(image_latent_ids, dim=0)
+        image_latent_ids = image_latent_ids.unsqueeze(0).expand(image_latents[0].size(0), -1, -1)
+
+        return image_latent_ids.to(device=image_latents[0].device)
+
+    def _get_rotary_emb(self, img_ids, txt_ids):
+        if img_ids.ndim == 3:
+            img_ids = img_ids[0]
+        if txt_ids.ndim == 3:
+            txt_ids = txt_ids[0]
+
+        with torch.autocast(device_type='cuda', dtype=torch.float32, enabled=False):
+            img_ids = img_ids.float()
+            image_rotary_emb = self.pos_embed(img_ids)
+            text_rotary_emb = self.pos_embed(txt_ids)
+            return (
+                torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+                torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+            )
+
+    @apply_lora_scale("joint_attention_kwargs")
+    def forward(
+            self,
+            x_t: torch.Tensor,
+            timestep: torch.Tensor,
+            encoder_hidden_states: torch.Tensor = None,
+            condition_latents: List[torch.Tensor] | None = None,
+            txt_ids: torch.Tensor = None,
+            guidance: torch.Tensor = None,
+            joint_attention_kwargs: dict[str, Any] | None = None):
+        x_t = self.patchify(x_t, self.patch_size)
+        img_ids = self._prepare_latent_ids(x_t)
+
+        bs, _, h, w = x_t.size()
+        x_t_packed = self.pack(x_t)
+        num_x_tokens = x_t_packed.size(1)
+        packed_ndim = x_t_packed.dim()
+
+        # scale and timestep calibration
+        calibration = self.asymflow_calibration(timestep, bs, packed_ndim)
+        hidden_states = x_t_packed * calibration.k.to(x_t_packed.dtype)
+
+        input_img_ids = img_ids
+        if condition_latents is not None:
+            condition_hidden_states = [self.patchify(z, self.patch_size) for z in condition_latents]
+            condition_latent_ids = self._prepare_condition_latent_ids(condition_hidden_states)
+            condition_hidden_states = [self.pack(z) / calibration.s for z in condition_hidden_states]
+            hidden_states = torch.cat([hidden_states] + condition_hidden_states, dim=1)
+            input_img_ids = torch.cat([img_ids, condition_latent_ids], dim=1)
+
+        hidden_states = self.x_embedder(hidden_states)
+
+        num_txt_tokens = encoder_hidden_states.shape[1]
+
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+        temb = self.time_guidance_embed(calibration.timestep.to(hidden_states.dtype) * 1000, guidance)
+
+        double_stream_mod_img = self.double_stream_modulation_img(temb)
+        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
+        single_stream_mod = self.single_stream_modulation(temb)
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+        concat_rotary_emb = self._get_rotary_emb(input_img_ids, txt_ids)
+
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    double_stream_mod_img,
+                    double_stream_mod_txt,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb_mod_img=double_stream_mod_img,
+                    temb_mod_txt=double_stream_mod_txt,
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for block in self.single_transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    None,
+                    single_stream_mod,
+                    concat_rotary_emb,
+                    joint_attention_kwargs,
+                )
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=None,
+                    temb_mod=single_stream_mod,
+                    image_rotary_emb=concat_rotary_emb,
+                    joint_attention_kwargs=joint_attention_kwargs,
+                )
+
+        hidden_states = hidden_states[:, num_txt_tokens:num_txt_tokens + num_x_tokens]
+        hidden_states = self.norm_out(hidden_states, temb)
+
+        u_a_packed = self.proj_out(hidden_states)
+
+        output_packed = self.asymflow_velocity(u_a_packed, x_t_packed, calibration)
+
+        output = self.unpack(output_packed.to(hidden_states.dtype), h, w)
+        output = self.unpatchify(output, self.patch_size)
+
+        return output

lakonlab/models/architectures/asymflow/common.py

@@ -0,0 +1,70 @@
+from typing import NamedTuple
+
+import torch
+import torch.nn.functional as F
+
+
+class AsymFlowCalibration(NamedTuple):
+    s: torch.Tensor
+    k: torch.Tensor
+    timestep: torch.Tensor
+    sigma: torch.Tensor
+
+
+class AsymFlowMixin:
+
+    train_sigma_min = 1e-6
+
+    def init_asymflow_buffers(self, patch_dim: int, base_rank: int):
+        assert patch_dim >= base_rank
+        eye = torch.eye(base_rank)
+        self.register_buffer('proj_buffer', F.pad(eye, (0, 0, 0, patch_dim - base_rank)))  # (patch_dim, base_rank)
+        self.register_buffer('scale_buffer', torch.tensor(1.0))
+
+    def asymflow_calibration(self, timestep, batch_size: int, ndim: int):
+        with torch.autocast(device_type='cuda', dtype=torch.float32, enabled=False):
+            timestep = timestep.float()
+            s = self.scale_buffer.float()
+            sigma = timestep / self.num_timesteps
+            k = 1 / (s + (1 - s) * sigma)
+            cal_timestep = timestep * k
+            sigma = sigma.expand(batch_size).reshape(batch_size, *((ndim - 1) * [1])).float()
+            k = k.reshape(batch_size, *((ndim - 1) * [1]))
+            return AsymFlowCalibration(
+                s=s,
+                k=k,
+                timestep=cal_timestep,
+                sigma=sigma,
+            )
+
+    @staticmethod
+    def orthogonal_decomposition(full_rank_state, proj_buffer):
+        subspace = full_rank_state @ proj_buffer @ proj_buffer.T
+        complement = full_rank_state - subspace
+        return subspace, complement
+
+    def asymflow_velocity(
+            self,
+            u_a_packed,
+            x_t_packed,
+            calibration: AsymFlowCalibration):
+        with torch.autocast(device_type='cuda', dtype=torch.float32, enabled=False):
+            sigma_min = self.train_sigma_min if self.training else self.sigma_min
+            u_a_packed = u_a_packed.float()
+            x_t_packed = x_t_packed.float()
+            proj_buffer = self.proj_buffer.float()
+            # orthogonal decomposition
+            u_a_subspace, u_a_complement = self.orthogonal_decomposition(u_a_packed, proj_buffer)
+            x_t_subspace, x_t_complement = self.orthogonal_decomposition(x_t_packed, proj_buffer)
+            # read calibration output
+            sk = calibration.s * calibration.k
+            sigma_clamped = calibration.sigma.clamp(min=sigma_min)
+            # low-rank subspace
+            u_subspace = (
+                sk * u_a_subspace
+                + (1 - sk) / sigma_clamped * x_t_subspace
+            )
+            # orthogonal complement
+            u_complement = (x_t_complement + calibration.s * u_a_complement) / sigma_clamped
+            # full velocity
+            return u_subspace + u_complement

lakonlab/models/architectures/autoencoders/__init__.py

@@ -0,0 +1,3 @@
+from .color_encoders import OklabColorEncoder, RGBColorEncoder
+
+__all__ = ['OklabColorEncoder', 'RGBColorEncoder']

lakonlab/models/architectures/autoencoders/color_encoders.py

@@ -0,0 +1,126 @@
+import torch
+
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+
+class OklabColorEncoder(ModelMixin, ConfigMixin):
+
+    @register_to_config
+    def __init__(
+        self,
+        use_affine_norm=True,
+        mean=(0.5, 0.0, 0.0),
+        std=0.21,
+    ):
+        super().__init__()
+        self.use_affine_norm = use_affine_norm
+        self.register_buffer('lrgb_to_lms', torch.tensor([
+            [0.4122214708, 0.5363325363, 0.0514459929],
+            [0.2119034982, 0.6806995451, 0.1073969566],
+            [0.0883024619, 0.2817188376, 0.6299787005]
+        ], dtype=torch.float32))
+        self.register_buffer('lms_to_oklab', torch.tensor([
+            [0.2104542553, 0.7936177850, -0.0040720468],
+            [1.9779984951, -2.4285922050, 0.4505937099],
+            [0.0259040371, 0.7827717662, -0.8086757660]
+        ], dtype=torch.float32))
+        self.register_buffer('oklab_to_lms', torch.linalg.inv(self.lms_to_oklab))
+        self.register_buffer('lms_to_lrgb', torch.linalg.inv(self.lrgb_to_lms))
+        if self.use_affine_norm:
+            self.register_buffer('affine_mean', torch.tensor(mean, dtype=torch.float32))
+            self.register_buffer('affine_std', torch.tensor(std, dtype=torch.float32))
+
+    @property
+    def dtype(self):
+        return self.lrgb_to_lms.dtype
+
+    @staticmethod
+    def srgb_to_lrgb(srgb):
+        a = 0.055
+        return torch.where(srgb <= 0.04045, srgb / 12.92, ((srgb + a) / (1 + a)) ** 2.4)
+
+    @staticmethod
+    def lrgb_to_srgb(lrgb):
+        lrgb = lrgb.clamp(min=0)
+        a = 0.055
+        return torch.where(lrgb <= 0.0031308, lrgb * 12.92, (1 + a) * (lrgb ** (1 / 2.4)) - a)
+
+    def lrgb_to_oklab(self, lrgb):
+        """
+        Args:
+            lrgb (torch.Tensor): Linear RGB, shape (N, 3, *)
+        """
+        lms = torch.einsum('ij,bj...->bi...', self.lrgb_to_lms, lrgb).clamp(min=0)
+        oklab = torch.einsum('ij,bj...->bi...', self.lms_to_oklab, lms.pow(1/3))
+        return oklab
+
+    def oklab_to_lrgb(self, oklab):
+        """
+        Args:
+            oklab (torch.Tensor): Oklab, shape (N, 3, *)
+        """
+        lms = torch.einsum('ij,bj...->bi...', self.oklab_to_lms, oklab).pow(3)
+        lrgb = torch.einsum('ij,bj...->bi...', self.lms_to_lrgb, lms)
+        return lrgb.clamp(0, 1)
+
+    def encode(self, img):
+        rgb = img / 2 + 0.5
+        lrgb = self.srgb_to_lrgb(rgb)
+        oklab = self.lrgb_to_oklab(lrgb)
+        if self.use_affine_norm:
+            n_dim = img.dim() - 2
+            mean = self.affine_mean.reshape(-1, *([1] * n_dim))
+            std = self.affine_std.reshape(-1, *([1] * n_dim))
+            oklab = (oklab - mean) / std
+        return oklab
+
+    def decode(self, oklab):
+        if self.use_affine_norm:
+            n_dim = oklab.dim() - 2
+            mean = self.affine_mean.reshape(-1, *([1] * n_dim))
+            std = self.affine_std.reshape(-1, *([1] * n_dim))
+            oklab = oklab * std + mean
+        lrgb = self.oklab_to_lrgb(oklab)
+        rgb = self.lrgb_to_srgb(lrgb)
+        img = rgb * 2 - 1
+        return img
+
+
+class RGBColorEncoder(ModelMixin, ConfigMixin):
+
+    @register_to_config
+    def __init__(
+        self,
+        use_affine_norm=False,
+        mean=(0.0, 0.0, 0.0),
+        std=1.0,
+    ):
+        super().__init__()
+        self.use_affine_norm = use_affine_norm
+        if self.use_affine_norm:
+            self.register_buffer('affine_mean', torch.tensor(mean, dtype=torch.float32))
+            self.register_buffer('affine_std', torch.tensor(std, dtype=torch.float32))
+
+    @property
+    def dtype(self):
+        if self.use_affine_norm:
+            return self.affine_mean.dtype
+        else:
+            return torch.float32
+
+    def encode(self, img):
+        if self.use_affine_norm:
+            n_dim = img.dim() - 2
+            mean = self.affine_mean.reshape(-1, *([1] * n_dim))
+            std = self.affine_std.reshape(-1, *([1] * n_dim))
+            img = (img - mean) / std
+        return img
+
+    def decode(self, img):
+        if self.use_affine_norm:
+            n_dim = img.dim() - 2
+            mean = self.affine_mean.reshape(-1, *([1] * n_dim))
+            std = self.affine_std.reshape(-1, *([1] * n_dim))
+            img = img * std + mean
+        return img.clamp(min=-1, max=1)

lakonlab/models/diffusions/__init__.py

@@ -0,0 +1 @@
+

lakonlab/models/diffusions/gaussian_flow.py

@@ -0,0 +1,20 @@
+from typing import Optional
+
+import torch
+
+
+@torch.jit.script
+def guidance_jit(
+        pos_mean, neg_mean, guidance_scale,
+        orthogonal: float = 1.0, parallel_dir: Optional[torch.Tensor] = None):
+    bias = (pos_mean - neg_mean) * (guidance_scale - 1)
+    if orthogonal:
+        dim = list(range(1, pos_mean.dim()))
+        if parallel_dir is None:
+            parallel_dir = pos_mean
+        bias = bias - ((bias * parallel_dir).mean(
+            dim=dim, keepdim=True
+        ) / (parallel_dir * parallel_dir).mean(
+            dim=dim, keepdim=True
+        ).clamp(min=1e-6) * parallel_dir).mul(orthogonal)
+    return bias

lakonlab/models/diffusions/schedulers/__init__.py

@@ -0,0 +1,3 @@
+from .flow_adapter import FlowAdapterScheduler
+
+__all__ = ['FlowAdapterScheduler']

lakonlab/models/diffusions/schedulers/flow_adapter.py

@@ -0,0 +1,252 @@
+# Copyright (c) 2025 Hansheng Chen
+
+import inspect
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union, List
+
+import numpy as np
+import torch
+import diffusers
+
+from diffusers.configuration_utils import register_to_config
+from diffusers.utils import BaseOutput
+from diffusers.schedulers import SchedulerMixin
+from diffusers.configuration_utils import ConfigMixin
+
+
+@dataclass
+class FlowWrapperSchedulerOutput(BaseOutput):
+    prev_sample: torch.FloatTensor
+
+
+class FlowAdapterScheduler(SchedulerMixin, ConfigMixin):
+
+    order = 1
+
+    @register_to_config
+    def __init__(
+            self,
+            num_train_timesteps: int = 1000,
+            shift: float = 1.0,
+            use_dynamic_shifting=False,
+            dynamic_shifting_type='exp',
+            base_seq_len=256,
+            max_seq_len=4096,
+            base_logshift=0.5,
+            max_logshift=1.15,
+            terminal_sigma=None,
+            base_scheduler='UniPCMultistep',
+            eps=1e-4,
+            **kwargs):
+
+        sigmas = torch.from_numpy(1 - np.linspace(
+            0, 1, num_train_timesteps, dtype=np.float32, endpoint=False))
+        self.sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+        self.timesteps = self.sigmas * num_train_timesteps
+        alphas = 1 - self.sigmas
+
+        base_scheduler_class = getattr(diffusers.schedulers, base_scheduler + 'Scheduler', None)
+
+        if base_scheduler_class is None:
+            raise AttributeError(f'Cannot find base_scheduler [{base_scheduler}].')
+        if base_scheduler in ['EulerDiscrete', 'EulerAncestralDiscrete']:
+            assert kwargs.get('prediction_type', 'epsilon') == 'epsilon'
+            kwargs['prediction_type'] = 'epsilon'
+            self.scales = ((alphas ** 2 + self.sigmas ** 2) / (
+                    1 + (self.sigmas / alphas.clamp(min=self.config.eps)) ** 2)).sqrt()
+        elif base_scheduler in [
+                'UniPCMultistep', 'DPMSolverSinglestep', 'DPMSolverMultistep', 'DEISMultistep', 'SASolver']:
+            self.scales = torch.ones_like(alphas)
+            assert kwargs.get('prediction_type', 'flow_prediction') == 'flow_prediction'
+            kwargs['prediction_type'] = 'flow_prediction'
+            kwargs['use_flow_sigmas'] = True
+        else:
+            raise AttributeError(f'Unsupported base_scheduler [{base_scheduler}].')
+
+        signatures = inspect.signature(base_scheduler_class).parameters.keys()
+        if 'final_sigmas_type' in signatures:
+            kwargs['final_sigmas_type'] = 'zero'
+        if 'lower_order_final' in signatures:
+            kwargs['lower_order_final'] = True
+
+        self.base_scheduler = base_scheduler_class(
+            num_train_timesteps=num_train_timesteps,
+            **kwargs)
+        self.base_scheduler.timesteps = self.timesteps
+        if self.config.base_scheduler in ['EulerDiscrete', 'EulerAncestralDiscrete']:
+            self.base_scheduler.sigmas = self.sigmas / alphas.clamp(min=self.config.eps)
+        elif self.config.base_scheduler in [
+                'UniPCMultistep', 'DPMSolverSinglestep', 'DPMSolverMultistep', 'DEISMultistep', 'SASolver']:
+            self.base_scheduler.sigmas = self.sigmas
+        else:
+            raise AttributeError(f'Unsupported base_scheduler [{self.config.base_scheduler}].')
+
+        self._step_index = None
+        self._begin_index = None
+
+    @property
+    def step_index(self):
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        self._begin_index = begin_index
+
+    def get_shift(self, seq_len=None):
+        if self.config.use_dynamic_shifting and seq_len is not None:
+            if self.config.dynamic_shifting_type == 'exp':
+                m = (self.config.max_logshift - self.config.base_logshift
+                     ) / (self.config.max_seq_len - self.config.base_seq_len)
+                logshift = (seq_len - self.config.base_seq_len) * m + self.config.base_logshift
+                if isinstance(logshift, torch.Tensor):
+                    shift = torch.exp(logshift)
+                else:
+                    shift = np.exp(logshift)
+            elif self.config.dynamic_shifting_type == 'sqrt':
+                max_shift = np.exp(self.config.max_logshift)
+                base_shift = np.exp(self.config.base_logshift)
+                sqrt_max_seq_len = np.sqrt(self.config.max_seq_len)
+                sqrt_base_seq_len = np.sqrt(self.config.base_seq_len)
+                m = (max_shift - base_shift) / (sqrt_max_seq_len - sqrt_base_seq_len)
+                shift = (np.sqrt(seq_len) - sqrt_base_seq_len) * m + base_shift
+            else:
+                raise ValueError(f'Unsupported dynamic_shifting_type [{self.config.dynamic_shifting_type}].')
+        else:
+            shift = self.config.shift
+        return shift
+
+    def stretch_to_terminal(self, sigma):
+        one_minus_sigma = 1 - sigma
+        stretched_sigma = 1 - (one_minus_sigma * (1 - self.config.terminal_sigma) / one_minus_sigma[-1])
+        return stretched_sigma
+
+    def set_timesteps(
+            self,
+            num_inference_steps: Optional[int] = None,
+            sigmas: Optional[List[float]] = None,
+            seq_len=None,
+            device=None):
+        if sigmas is None:
+            assert num_inference_steps is not None, 'Either num_inference_steps or sigmas must be provided.'
+            self.num_inference_steps = num_inference_steps
+            sigmas = np.linspace(1, 0, num_inference_steps, dtype=np.float32, endpoint=False)
+        else:
+            if num_inference_steps is not None:
+                assert len(sigmas) == num_inference_steps
+            self.num_inference_steps = len(sigmas)
+            sigmas = np.array(sigmas, dtype=np.float32)
+
+        sigmas = torch.from_numpy(sigmas)
+        shift = self.get_shift(seq_len=seq_len)
+        sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
+
+        if self.config.terminal_sigma is not None:
+            sigmas = self.stretch_to_terminal(sigmas)
+
+        self.timesteps = (sigmas * self.config.num_train_timesteps).to(device)
+        if self.config.base_scheduler in ['DEISMultistep', 'SASolver']:
+            self.sigmas = torch.cat(
+                [sigmas, torch.tensor([self.config.eps], dtype=torch.float32, device=sigmas.device)])
+        else:
+            self.sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)])
+        alphas = 1 - self.sigmas
+
+        self.base_scheduler.set_timesteps(num_inference_steps, device=device)
+
+        self.base_scheduler.timesteps = self.timesteps
+        if self.config.base_scheduler in ['EulerDiscrete', 'EulerAncestralDiscrete']:
+            self.base_scheduler.sigmas = self.sigmas / alphas.clamp(min=self.config.eps)
+            self.scales = ((alphas ** 2 + self.sigmas ** 2) / (
+                    1 + (self.sigmas / alphas.clamp(min=self.config.eps)) ** 2)).sqrt()
+        elif self.config.base_scheduler in [
+                'UniPCMultistep', 'DPMSolverSinglestep', 'DPMSolverMultistep', 'DEISMultistep', 'SASolver']:
+            self.base_scheduler.sigmas = self.sigmas.clamp(max=1 - self.config.eps)
+            self.scales = torch.ones_like(alphas)
+        else:
+            raise AttributeError(f'Unsupported base_scheduler [{self.config.base_scheduler}].')
+
+        self._step_index = None
+        self._begin_index = None
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        indices = (schedule_timesteps == timestep).nonzero()
+
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        pos = 1 if len(indices) > 1 else 0
+
+        return indices[pos].item()
+
+    def _init_step_index(self, timestep):
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+            self,
+            model_output: torch.FloatTensor,
+            timestep: Union[float, torch.FloatTensor],
+            sample: torch.FloatTensor,
+            generator: Optional[torch.Generator] = None,
+            return_dict: bool = True,
+            prediction_type='u',
+            eps=1e-6) -> Union[FlowWrapperSchedulerOutput, Tuple]:
+        assert prediction_type in ['u', 'x0']
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Upcast to avoid precision issues when computing prev_sample
+        ori_dtype = model_output.dtype
+        sample = sample.to(torch.float32)
+        model_output = model_output.to(torch.float32)
+
+        sigma = self.sigmas[self.step_index]
+        alpha = 1 - sigma
+        scale = self.scales[self.step_index]
+        next_scale = self.scales[self.step_index + 1]
+
+        if hasattr(self.base_scheduler, 'is_scale_input_called'):
+            self.base_scheduler.is_scale_input_called = True
+        kwargs = dict(return_dict=False)
+        if generator is not None:
+            kwargs.update(generator=generator)
+
+        if self.config.base_scheduler in [
+                'UniPCMultistep', 'DPMSolverSinglestep', 'DPMSolverMultistep', 'DEISMultistep', 'SASolver']:
+            if prediction_type == 'u':
+                model_output = model_output
+            else:
+                model_output = (sample - model_output) / sigma.clamp(min=eps)
+        else:  # to epsilon
+            if prediction_type == 'u':
+                model_output = sample + alpha * model_output
+            else:
+                model_output = (sample - alpha * model_output) / sigma.clamp(min=eps)
+        prev_sample = self.base_scheduler.step(
+            model_output,
+            timestep,
+            sample / scale,
+            **kwargs
+        )[0] * next_scale
+
+        prev_sample = prev_sample.to(ori_dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return FlowWrapperSchedulerOutput(prev_sample=prev_sample)

lakonlab/pipelines/__init__.py

@@ -0,0 +1 @@
+

lakonlab/pipelines/pipeline_pixelflux2_klein.py

@@ -0,0 +1,310 @@
+from typing import Any, Callable
+
+import PIL
+import torch
+
+from transformers import Qwen3ForCausalLM, Qwen2TokenizerFast
+from diffusers.utils import is_torch_xla_available
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.models import Flux2Transformer2DModel
+from diffusers.pipelines.flux2.pipeline_flux2_klein import (
+    Flux2KleinPipeline, Flux2PipelineOutput, Flux2ImageProcessor)
+from .utils import LakonLabMixin
+from lakonlab.models.diffusions.gaussian_flow import guidance_jit
+from lakonlab.models.diffusions.schedulers import FlowAdapterScheduler
+from lakonlab.models.architectures import OklabColorEncoder
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+class PixelFlux2KleinPipeline(Flux2KleinPipeline, LakonLabMixin):
+
+    model_cpu_offload_seq = "text_encoder->transformer"
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowAdapterScheduler,
+        vae: OklabColorEncoder,
+        text_encoder: Qwen3ForCausalLM,
+        tokenizer: Qwen2TokenizerFast,
+        transformer: Flux2Transformer2DModel,
+        is_distilled: bool = False,
+    ):
+        super(Flux2KleinPipeline, self).__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            scheduler=scheduler,
+            transformer=transformer,
+        )
+
+        self.register_to_config(is_distilled=is_distilled)
+
+        self.vae_scale_factor = 1
+        self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 16)
+        self.tokenizer_max_length = 512
+        self.default_sample_size = 1024
+
+    def prepare_latents(
+        self,
+        batch_size,
+        height,
+        width,
+        dtype,
+        device,
+        generator: torch.Generator,
+        latents: torch.Tensor | None = None,
+    ):
+        height = 16 * (int(height) // (self.vae_scale_factor * 16))
+        width = 16 * (int(width) // (self.vae_scale_factor * 16))
+
+        shape = (batch_size, 3, height, width)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+
+        return latents
+
+    def prepare_image_latents(
+        self,
+        images: list[torch.Tensor],
+        batch_size,
+        device,
+        dtype,
+    ):
+        image_latents = []
+        for image in images:
+            image = image.to(device=device, dtype=dtype)
+            imagge_latent = self.vae.encode(image).to(self.transformer.dtype)
+            image_latents.append(imagge_latent.repeat(batch_size, 1, 1, 1))  # (bs, 3, 1024, 1024)
+        return image_latents
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        image: list[PIL.Image.Image] | PIL.Image.Image | None = None,
+        prompt: str | list[str] = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int = 50,
+        sigmas: list[float] | None = None,
+        guidance_scale: float = 4.0,
+        orthogonal_guidance: float = 1.0,
+        clamp_denoised: bool = True,
+        num_images_per_prompt: int = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: str | list[str] | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        callback_on_step_end: Callable[[int, int, dict], None] | None = None,
+        callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+        max_sequence_length: int = 512,
+        text_encoder_out_layers: tuple[int] = (9, 18, 27),
+    ):
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt=prompt,
+            height=height,
+            width=width,
+            prompt_embeds=prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            guidance_scale=guidance_scale,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. prepare text embeddings
+        prompt_embeds, text_ids = self.encode_prompt(
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            text_encoder_out_layers=text_encoder_out_layers,
+        )
+
+        if self.do_classifier_free_guidance:
+            if negative_prompt is None:
+                negative_prompt = ""
+            if prompt is not None and isinstance(prompt, list) and not isinstance(negative_prompt, list):
+                negative_prompt = [negative_prompt] * len(prompt)
+            negative_prompt_embeds, negative_text_ids = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                text_encoder_out_layers=text_encoder_out_layers,
+            )
+            guidance_scale = torch.tensor(guidance_scale, device=device, dtype=torch.float32)
+
+        # 4. process images
+        if image is not None and not isinstance(image, list):
+            image = [image]
+
+        condition_images = None
+        if image is not None:
+            for img in image:
+                self.image_processor.check_image_input(img)
+
+            condition_images = []
+            for img in image:
+                image_width, image_height = img.size
+                if image_width * image_height > 1024 * 1024:
+                    img = self.image_processor._resize_to_target_area(img, 1024 * 1024)
+                    image_width, image_height = img.size
+
+                multiple_of = self.vae_scale_factor * 16
+                image_width = (image_width // multiple_of) * multiple_of
+                image_height = (image_height // multiple_of) * multiple_of
+                img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop")
+                condition_images.append(img)
+                height = height or image_height
+                width = width or image_width
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 5. prepare latent variables
+        latents = self.prepare_latents(
+            batch_size=batch_size * num_images_per_prompt,
+            height=height,
+            width=width,
+            dtype=torch.float32,
+            device=device,
+            generator=generator,
+            latents=latents,
+        )
+
+        image_latents = None
+        if condition_images is not None:
+            image_latents = self.prepare_image_latents(
+                images=condition_images,
+                batch_size=batch_size * num_images_per_prompt,
+                device=device,
+                dtype=self.vae.dtype,
+            )
+
+        # 6. Prepare timesteps
+        image_seq_len = latents.shape[2:].numel()
+        self.scheduler.set_timesteps(
+            num_inference_steps, seq_len=image_seq_len, device=self._execution_device)
+        timesteps = self.scheduler.timesteps
+        self._num_timesteps = len(timesteps)
+
+        # 7. Denoising loop
+        # We set the index here to remove DtoH sync, helpful especially during compilation.
+        # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                _t = t / 1000
+                timestep = _t.expand(latents.shape[0]).to(latents.dtype)
+                latent_model_input = latents.to(self.transformer.dtype)
+
+                with self.transformer.cache_context("cond"):
+                    denoising_output = self.transformer(
+                        x_t=latent_model_input,  # (B, 3, H, W)
+                        timestep=timestep,
+                        encoder_hidden_states=prompt_embeds,
+                        condition_latents=image_latents,
+                        txt_ids=text_ids,  # B, text_seq_len, 4
+                        guidance=None,
+                        joint_attention_kwargs=self.attention_kwargs,
+                    ).float()
+
+                if self.do_classifier_free_guidance:
+                    with self.transformer.cache_context("uncond"):
+                        neg_denoising_output = self.transformer(
+                            x_t=latent_model_input,  # (B, 3, H, W)
+                            timestep=timestep,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            condition_latents=image_latents,
+                            txt_ids=negative_text_ids,
+                            guidance=None,
+                            joint_attention_kwargs=self._attention_kwargs,
+                        ).float()
+                    cfg_bias = guidance_jit(
+                        denoising_output,
+                        neg_denoising_output,
+                        guidance_scale,
+                        orthogonal_guidance,
+                        latents - denoising_output * _t)
+                    denoising_output = denoising_output + cfg_bias
+
+                if clamp_denoised:
+                    denoised = latents - denoising_output * _t
+                    image = self.vae.decode(denoised.to(self.vae.dtype)).clamp(-1, 1)
+                    denoised = self.vae.encode(image).to(latents.dtype)
+                    denoising_output = (latents - denoised) / _t.clamp(min=1e-4)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents = self.scheduler.step(denoising_output, t, latents, return_dict=False)[0]
+
+                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, 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:
+            image = self.vae.decode(latents.to(self.vae.dtype))
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return Flux2PipelineOutput(images=image)

lakonlab/pipelines/prompt_rewriters/qwen3_vl.py

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

lakonlab/pipelines/prompt_rewriters/system_prompts/default_text_only.txt

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

lakonlab/pipelines/prompt_rewriters/system_prompts/default_with_images.txt

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

lakonlab/pipelines/utils.py

@@ -0,0 +1,209 @@
+import os
+from typing import Optional, Union
+
+import accelerate
+import torch
+from diffusers.models import AutoModel
+from diffusers.models.modeling_utils import (
+    _LOW_CPU_MEM_USAGE_DEFAULT,
+    ContextManagers,
+    load_state_dict,
+    no_init_weights,
+)
+from diffusers.quantizers import DiffusersAutoQuantizer
+from diffusers.utils import (
+    SAFETENSORS_WEIGHTS_NAME,
+    WEIGHTS_NAME,
+    _add_variant,
+    _get_model_file,
+    is_accelerate_available,
+    is_torch_version,
+    logging,
+)
+from diffusers.utils.torch_utils import empty_device_cache
+
+from lakonlab.models.architectures.asymflow.asymflux2 import _AsymFlux2Transformer2DModel
+
+
+LOCAL_CLASS_MAPPING = {
+    'AsymFlux2Transformer2DModel': _AsymFlux2Transformer2DModel,
+}
+
+logger = logging.get_logger(__name__)
+
+
+def assign_param(module, tensor_name: str, param: torch.nn.Parameter):
+    if '.' in tensor_name:
+        splits = tensor_name.split('.')
+        for split in splits[:-1]:
+            module = getattr(module, split)
+        tensor_name = splits[-1]
+    module._parameters[tensor_name] = param
+
+
+class LakonLabMixin:
+
+    def load_lakonlab_adapter(
+            self,
+            pretrained_model_name_or_path: Union[str, os.PathLike],
+            target_module_name: str = 'transformer',
+            adapter_name: Optional[str] = None,
+            **kwargs):
+        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('accelerate is not available; using low_cpu_mem_usage=False.')
+
+        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.')
+
+        user_agent = {
+            'diffusers': 'demo',
+            'file_type': 'model',
+            'framework': 'pytorch',
+        }
+
+        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)
+        model_cls = LOCAL_CLASS_MAPPING[config['_class_name']]
+
+        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:
+                if not allow_pickle:
+                    raise
+
+        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,
+            )
+
+        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)
+
+        overwrite_state_dict = {}
+        lora_state_dict = {}
+        adapter_state_dict = load_state_dict(model_file, disable_mmap=disable_mmap)
+        for key, value in adapter_state_dict.items():
+            value = value.to(dtype=torch_dtype, device=device)
+            key = key.removeprefix(f'{target_module_name}.')
+            if 'lora' in key:
+                lora_state_dict[key] = value
+            else:
+                overwrite_state_dict[key] = value
+
+        pre_quantized = (
+            'quantization_config' in base_module.config
+            and base_module.config['quantization_config'] is not None)
+        if pre_quantized:
+            config['quantization_config'] = base_module.config.quantization_config
+            hf_quantizer = DiffusersAutoQuantizer.from_config(
+                config['quantization_config'], pre_quantized=True)
+            hf_quantizer.validate_environment(torch_dtype=torch_dtype)
+            torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype)
+            if low_cpu_mem_usage is None:
+                low_cpu_mem_usage = True
+            elif not low_cpu_mem_usage:
+                raise ValueError('low_cpu_mem_usage cannot be False with quantization.')
+        else:
+            hf_quantizer = None
+
+        keep_in_fp32_modules = []
+        for key in overwrite_state_dict.keys():
+            module_name = key.rsplit('.', 1)[0]
+            if module_name and module_name not in keep_in_fp32_modules:
+                keep_in_fp32_modules.append(module_name)
+
+        init_contexts = [no_init_weights()]
+        if low_cpu_mem_usage:
+            init_contexts.append(accelerate.init_empty_weights())
+
+        with ContextManagers(init_contexts):
+            adapter_module = model_cls.from_config(config).eval()
+
+        torch.set_default_dtype(dtype_orig)
+
+        if hf_quantizer is not None:
+            hf_quantizer.preprocess_model(
+                model=adapter_module, device_map=None, keep_in_fp32_modules=keep_in_fp32_modules)
+
+        base_state_dict = base_module.state_dict()
+        base_state_dict.update(overwrite_state_dict)
+        empty_state_dict = adapter_module.state_dict()
+        for param_name, param in base_state_dict.items():
+            if param_name not in empty_state_dict:
+                continue
+            if hf_quantizer is not None and hf_quantizer.check_if_quantized_param(
+                    adapter_module, param, param_name, base_state_dict, param_device=device):
+                hf_quantizer.create_quantized_param(
+                    adapter_module, param_name=param_name, param_value=param,
+                    target_device=device, state_dict=base_state_dict,
+                    unexpected_keys=[], dtype=torch_dtype)
+            else:
+                assign_param(adapter_module, param_name, param)
+
+        empty_device_cache()
+
+        if hf_quantizer is not None:
+            hf_quantizer.postprocess_model(adapter_module)
+            adapter_module.hf_quantizer = hf_quantizer
+
+        if lora_state_dict:
+            if adapter_name is None:
+                adapter_name = f'{target_module_name}_lakonlab'
+            adapter_module.load_lora_adapter(
+                lora_state_dict, prefix=None, adapter_name=adapter_name,
+                low_cpu_mem_usage=low_cpu_mem_usage)
+        else:
+            adapter_name = None
+
+        setattr(self, target_module_name, adapter_module)
+        return adapter_name

lakonlab/ui/gradio/create_text_to_img.py

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

lakonlab/ui/gradio/shared_opts.py

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

lakonlab/ui/gradio/style.css

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

requirements.txt

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