scripts/train_omnigen.py

from collections import defaultdict
import contextlib
import datetime
from concurrent import futures
import hashlib
import itertools
import json
import math
import os
import random
import sys
import tempfile
import time

from absl import app, flags
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from huggingface_hub import snapshot_download
from ml_collections import config_flags
import numpy as np
from peft import LoraConfig, PeftModel, get_peft_model
from PIL import Image, ImageDraw
import torch
from torch.utils.data import DataLoader, Dataset, Sampler
import tqdm
import wandb

from diffusers.models import AutoencoderKL
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_outputs import BaseModelOutputWithPast

import flow_grpo.rewards
from flow_grpo.dataset_paths import DatasetPathResolver, resolve_metadata_paths
from flow_grpo.ema import EMAModuleWrapper
from flow_grpo.omnigen_patch.omnigen_pipeline_with_logprob import (
    build_group_ids,
    compute_cfg_model_output,
    compute_cfg_model_output_joint,
    filter_branch_inputs,
    omnigen_sde_step_with_logprob,
    pipeline_with_logprob,
    pipeline_with_logprob_joint_image_reward,
)
from flow_grpo.omnigen_patch.joint_model_loader import load_joint_omnigen_components_for_rl
from flow_grpo.stat_tracking import PerPromptStatTracker
from flow_grpo.utils.run_metadata import write_run_metadata


DEFAULT_OMNIGEN_CODE_ROOT = "/NAS_REMOTE/vicky/wt/codes/cxr_syn/gen_code"
OMNIGEN_CODE_ROOT = os.environ.get("OMNIGEN_CODE_ROOT", DEFAULT_OMNIGEN_CODE_ROOT)
if os.path.exists(OMNIGEN_CODE_ROOT) and OMNIGEN_CODE_ROOT not in sys.path:
    sys.path.insert(0, OMNIGEN_CODE_ROOT)

import transformers.cache_utils as hf_cache_utils  # noqa: E402

if not hasattr(hf_cache_utils, "OffloadedCache"):
    hf_cache_utils.OffloadedCache = hf_cache_utils.DynamicCache

from OmniGen import OmniGen, OmniGenProcessor  # noqa: E402
from OmniGen.transformer import Phi3Transformer  # noqa: E402


tqdm = tqdm.tqdm
FLAGS = flags.FLAGS
config_flags.DEFINE_config_file("config", "config/base.py", "Training configuration.")
logger = get_logger(__name__)


def _debug_flag(name: str) -> bool:
    return os.environ.get(name, "0").strip().lower() in {"1", "true", "yes", "on"}


def _to_rgb_pil(image):
    if isinstance(image, Image.Image):
        return image.convert("RGB")

    if isinstance(image, torch.Tensor):
        image = image.detach().cpu()
        if image.ndim == 3 and image.shape[0] in (1, 3):
            if image.shape[0] == 1:
                image = image.repeat(3, 1, 1)
            image = image.clamp(0, 1).permute(1, 2, 0).numpy()
        elif image.ndim == 2:
            image = image.numpy()
        else:
            image = image.numpy()

    array = np.asarray(image)
    if array.ndim == 3 and array.shape[0] in (1, 3) and array.shape[-1] not in (1, 3):
        array = np.transpose(array, (1, 2, 0))

    if array.dtype != np.uint8:
        array = np.clip(array, 0, 1) if np.issubdtype(array.dtype, np.floating) else np.clip(array, 0, 255)
        array = (array * 255).round().astype(np.uint8) if array.max(initial=0) <= 1.0 else array.astype(np.uint8)

    if array.ndim == 2:
        return Image.fromarray(array, mode="L").convert("RGB")
    if array.ndim == 3 and array.shape[-1] == 1:
        return Image.fromarray(array[..., 0], mode="L").convert("RGB")
    return Image.fromarray(array).convert("RGB")


def _load_optional_image(image_path):
    if not image_path:
        return None
    if not os.path.exists(image_path):
        return None
    return Image.open(image_path).convert("RGB")


def _truncate_text(text, max_length=220):
    text = str(text).strip().replace("\n", " ")
    if len(text) <= max_length:
        return text
    return text[: max_length - 3] + "..."


def _build_visual_panel(input_image, output_image, gt_image=None):
    panels = [("Input", _to_rgb_pil(input_image)), ("Output", _to_rgb_pil(output_image))]
    if gt_image is not None:
        panels.append(("GT", _to_rgb_pil(gt_image)))

    target_width = max(image.width for _, image in panels)
    target_height = max(image.height for _, image in panels)
    resized = [
        (
            label,
            image.resize((target_width, target_height), resample=Image.Resampling.BILINEAR),
        )
        for label, image in panels
    ]

    gap = 12
    title_height = 24
    canvas_width = len(resized) * target_width + (len(resized) - 1) * gap
    canvas_height = title_height + target_height
    canvas = Image.new("RGB", (canvas_width, canvas_height), color=(255, 255, 255))
    draw = ImageDraw.Draw(canvas)

    x = 0
    for label, image in resized:
        draw.text((x, 4), label, fill=(0, 0, 0))
        canvas.paste(image, (x, title_height))
        x += target_width + gap

    return canvas


def _build_wandb_examples(
    output_images,
    prompts,
    metadatas,
    ref_images,
    *,
    instructions=None,
    rewards=None,
    max_samples=8,
    note=None,
):
    num_samples = min(max_samples, len(output_images))
    examples = []
    for idx in range(num_samples):
        metadata = metadatas[idx]
        gt_image = _load_optional_image(metadata.get("gt_image"))
        panel = _build_visual_panel(ref_images[idx], output_images[idx], gt_image=gt_image)

        caption_parts = []
        sample_id = metadata.get("sample_id")
        if sample_id:
            caption_parts.append(str(sample_id))
        caption_parts.append(f"prompt: {_truncate_text(prompts[idx])}")

        if instructions is not None and instructions[idx] != prompts[idx]:
            caption_parts.append(f"instruction: {_truncate_text(instructions[idx])}")

        if rewards is not None:
            reward_parts = []
            for key, value in rewards.items():
                current_value = float(value[idx])
                if current_value != -10:
                    reward_parts.append(f"{key}: {current_value:.2f}")
            if reward_parts:
                caption_parts.append(" | ".join(reward_parts))

        if note:
            caption_parts.append(note)

        examples.append(wandb.Image(panel, caption="\n".join(caption_parts)))

    return examples


def patch_omnigen_phi3_forward():
    if getattr(Phi3Transformer.forward, "_flow_grpo_patched", False):
        return

    def _patched_forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        inputs_embeds=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        cache_position=None,
        offload_model=False,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if self.gradient_checkpointing and self.training and use_cache:
            logger.warning_once(
                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
            )
            use_cache = False

        return_legacy_cache = False
        if use_cache and not isinstance(past_key_values, Cache):
            return_legacy_cache = True
            if past_key_values is None:
                past_key_values = DynamicCache(config=self.config)
            else:
                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
                logger.warning_once(
                    "We detected tuple-based `past_key_values`. This is deprecated and was converted to DynamicCache."
                )

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens,
                past_seen_tokens + inputs_embeds.shape[1],
                device=inputs_embeds.device,
            )
        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)

        if attention_mask is not None and attention_mask.dim() == 3:
            dtype = inputs_embeds.dtype
            min_dtype = torch.finfo(dtype).min
            attention_mask = (1 - attention_mask) * min_dtype
            attention_mask = attention_mask.unsqueeze(1).to(inputs_embeds.dtype)
        else:
            raise Exception("attention_mask parameter was unavailable or invalid")

        hidden_states = inputs_embeds
        has_external_position_embeddings = hasattr(self, "rotary_emb")
        position_embeddings = (
            self.rotary_emb(hidden_states, position_ids=position_ids)
            if has_external_position_embeddings
            else None
        )

        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = None

        for layer_idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            if self.gradient_checkpointing and self.training:
                if has_external_position_embeddings:
                    layer_outputs = self._gradient_checkpointing_func(
                        decoder_layer.__call__,
                        hidden_states,
                        attention_mask,
                        position_ids,
                        past_key_values,
                        use_cache,
                        position_embeddings,
                    )
                else:
                    layer_outputs = self._gradient_checkpointing_func(
                        decoder_layer.__call__,
                        hidden_states,
                        attention_mask,
                        position_ids,
                        past_key_values,
                        output_attentions,
                        use_cache,
                        cache_position,
                    )
            else:
                if offload_model and not self.training:
                    self.get_offlaod_layer(layer_idx, device=inputs_embeds.device)
                if has_external_position_embeddings:
                    layer_outputs = decoder_layer(
                        hidden_states,
                        attention_mask=attention_mask,
                        position_ids=position_ids,
                        past_key_values=past_key_values,
                        output_attentions=output_attentions,
                        use_cache=use_cache,
                        position_embeddings=position_embeddings,
                        cache_position=cache_position,
                    )
                else:
                    layer_outputs = decoder_layer(
                        hidden_states,
                        attention_mask=attention_mask,
                        position_ids=position_ids,
                        past_key_value=past_key_values,
                        output_attentions=output_attentions,
                        use_cache=use_cache,
                        cache_position=cache_position,
                    )

            if isinstance(layer_outputs, tuple):
                hidden_states = layer_outputs[0]
                if use_cache and not has_external_position_embeddings:
                    next_decoder_cache = layer_outputs[2 if output_attentions else 1]
                if output_attentions and len(layer_outputs) > 1:
                    all_self_attns += (layer_outputs[1],)
            else:
                hidden_states = layer_outputs

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = (past_key_values if has_external_position_embeddings else next_decoder_cache) if use_cache else None
        if return_legacy_cache and next_cache is not None:
            next_cache = next_cache.to_legacy_cache()

        if not return_dict:
            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

    _patched_forward._flow_grpo_patched = True
    Phi3Transformer.forward = _patched_forward


patch_omnigen_phi3_forward()


class RadiomicsEditDataset(Dataset):
    def __init__(self, dataset, split="train", condition_dropout_prob=0.0):
        if isinstance(dataset, dict) or hasattr(dataset, "get"):
            key = f"{split}_jsonl"
            self.file_path = dataset.get(key) or dataset.get("jsonl")
            if self.file_path is None and dataset.get("root"):
                self.file_path = os.path.join(dataset.get("root"), f"{split}_metadata.jsonl")
            if self.file_path is None:
                raise ValueError(f"Dataset config is missing {key!r} for split={split}.")
        else:
            self.file_path = os.path.join(dataset, f"{split}_metadata.jsonl")
        self.path_resolver = DatasetPathResolver.from_dataset_config(dataset)
        self.condition_dropout_prob = condition_dropout_prob if split == "train" else 0.0
        with open(self.file_path, "r", encoding="utf-8") as f:
            self.metadatas = [json.loads(line) for line in f if line.strip()]

    def __len__(self):
        return len(self.metadatas)

    def __getitem__(self, idx):
        metadata = resolve_metadata_paths(
            dict(self.metadatas[idx]),
            self.path_resolver,
            required_keys=("input_images", "output_image", "output_mask"),
        )
        input_images = metadata.get("input_images") or []
        if len(input_images) == 0:
            raise ValueError("OmniGen edit training requires at least one input image.")

        instruction = metadata.get("instruction")
        if not instruction:
            instruction = f"<img><|image_1|></img> {metadata['prompt']}"

        primary_input_image = input_images[0]
        prompt = metadata.get("prompt") or metadata.get("instruction", "")
        output_image = metadata.get("output_image") or metadata.get("gt_image")
        group_key = str(
            metadata.get("sample_id")
            or output_image
            or (f"{primary_input_image}::{output_image}" if output_image else None)
            or f"{prompt}::{primary_input_image}"
        )

        if random.random() < self.condition_dropout_prob:
            instruction = "<cfg>"
            input_images = []

        return {
            "prompt": prompt,
            "instruction": instruction,
            "metadata": metadata,
            "input_image_paths": input_images,
            "ref_image": Image.open(primary_input_image).convert("RGB"),
            "group_key": group_key,
        }

    @staticmethod
    def collate_fn(examples):
        prompts = [example["prompt"] for example in examples]
        instructions = [example["instruction"] for example in examples]
        metadatas = [example["metadata"] for example in examples]
        input_image_paths = [example["input_image_paths"] for example in examples]
        ref_images = [example["ref_image"] for example in examples]
        group_keys = [example["group_key"] for example in examples]
        return prompts, instructions, metadatas, input_image_paths, ref_images, group_keys


class DistributedKRepeatSampler(Sampler):
    def __init__(self, dataset, batch_size, k, num_replicas, rank, num_batches_per_epoch, seed=0):
        self.dataset = dataset
        self.batch_size = batch_size
        self.k = k
        self.num_replicas = num_replicas
        self.rank = rank
        self.num_batches_per_epoch = num_batches_per_epoch
        self.seed = seed
        self.total_samples = self.num_replicas * self.batch_size
        self.group_span = self.k // math.gcd(self.total_samples, self.k)
        if self.num_batches_per_epoch % self.group_span != 0:
            raise ValueError(
                "num_batches_per_epoch must be divisible by the sampler group span. "
                f"Got num_batches_per_epoch={num_batches_per_epoch}, group_span={self.group_span}, "
                f"k={k}, num_replicas={num_replicas}, batch_size={batch_size}."
            )
        self.groups_per_epoch = self.num_batches_per_epoch // self.group_span
        self.prompts_per_group_pack = self.total_samples * self.group_span // self.k
        self.epoch = 0
        self._scheduled_epoch = None
        self._scheduled_batches = []
        self._batch_ptr = 0

    def _sample_unique_indices(self, generator, count):
        indices = []
        dataset_size = len(self.dataset)
        while len(indices) < count:
            indices.extend(torch.randperm(dataset_size, generator=generator).tolist())
        return indices[:count]

    def _build_epoch_batches(self):
        generator = torch.Generator()
        generator.manual_seed(self.seed + self.epoch)

        prompts_needed = self.groups_per_epoch * self.prompts_per_group_pack
        prompt_indices = self._sample_unique_indices(generator, prompts_needed)
        scheduled_batches = []

        for pack_index in range(self.groups_per_epoch):
            start = pack_index * self.prompts_per_group_pack
            end = start + self.prompts_per_group_pack
            current_indices = prompt_indices[start:end]
            repeated_indices = [idx for idx in current_indices for _ in range(self.k)]
            shuffled_indices = torch.randperm(len(repeated_indices), generator=generator).tolist()
            shuffled_samples = [repeated_indices[i] for i in shuffled_indices]

            for batch_offset in range(self.group_span):
                batch_start = batch_offset * self.total_samples
                batch_end = batch_start + self.total_samples
                global_batch = shuffled_samples[batch_start:batch_end]
                replica_start = self.rank * self.batch_size
                replica_end = replica_start + self.batch_size
                scheduled_batches.append(global_batch[replica_start:replica_end])

        return scheduled_batches

    def __iter__(self):
        while True:
            if self._scheduled_epoch != self.epoch or self._batch_ptr >= len(self._scheduled_batches):
                self._scheduled_batches = self._build_epoch_batches()
                self._scheduled_epoch = self.epoch
                self._batch_ptr = 0

            batch_indices = self._scheduled_batches[self._batch_ptr]
            self._batch_ptr += 1
            yield batch_indices

    def set_epoch(self, epoch):
        self.epoch = epoch
        self._scheduled_epoch = None
        self._scheduled_batches = []
        self._batch_ptr = 0


def create_generator(group_keys, base_seed, device):
    generators = []
    generator_device = device.type if isinstance(device, torch.device) else str(device)
    for key in group_keys:
        digest = hashlib.sha256(str(key).encode("utf-8")).digest()
        key_hash = int.from_bytes(digest[:4], "big")
        seed = (base_seed + key_hash) % (2**31)
        generators.append(torch.Generator(device=generator_device).manual_seed(seed))
    return generators


def calculate_zero_std_ratio(group_keys, gathered_rewards):
    group_array = np.array(group_keys)
    unique_groups, inverse_indices, counts = np.unique(
        group_array, return_inverse=True, return_counts=True
    )
    del unique_groups
    grouped_rewards = gathered_rewards["ori_avg"][np.argsort(inverse_indices)]
    split_indices = np.cumsum(counts)[:-1]
    reward_groups = np.split(grouped_rewards, split_indices)
    group_std = np.array([np.std(group) for group in reward_groups])
    zero_std_ratio = np.count_nonzero(group_std == 0) / len(group_std)
    return zero_std_ratio, group_std.mean()


def requires_grad(model, flag=True):
    for param in model.parameters():
        param.requires_grad = flag


def count_parameters(model):
    total = 0
    trainable = 0
    for param in model.parameters():
        numel = param.numel()
        total += numel
        if param.requires_grad:
            trainable += numel
    return trainable, total


def unwrap_model(model, accelerator):
    model = accelerator.unwrap_model(model)
    return getattr(model, "_orig_mod", model)


def disable_adapter(model, accelerator):
    return unwrap_model(model, accelerator).disable_adapter()


def resolve_model_root(model_name_or_path):
    if os.path.exists(model_name_or_path):
        return model_name_or_path
    cache_dir = os.getenv("HF_HUB_CACHE", "/tmp/flow_grpo_hf_cache/hub")
    os.makedirs(cache_dir, exist_ok=True)
    return snapshot_download(
        repo_id=model_name_or_path,
        cache_dir=cache_dir,
        ignore_patterns=["flax_model.msgpack", "rust_model.ot", "tf_model.h5"],
    )


def patch_omnigen_config_compat(model_root):
    config_path = os.path.join(model_root, "config.json")
    if not os.path.exists(config_path):
        return

    with open(config_path, "r", encoding="utf-8") as f:
        config = json.load(f)

    rope_scaling = config.get("rope_scaling")
    original_max_position_embeddings = config.get("original_max_position_embeddings")
    updated = False

    if isinstance(rope_scaling, dict):
        if rope_scaling.get("type") in {"su", "yarn"}:
            rope_scaling["type"] = "longrope"
            updated = True
        allowed_rope_keys = {"type", "short_factor", "long_factor"}
        extra_rope_keys = set(rope_scaling) - allowed_rope_keys
        for key in extra_rope_keys:
            rope_scaling.pop(key)
        updated = updated or bool(extra_rope_keys)

    if updated:
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
            f.write("\n")


def load_omnigen_model(config, device, weight_dtype, *, enable_gradient_checkpointing=False):
    model_root = resolve_model_root(config.pretrained.model)
    patch_omnigen_config_compat(model_root)
    model = OmniGen.from_pretrained(model_root)
    model.llm.config.use_cache = False
    if hasattr(model.llm, "gradient_checkpointing_enable") and enable_gradient_checkpointing:
        try:
            model.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
        except TypeError:
            model.llm.gradient_checkpointing_enable()
    model.to(device=device, dtype=weight_dtype)
    return model, model_root


def load_omnigen_components(config, device, weight_dtype):
    model, model_root = load_omnigen_model(
        config,
        device,
        weight_dtype,
        enable_gradient_checkpointing=config.activation_checkpointing,
    )
    processor = OmniGenProcessor.from_pretrained(model_root)

    vae_path = getattr(config.pretrained, "vae_path", None)
    if vae_path:
        vae = AutoencoderKL.from_pretrained(vae_path)
    else:
        local_vae_path = os.path.join(model_root, "vae")
        if os.path.exists(local_vae_path):
            vae = AutoencoderKL.from_pretrained(local_vae_path)
        else:
            vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")

    vae.to(device=device, dtype=torch.float32)
    return model, vae, processor


def merge_lora_into_base_model(model, lora_path, weight_dtype, *, trainable=True):
    logger.info("Merging OmniGen LoRA into base weights from %s", lora_path)
    merged_model = PeftModel.from_pretrained(
        model,
        lora_path,
        is_trainable=False,
    )
    if hasattr(merged_model, "set_adapter"):
        merged_model.set_adapter("default")
    if not hasattr(merged_model, "merge_and_unload"):
        raise AttributeError(
            "Loaded OmniGen LoRA does not support merge_and_unload(); "
            "cannot start full-parameter training from adapter weights."
        )
    model = merged_model.merge_and_unload()
    model.to(dtype=weight_dtype)
    requires_grad(model, trainable)
    return model


def load_lora_adapter_config(lora_path):
    adapter_config_path = os.path.join(lora_path, "adapter_config.json")
    try:
        with open(adapter_config_path, "r", encoding="utf-8") as handle:
            return json.load(handle)
    except OSError:
        logger.warning("Could not read LoRA adapter config from %s", adapter_config_path)
    except json.JSONDecodeError:
        logger.warning("Invalid LoRA adapter config JSON at %s", adapter_config_path)
    return None


def load_full_finetune_reference_model(config, device, weight_dtype):
    # Keep a frozen copy of the initial full-finetuning policy for KL regularization.
    reference_model, _ = load_omnigen_model(config, device, weight_dtype)
    if config.train.lora_path:
        reference_model = merge_lora_into_base_model(
            reference_model,
            config.train.lora_path,
            weight_dtype,
            trainable=False,
        )
    else:
        requires_grad(reference_model, False)
    reference_model.eval()
    return reference_model


def compute_log_prob(model, sample, timestep_index, config):
    current_latents = sample["latents"][:, timestep_index]
    current_timestep = sample["timesteps"][:, timestep_index]
    next_timestep = sample["next_timesteps"][:, timestep_index]
    if sample.get("use_joint_mask", False):
        current_mask_latents = sample["mask_latents"][:, timestep_index]
        model_output, _ = compute_cfg_model_output_joint(
            model,
            current_latents,
            current_mask_latents,
            current_timestep,
            sample["branch_inputs"],
            cfg_scale=config.sample.guidance_scale,
            img_cfg_scale=config.sample.img_guidance_scale,
            use_img_guidance=config.sample.use_img_guidance,
        )
    else:
        model_output = compute_cfg_model_output(
            model,
            current_latents,
            current_timestep,
            sample["branch_inputs"],
            cfg_scale=config.sample.guidance_scale,
            img_cfg_scale=config.sample.img_guidance_scale,
            use_img_guidance=config.sample.use_img_guidance,
        )
    stochastic_step_mask = sample.get("stochastic_step_mask")
    use_stochastic_step = True
    if stochastic_step_mask is not None:
        use_stochastic_step = bool(stochastic_step_mask[timestep_index].detach().cpu().item())
    return omnigen_sde_step_with_logprob(
        model_output,
        current_timestep,
        next_timestep,
        current_latents,
        prev_sample=sample["next_latents"][:, timestep_index],
        noise_level=config.sample.noise_level if use_stochastic_step else 0.0,
        sde_type=config.sample.sde_type if use_stochastic_step else "ode",
        logprob_mode=getattr(config.sample, "omnigen_logprob_mode", "mse"),
    )


def resolve_train_timestep_indices(config):
    num_steps = int(config.sample.num_steps)
    if bool(getattr(config.sample, "sde_window_enabled", False)) and bool(
        getattr(config.sample, "sde_window_train_only", True)
    ):
        mode = getattr(config.sample, "sde_window_mode", "fixed")
        if mode != "fixed":
            raise ValueError(f"Unsupported OmniGen SDE window mode: {mode!r}")
        start = int(getattr(config.sample, "sde_window_start", 0))
        end = int(getattr(config.sample, "sde_window_end", num_steps))
        if start < 0 or end < 0 or start >= end or end > num_steps:
            raise ValueError(
                "Invalid SDE window train range. Expected 0 <= start < end <= num_steps, "
                f"got start={start}, end={end}, num_steps={num_steps}."
            )
        return list(range(start, end))

    num_train_timesteps = max(1, int(num_steps * config.train.timestep_fraction))
    return list(range(num_train_timesteps))


def save_ckpt(save_dir, model, global_step, accelerator, ema, trainable_parameters, config):
    save_root = os.path.join(save_dir, "checkpoints", f"checkpoint-{global_step}")
    adapter_root = os.path.join(save_root, "lora")
    os.makedirs(adapter_root, exist_ok=True)

    if not accelerator.is_main_process:
        return

    if config.train.ema:
        ema.copy_ema_to(trainable_parameters, store_temp=True)

    unwrapped = unwrap_model(model, accelerator)
    if config.use_lora:
        unwrapped.save_pretrained(adapter_root)
    else:
        torch.save(unwrapped.state_dict(), os.path.join(save_root, "model.pt"))

    if config.train.ema:
        ema.copy_temp_to(trainable_parameters)


def eval(
    model,
    vae,
    processor,
    test_dataloader,
    config,
    accelerator,
    global_step,
    reward_fn,
    executor,
    autocast,
    weight_dtype,
    ema,
    trainable_parameters,
):
    eval_start_time = time.time()
    if config.train.ema:
        ema.copy_ema_to(trainable_parameters, store_temp=True)

    all_rewards = defaultdict(list)
    eval_gallery = []
    max_logged_eval_samples = 15
    max_eval_batches = getattr(config.sample, "max_eval_batches", None)
    eval_dataloader_len = len(test_dataloader)
    if max_eval_batches is None:
        eval_iterable = test_dataloader
        eval_total = eval_dataloader_len
    else:
        eval_iterable = itertools.islice(test_dataloader, max_eval_batches)
        eval_total = min(max_eval_batches, eval_dataloader_len)

    if accelerator.is_main_process:
        logger.info(
            "Starting train-time Eval: this is no-grad monitoring, not optimizer training. "
            "max_eval_batches=%s, eval_num_steps=%s, eval_noise_level=0.0, eval_batches=%s, eval_dataset_size=%s",
            max_eval_batches,
            config.sample.eval_num_steps,
            eval_total,
            len(getattr(test_dataloader, "dataset", [])),
        )

    for batch_index, test_batch in enumerate(
        tqdm(
            eval_iterable,
            desc="Eval",
            total=eval_total,
            disable=not accelerator.is_local_main_process,
            dynamic_ncols=True,
        )
    ):
        prompts, instructions, prompt_metadata, input_image_paths, ref_images, _ = test_batch
        with autocast():
            with torch.no_grad():
                eval_pipeline_fn = pipeline_with_logprob_joint_image_reward if bool(getattr(config, "use_joint_mask", False)) else pipeline_with_logprob
                collected = eval_pipeline_fn(
                    model,
                    vae,
                    processor,
                    instructions,
                    input_image_paths,
                    height=config.resolution,
                    width=config.resolution,
                    num_inference_steps=config.sample.eval_num_steps,
                    guidance_scale=config.sample.eval_guidance_scale,
                    img_guidance_scale=config.sample.eval_img_guidance_scale,
                    max_input_image_size=config.sample.max_input_image_size,
                    use_img_guidance=config.sample.use_img_guidance,
                    use_input_image_size_as_output=config.sample.use_input_image_size_as_output,
                    dtype=weight_dtype,
                    output_type="pt",
                    noise_level=0.0,
                    sde_type=config.sample.sde_type,
                    mask_scale_factor=getattr(getattr(config, "joint", {}), "mask_scale_factor", 1.0),
                    sde_window_enabled=False,
                )
                processor = collected["processor"]

        images = collected["images"]
        rewards_future = executor.submit(
            reward_fn,
            images,
            prompts,
            prompt_metadata,
            ref_images,
            only_strict=False,
        )
        time.sleep(0)
        rewards, _ = rewards_future.result()

        for key, value in rewards.items():
            rewards_gather = accelerator.gather(
                torch.as_tensor(value, device=accelerator.device).float()
            ).cpu().numpy()
            all_rewards[key].append(rewards_gather)

        local_images = images.float().cpu().numpy()
        for sample_index in range(len(local_images)):
            if len(eval_gallery) >= max_logged_eval_samples:
                break
            eval_gallery.append(
                {
                    "image": local_images[sample_index],
                    "prompt": prompts[sample_index],
                    "instruction": instructions[sample_index],
                    "metadata": prompt_metadata[sample_index],
                    "ref_image": ref_images[sample_index],
                    "rewards": {
                        key: float(np.asarray(value)[sample_index])
                        for key, value in rewards.items()
                    },
                }
            )

    if not all_rewards:
        if config.train.ema:
            ema.copy_temp_to(trainable_parameters)
        return processor

    all_rewards = {key: np.concatenate(value) for key, value in all_rewards.items()}
    if accelerator.is_main_process and eval_gallery:
        wandb.log(
            {
                "eval_images": _build_wandb_examples(
                    [item["image"] for item in eval_gallery],
                    [item["prompt"] for item in eval_gallery],
                    [item["metadata"] for item in eval_gallery],
                    [item["ref_image"] for item in eval_gallery],
                    instructions=[item["instruction"] for item in eval_gallery],
                    rewards={
                        key: np.array([item["rewards"].get(key, -10.0) for item in eval_gallery], dtype=np.float32)
                        for key in eval_gallery[0]["rewards"].keys()
                    },
                    max_samples=max_logged_eval_samples,
                    note="eval output",
                ),
                **{
                    f"eval_reward_{key}": np.mean(value[value != -10])
                    for key, value in all_rewards.items()
                },
            },
            step=global_step,
        )

    if config.train.ema:
        ema.copy_temp_to(trainable_parameters)
    if accelerator.is_main_process:
        logger.info("Train-time Eval finished in %.1fs", time.time() - eval_start_time)
    return processor


def main(_):
    config = FLAGS.config
    unique_id = datetime.datetime.now().strftime("%Y.%m.%d_%H.%M.%S")
    config.run_name = f"{config.run_name}_{unique_id}" if config.run_name else unique_id

    train_timestep_indices = resolve_train_timestep_indices(config)
    num_train_timesteps = len(train_timestep_indices)
    accelerator_config = ProjectConfiguration(
        project_dir=os.path.join(config.logdir, config.run_name),
        automatic_checkpoint_naming=True,
        total_limit=config.num_checkpoint_limit,
    )
    accelerator = Accelerator(
        mixed_precision=config.mixed_precision,
        project_config=accelerator_config,
        gradient_accumulation_steps=config.train.gradient_accumulation_steps * num_train_timesteps,
    )

    if accelerator.is_main_process:
        wandb.init(
            project=os.environ.get("WANDB_PROJECT", "flow_grpo"),
            name=os.environ.get("WANDB_NAME"),
            mode=os.environ.get("WANDB_MODE", "online"),
        )
        write_run_metadata(
            config.save_dir,
            config,
            config_entry=os.environ.get("CONFIG_ENTRY"),
            extra={
                "base_model": getattr(config.pretrained, "model", None),
                "omnigen_code_root": getattr(config.pretrained, "local_code_root", None),
                "sft_lora_path": getattr(getattr(config, "joint", {}), "sft_lora_path", None),
                "mask_modules_path": getattr(getattr(config, "joint", {}), "mask_modules_path", None),
                "reward_mask_channels": os.environ.get("REWARD_MASK_CHANNELS"),
                "rl_lora_target_modules": getattr(config.train, "rl_lora_target_modules", None),
                "train_jsonl": os.environ.get("TRAIN_JSONL"),
                "test_jsonl": os.environ.get("TEST_JSONL"),
                "save_dir": config.save_dir,
                "sde_window_enabled": bool(getattr(config.sample, "sde_window_enabled", False)),
                "sde_window_start": getattr(config.sample, "sde_window_start", None),
                "sde_window_end": getattr(config.sample, "sde_window_end", None),
                "sde_window_mode": getattr(config.sample, "sde_window_mode", None),
                "omnigen_logprob_mode": getattr(config.sample, "omnigen_logprob_mode", "mse"),
                "train_timestep_indices": train_timestep_indices,
            },
        )
    if accelerator.is_main_process:
        if _debug_flag("DEBUG_VERBOSE"):
            logger.info(f"\n{config}")
        else:
            logger.info(
                "Resolved joint config summary: max_eval_batches=%s, eval_num_steps=%s, num_steps=%s, "
                "noise_level=%s, train_batch_size=%s, test_batch_size=%s, num_image_per_prompt=%s, beta=%s, "
                "sde_window_enabled=%s, sde_window=[%s,%s), train_timesteps=%s, logprob_mode=%s",
                getattr(config.sample, "max_eval_batches", None),
                config.sample.eval_num_steps,
                config.sample.num_steps,
                config.sample.noise_level,
                getattr(config.sample, "train_batch_size", None),
                getattr(config.sample, "test_batch_size", None),
                config.sample.num_image_per_prompt,
                getattr(config.train, "beta", None),
                bool(getattr(config.sample, "sde_window_enabled", False)),
                getattr(config.sample, "sde_window_start", None),
                getattr(config.sample, "sde_window_end", None),
                train_timestep_indices,
                getattr(config.sample, "omnigen_logprob_mode", "mse"),
            )
    set_seed(config.seed, device_specific=True)

    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    model_load_start = time.time()
    use_joint_mask = bool(getattr(config, "use_joint_mask", False))
    joint_helpers = {}
    if use_joint_mask:
        model, vae, processor, joint_helpers = load_joint_omnigen_components_for_rl(
            config,
            device=accelerator.device,
            weight_dtype=weight_dtype,
            attach_rl_lora=True,
        )
    else:
        model, vae, processor = load_omnigen_components(config, accelerator.device, weight_dtype)
        requires_grad(vae, False)

        if config.use_lora:
            merge_lora_path = getattr(config.train, "merge_lora_path", None)
            if merge_lora_path:
                model = merge_lora_into_base_model(
                    model,
                    merge_lora_path,
                    weight_dtype,
                    trainable=False,
                )
            requires_grad(model, False)
            lora_rank = getattr(config.train, "lora_rank", 8)
            lora_alpha = getattr(config.train, "lora_alpha", lora_rank)
            lora_config = LoraConfig(
                r=lora_rank,
                lora_alpha=lora_alpha,
                init_lora_weights="gaussian",
                target_modules=getattr(config.train, "rl_lora_target_modules", ["qkv_proj", "o_proj"]),
            )
            model.llm.enable_input_require_grads()
            if config.train.lora_path:
                logger.info("Loading pretrained OmniGen LoRA from %s", config.train.lora_path)
                saved_lora_config = load_lora_adapter_config(config.train.lora_path)
                if saved_lora_config is not None:
                    saved_rank = saved_lora_config.get("r")
                    saved_alpha = saved_lora_config.get("lora_alpha")
                    if saved_rank != lora_rank or saved_alpha != lora_alpha:
                        logger.warning(
                            "Loaded LoRA adapter was saved with r=%s, alpha=%s; "
                            "requested r=%s, alpha=%s only applies when creating a fresh adapter.",
                            saved_rank,
                            saved_alpha,
                            lora_rank,
                            lora_alpha,
                        )
                model = PeftModel.from_pretrained(
                    model,
                    config.train.lora_path,
                    is_trainable=True,
                )
                model.set_adapter("default")
            else:
                model = get_peft_model(model, lora_config)
            model.to(dtype=weight_dtype)
        elif config.train.lora_path:
            model = merge_lora_into_base_model(model, config.train.lora_path, weight_dtype)

    if accelerator.is_main_process:
        logger.info("Model/component load finished in %.1fs", time.time() - model_load_start)
    trainable_parameters = [p for p in model.parameters() if p.requires_grad]
    trainable_param_count, total_param_count = count_parameters(model)
    ema = None
    if config.train.ema:
        ema = EMAModuleWrapper(
            trainable_parameters,
            decay=0.9,
            update_step_interval=8,
            device=accelerator.device,
        )

    if config.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True

    if config.train.use_8bit_adam:
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
            )

        optimizer_cls = bnb.optim.AdamW8bit
    else:
        optimizer_cls = torch.optim.AdamW

    optimizer = optimizer_cls(
        trainable_parameters,
        lr=config.train.learning_rate,
        betas=(config.train.adam_beta1, config.train.adam_beta2),
        weight_decay=config.train.adam_weight_decay,
        eps=config.train.adam_epsilon,
    )

    train_dataset = RadiomicsEditDataset(
        config.dataset,
        "train",
        condition_dropout_prob=getattr(config.train, "condition_dropout_prob", 0.0),
    )
    test_dataset = RadiomicsEditDataset(config.dataset, "test")
    train_sampler = DistributedKRepeatSampler(
        dataset=train_dataset,
        batch_size=config.sample.train_batch_size,
        k=config.sample.num_image_per_prompt,
        num_replicas=accelerator.num_processes,
        rank=accelerator.process_index,
        num_batches_per_epoch=config.sample.num_batches_per_epoch,
        seed=42,
    )
    if config.sample.same_latent and train_sampler.group_span > 1:
        raise ValueError(
            "same_latent=True currently requires each prompt group to fit inside one global sampling batch. "
            f"Got group_span={train_sampler.group_span}, k={config.sample.num_image_per_prompt}, "
            f"num_replicas={accelerator.num_processes}, train_batch_size={config.sample.train_batch_size}."
        )

    train_dataloader = DataLoader(
        train_dataset,
        batch_sampler=train_sampler,
        num_workers=0,
        collate_fn=RadiomicsEditDataset.collate_fn,
    )
    test_dataloader = DataLoader(
        test_dataset,
        batch_size=config.sample.test_batch_size,
        shuffle=False,
        num_workers=2,
        collate_fn=RadiomicsEditDataset.collate_fn,
    )

    if config.sample.num_image_per_prompt == 1:
        config.per_prompt_stat_tracking = False
    if config.per_prompt_stat_tracking:
        stat_tracker = PerPromptStatTracker(config.sample.global_std)

    autocast = contextlib.nullcontext if config.use_lora else accelerator.autocast
    reward_fn = getattr(flow_grpo.rewards, "multi_score")(accelerator.device, config.reward_fn)
    eval_reward_fn = getattr(flow_grpo.rewards, "multi_score")(accelerator.device, config.reward_fn)

    model, optimizer, train_dataloader, test_dataloader = accelerator.prepare(
        model, optimizer, train_dataloader, test_dataloader
    )
    reference_model = None
    kl_reference_mode = "disabled"
    if config.train.beta > 0:
        if config.use_lora:
            kl_reference_mode = "base_model_without_adapter"
        else:
            logger.info("Loading frozen OmniGen reference model for full-finetune KL regularization.")
            reference_model = load_full_finetune_reference_model(
                config,
                accelerator.device,
                weight_dtype,
            )
            kl_reference_mode = "frozen_initial_policy"
    trainable_parameters = [p for p in model.parameters() if p.requires_grad]
    if config.train.ema:
        ema = EMAModuleWrapper(
            trainable_parameters,
            decay=0.9,
            update_step_interval=8,
            device=accelerator.device,
        )

    executor = futures.ThreadPoolExecutor(max_workers=8)
    samples_per_epoch = (
        config.sample.train_batch_size
        * accelerator.num_processes
        * config.sample.num_batches_per_epoch
    )
    total_train_batch_size = (
        config.train.batch_size
        * accelerator.num_processes
        * config.train.gradient_accumulation_steps
    )
    logger.info("***** Running OmniGen GRPO training *****")
    logger.info(f"  Finetuning mode = {'LoRA' if config.use_lora else 'full'}")
    if config.use_lora:
        logger.info(
            "  LoRA rank / alpha = %s / %s",
            getattr(config.train, "lora_rank", 8),
            getattr(config.train, "lora_alpha", getattr(config.train, "lora_rank", 8)),
        )
        logger.info("  Merged base LoRA path = %s", getattr(config.train, "merge_lora_path", None))
        logger.info("  Trainable LoRA path = %s", getattr(config.train, "lora_path", None))
    logger.info("  KL reference mode = %s", kl_reference_mode)
    trainable_tensor_count = sum(1 for p in model.parameters() if p.requires_grad)
    logger.info(
        "  Trainable parameters = %s / %s (%.2f%%); trainable tensors=%s; LoRA targets=%s",
        f"{trainable_param_count:,}",
        f"{total_param_count:,}",
        100.0 * trainable_param_count / max(total_param_count, 1),
        trainable_tensor_count,
        getattr(config.train, "rl_lora_target_modules", None),
    )
    if _debug_flag("DEBUG_MODEL_PARAMS") and accelerator.is_main_process:
        logger.info("  Trainable parameter names = %s", [name for name, param in model.named_parameters() if param.requires_grad])
    logger.info(f"  Sample batch size per device = {config.sample.train_batch_size}")
    logger.info(f"  Train batch size per device = {config.train.batch_size}")
    logger.info(f"  Gradient accumulation steps = {config.train.gradient_accumulation_steps}")
    logger.info(f"  Samples per epoch = {samples_per_epoch}")
    logger.info(f"  Effective train batch size = {total_train_batch_size}")

    epoch = 0
    global_step = 0
    train_iter = iter(train_dataloader)
    max_train_steps = int(getattr(config.train, "max_train_steps", 0) or 0)

    while epoch < config.num_epochs:
        if max_train_steps > 0 and global_step >= max_train_steps:
            logger.info("Reached max_train_steps=%s at global_step=%s.", max_train_steps, global_step)
            break
        train_sampler.set_epoch(epoch)
        model.eval()
        # Disabled per request while we focus on train rollouts and throughput.
        if epoch % config.eval_freq == 0:
            processor = eval(
                model,
                vae,
                processor,
                test_dataloader,
                config,
                accelerator,
                global_step,
                eval_reward_fn,
                executor,
                autocast,
                weight_dtype,
                ema,
                trainable_parameters,
            )
        if epoch % config.save_freq == 0 and epoch > 0:
            save_ckpt(config.save_dir, model, global_step, accelerator, ema, trainable_parameters, config)

        if accelerator.is_main_process:
            logger.info("Starting no-grad rollout sampling for epoch %s", epoch)
        rollout_start_time = time.time()
        reward_wait_start_time = None
        update_start_time = None
        samples = []
        last_images = None
        last_prompts = None
        last_instructions = None
        last_rewards = None
        last_ref_images = None
        last_metadata = None
        last_input_image_paths = None

        for batch_index in tqdm(
            range(config.sample.num_batches_per_epoch),
            desc=f"Epoch {epoch}: sampling",
            disable=not accelerator.is_local_main_process,
            dynamic_ncols=True,
        ):
            prompts, instructions, prompt_metadata, input_image_paths, ref_images, group_keys = next(train_iter)
            group_ids = build_group_ids(group_keys, device=accelerator.device)
            generator = (
                create_generator(
                    group_keys,
                    base_seed=epoch * 10000 + batch_index,
                    device=accelerator.device,
                )
                if config.sample.same_latent
                else None
            )

            with autocast():
                with torch.no_grad():
                    pipeline_fn = pipeline_with_logprob_joint_image_reward if use_joint_mask else pipeline_with_logprob
                    collected = pipeline_fn(
                        model,
                        vae,
                        processor,
                        instructions,
                        input_image_paths,
                        height=config.resolution,
                        width=config.resolution,
                        num_inference_steps=config.sample.num_steps,
                        guidance_scale=config.sample.guidance_scale,
                        img_guidance_scale=config.sample.img_guidance_scale,
                        max_input_image_size=config.sample.max_input_image_size,
                        use_img_guidance=config.sample.use_img_guidance,
                        use_input_image_size_as_output=config.sample.use_input_image_size_as_output,
                        dtype=weight_dtype,
                        generator=generator,
                        output_type="pt",
                        noise_level=config.sample.noise_level,
                        sde_type=config.sample.sde_type,
                        mask_scale_factor=getattr(getattr(config, "joint", {}), "mask_scale_factor", 1.0),
                        sde_window_enabled=bool(getattr(config.sample, "sde_window_enabled", False)),
                        sde_window_start=int(getattr(config.sample, "sde_window_start", 0)),
                        sde_window_end=int(getattr(config.sample, "sde_window_end", config.sample.num_steps)),
                        sde_window_mode=getattr(config.sample, "sde_window_mode", "fixed"),
                        logprob_mode=getattr(config.sample, "omnigen_logprob_mode", "mse"),
                    )
                    processor = collected["processor"]

            latents = torch.stack(collected["all_latents"], dim=1)
            mask_latents = (
                torch.stack(collected["all_mask_latents"], dim=1)
                if use_joint_mask and "all_mask_latents" in collected
                else None
            )
            log_probs = torch.stack(collected["all_log_probs"], dim=1)
            timesteps = torch.stack(collected["all_timesteps"], dim=1)
            next_timesteps = torch.stack(collected["all_next_timesteps"], dim=1)
            std_dev_ts = torch.stack(collected["all_std_dev_ts"], dim=1)
            stochastic_step_mask = collected["stochastic_step_mask"].to(accelerator.device)
            images = collected["images"]
            rewards_future = executor.submit(
                reward_fn,
                images,
                prompts,
                prompt_metadata,
                ref_images,
                only_strict=True,
            )
            time.sleep(0)

            samples.append(
                {
                    "group_ids": group_ids,
                    "branch_inputs": collected["branch_inputs"],
                    "timesteps": timesteps,
                    "next_timesteps": next_timesteps,
                    "latents": latents[:, :-1],
                    "next_latents": latents[:, 1:],
                    "mask_latents": mask_latents[:, :-1] if mask_latents is not None else None,
                    "next_mask_latents": mask_latents[:, 1:] if mask_latents is not None else None,
                    "use_joint_mask": use_joint_mask,
                    "log_probs": log_probs,
                    "std_dev_ts": std_dev_ts,
                    "stochastic_step_mask": stochastic_step_mask,
                    "rewards": rewards_future,
                }
            )
            last_images = images
            last_prompts = list(prompts)
            last_instructions = list(instructions)
            last_ref_images = list(ref_images)
            last_metadata = list(prompt_metadata)
            last_input_image_paths = [list(paths) for paths in input_image_paths]

        if accelerator.is_main_process:
            logger.info("No-grad rollout sampling finished in %.1fs; waiting for radiomics rewards", time.time() - rollout_start_time)
        reward_wait_start_time = time.time()
        for sample in tqdm(
            samples,
            desc="Waiting for rewards",
            disable=not accelerator.is_local_main_process,
            dynamic_ncols=True,
        ):
            rewards, _ = sample["rewards"].result()
            sample["rewards"] = {
                key: torch.as_tensor(value, device=accelerator.device).float()
                for key, value in rewards.items()
            }
            last_rewards = sample["rewards"]

        if accelerator.is_main_process:
            logger.info("Radiomics reward wait finished in %.1fs", time.time() - reward_wait_start_time)
        reward_tensors = {
            key: torch.cat([sample["rewards"][key] for sample in samples], dim=0)
            for key in samples[0]["rewards"].keys()
        }
        reward_tensors["ori_avg"] = reward_tensors["avg"]
        reward_tensors["avg"] = reward_tensors["avg"].unsqueeze(1).repeat(1, config.sample.num_steps)
        gathered_rewards = {
            key: accelerator.gather(value).float().cpu().numpy()
            for key, value in reward_tensors.items()
        }
        gathered_group_ids = accelerator.gather(
            torch.cat([sample["group_ids"] for sample in samples], dim=0)
        ).cpu().numpy()

        if epoch % 10 == 0 and last_images is not None:
            preview_images = None
            with autocast():
                with torch.no_grad():
                    preview_pipeline_fn = pipeline_with_logprob_joint_image_reward if use_joint_mask else pipeline_with_logprob
                    preview_collected = preview_pipeline_fn(
                        model,
                        vae,
                        processor,
                        last_instructions,
                        last_input_image_paths,
                        height=config.resolution,
                        width=config.resolution,
                        num_inference_steps=config.sample.eval_num_steps,
                        guidance_scale=config.sample.eval_guidance_scale,
                        img_guidance_scale=config.sample.eval_img_guidance_scale,
                        max_input_image_size=config.sample.max_input_image_size,
                        use_img_guidance=config.sample.use_img_guidance,
                        use_input_image_size_as_output=config.sample.use_input_image_size_as_output,
                        dtype=weight_dtype,
                        output_type="pt",
                        noise_level=0.0,
                        sde_type=config.sample.sde_type,
                        mask_scale_factor=getattr(getattr(config, "joint", {}), "mask_scale_factor", 1.0),
                        sde_window_enabled=False,
                    )
                    if accelerator.is_main_process:
                        preview_images = preview_collected["images"].float().cpu().numpy()

            if accelerator.is_main_process:
                reward_arrays = {
                    key: value.detach().float().cpu().numpy()
                    for key, value in last_rewards.items()
                }
                wandb.log(
                    {
                        "images": _build_wandb_examples(
                            preview_images,
                            last_prompts,
                            last_metadata,
                            last_ref_images,
                            instructions=last_instructions,
                            max_samples=8,
                            note="deterministic preview (noise_level=0.0)",
                        ),
                        "train_rollout_images": _build_wandb_examples(
                            last_images.float().cpu().numpy(),
                            last_prompts,
                            last_metadata,
                            last_ref_images,
                            instructions=last_instructions,
                            rewards=reward_arrays,
                            max_samples=8,
                            note=f"stochastic rollout (noise_level={config.sample.noise_level})",
                        ),
                    },
                    step=global_step,
                )
            accelerator.wait_for_everyone()

        reward_summary = None
        if accelerator.is_main_process:
            reward_summary = {}
            extra_logs = {}
            if "radiomics_missing" in gathered_rewards:
                extra_logs["reward/radiomics_missing_score_rate"] = float(np.mean(gathered_rewards["radiomics_missing"]))
            if "radiomics_roi_pixels" in gathered_rewards:
                extra_logs["reward/roi_pixels_mean"] = float(np.mean(gathered_rewards["radiomics_roi_pixels"]))
            wandb.log(
                {
                    "epoch": epoch,
                    "config/use_joint_mask": float(bool(getattr(config, "use_joint_mask", False))),
                    "config/sde_window_enabled": float(bool(getattr(config.sample, "sde_window_enabled", False))),
                    "config/sde_window_start": float(getattr(config.sample, "sde_window_start", 0)),
                    "config/sde_window_end": float(getattr(config.sample, "sde_window_end", config.sample.num_steps)),
                    "config/sde_window_count": float(num_train_timesteps),
                    "config/omnigen_logprob_mode": getattr(config.sample, "omnigen_logprob_mode", "mse"),
                    **extra_logs,
                    **{
                        f"reward_{key}": value.mean()
                        for key, value in gathered_rewards.items()
                        if "_strict_accuracy" not in key and "_accuracy" not in key
                    },
                },
                step=global_step,
            )
            for key, value in gathered_rewards.items():
                if key == "ori_avg" or "_strict_accuracy" in key or "_accuracy" in key:
                    continue
                valid = value != -10
                reward_summary[key] = float(np.mean(value[valid])) if np.any(valid) else float(np.mean(value))

        if config.per_prompt_stat_tracking:
            group_key_strings = [str(group_id) for group_id in gathered_group_ids.tolist()]
            advantages = stat_tracker.update(group_key_strings, gathered_rewards["avg"])
            group_size, trained_prompt_num = stat_tracker.get_stats()
            zero_std_ratio, reward_std_mean = calculate_zero_std_ratio(
                group_key_strings,
                gathered_rewards,
            )
            if accelerator.is_main_process:
                wandb.log(
                    {
                        "group_size": group_size,
                        "trained_prompt_num": trained_prompt_num,
                        "zero_std_ratio": zero_std_ratio,
                        "reward_std_mean": reward_std_mean,
                    },
                    step=global_step,
                )
            stat_tracker.clear()
        else:
            group_size = None
            trained_prompt_num = None
            zero_std_ratio = None
            reward_std_mean = None
            advantages = (
                gathered_rewards["avg"] - gathered_rewards["avg"].mean()
            ) / (gathered_rewards["avg"].std() + 1e-4)

        if accelerator.is_main_process and reward_summary is not None:
            summary_parts = [f"{key}={value:.4f}" for key, value in sorted(reward_summary.items())]
            if zero_std_ratio is not None:
                summary_parts.append(f"zero_std_ratio={zero_std_ratio:.4f}")
            if reward_std_mean is not None:
                summary_parts.append(f"reward_std_mean={reward_std_mean:.4f}")
            if group_size is not None:
                summary_parts.append(f"group_size={group_size:.2f}")
            if trained_prompt_num is not None:
                summary_parts.append(f"trained_prompt_num={trained_prompt_num}")
            logger.info("Epoch %s reward summary: %s", epoch, ", ".join(summary_parts))

        local_advantages = (
            torch.as_tensor(advantages)
            .reshape(accelerator.num_processes, -1, advantages.shape[-1])[accelerator.process_index]
            .to(accelerator.device)
        )
        offset = 0
        for sample in samples:
            batch_size = sample["latents"].shape[0]
            sample["advantages"] = local_advantages[offset : offset + batch_size]
            offset += batch_size

            sample["train_mask"] = sample["advantages"].abs().sum(dim=1) != 0
            del sample["rewards"]

        if len(samples) == 0:
            epoch += 1
            continue

        if accelerator.is_main_process:
            logger.info("Starting RL training updates")
        update_start_time = time.time()
        for inner_epoch in range(config.train.num_inner_epochs):
            random.shuffle(samples)
            model.train()
            info = defaultdict(list)

            for sample in tqdm(
                samples,
                desc=f"Epoch {epoch}.{inner_epoch}: training",
                disable=not accelerator.is_local_main_process,
                dynamic_ncols=True,
            ):
                for timestep_index in tqdm(
                    train_timestep_indices,
                    desc="Timestep",
                    leave=False,
                    disable=not accelerator.is_local_main_process,
                    dynamic_ncols=True,
                ):
                    with accelerator.accumulate(model):
                        with autocast():
                            _, log_prob, prev_sample_mean, std_dev_t = compute_log_prob(
                                model, sample, timestep_index, config
                            )
                            if config.train.beta > 0:
                                with torch.no_grad():
                                    if reference_model is None:
                                        with disable_adapter(model, accelerator):
                                            _, _, prev_sample_mean_ref, _ = compute_log_prob(
                                                model, sample, timestep_index, config
                                            )
                                    else:
                                        _, _, prev_sample_mean_ref, _ = compute_log_prob(
                                            reference_model, sample, timestep_index, config
                                        )

                        advantages_t = torch.clamp(
                            sample["advantages"][:, timestep_index],
                            -config.train.adv_clip_max,
                            config.train.adv_clip_max,
                        ).to(device=log_prob.device, dtype=log_prob.dtype)
                        train_mask = sample["train_mask"].to(device=log_prob.device, dtype=log_prob.dtype)
                        train_weight_sum = train_mask.sum().clamp_min(1.0)
                        ratio = torch.exp(log_prob - sample["log_probs"][:, timestep_index])
                        log_prob_delta = log_prob - sample["log_probs"][:, timestep_index]
                        unclipped_loss = -advantages_t * ratio
                        clipped_loss = -advantages_t * torch.clamp(
                            ratio,
                            1.0 - config.train.clip_range,
                            1.0 + config.train.clip_range,
                        )
                        policy_loss = (
                            torch.maximum(unclipped_loss, clipped_loss) * train_mask
                        ).sum() / train_weight_sum
                        if config.train.beta > 0:
                            kl_per_sample = (
                                ((prev_sample_mean - prev_sample_mean_ref) ** 2)
                                .mean(dim=tuple(range(1, prev_sample_mean.ndim)), keepdim=False)
                                / (2 * (std_dev_t.view(std_dev_t.shape[0], -1).mean(dim=1) ** 2 + 1e-6))
                            )
                            kl_loss = (kl_per_sample * train_mask).sum() / train_weight_sum
                            loss = policy_loss + config.train.beta * kl_loss
                        else:
                            kl_loss = None
                            loss = policy_loss

                        info["approx_kl"].append(
                            0.5
                            * (
                                ((log_prob - sample["log_probs"][:, timestep_index]) ** 2) * train_mask
                            ).sum()
                            / train_weight_sum
                        )
                        info["clipfrac"].append(
                            (
                                (torch.abs(ratio - 1.0) > config.train.clip_range).float()
                                * train_mask
                            ).sum()
                            / train_weight_sum
                        )
                        info["clipfrac_gt_one"].append(
                            ((ratio - 1.0 > config.train.clip_range).float() * train_mask).sum()
                            / train_weight_sum
                        )
                        info["clipfrac_lt_one"].append(
                            ((1.0 - ratio > config.train.clip_range).float() * train_mask).sum()
                            / train_weight_sum
                        )
                        info["ratio_mean"].append((ratio * train_mask).sum() / train_weight_sum)
                        info["ratio_std"].append(
                            torch.sqrt(
                                ((((ratio - ((ratio * train_mask).sum() / train_weight_sum)) ** 2) * train_mask).sum() / train_weight_sum)
                                + 1e-12
                            )
                        )
                        info["log_prob_delta_mean"].append((log_prob_delta * train_mask).sum() / train_weight_sum)
                        info["std_dev_t_mean"].append(
                            (std_dev_t.view(std_dev_t.shape[0], -1).mean(dim=1) * train_mask).sum() / train_weight_sum
                        )
                        info[f"per_step/approx_kl_step_{timestep_index}"].append(info["approx_kl"][-1])
                        info[f"per_step/clipfrac_step_{timestep_index}"].append(info["clipfrac"][-1])
                        info[f"per_step/ratio_mean_step_{timestep_index}"].append(info["ratio_mean"][-1])
                        info[f"per_step/std_dev_t_step_{timestep_index}"].append(info["std_dev_t_mean"][-1])
                        info["policy_loss"].append(policy_loss)
                        info["loss"].append(loss)
                        if kl_loss is not None:
                            info["kl_loss"].append(kl_loss)

                        accelerator.backward(loss)
                        if accelerator.sync_gradients:
                            accelerator.clip_grad_norm_(model.parameters(), config.train.max_grad_norm)
                        optimizer.step()
                        optimizer.zero_grad()

                    if accelerator.sync_gradients:
                        reduced = {key: torch.mean(torch.stack(value)) for key, value in info.items()}
                        reduced = accelerator.reduce(reduced, reduction="mean")
                        reduced.update({"epoch": epoch, "inner_epoch": inner_epoch})
                        if accelerator.is_main_process:
                            wandb.log(reduced, step=global_step)
                        global_step += 1
                        if max_train_steps > 0 and global_step >= max_train_steps:
                            break
                        info = defaultdict(list)

                if config.train.ema:
                    ema.step(trainable_parameters, global_step)
                if max_train_steps > 0 and global_step >= max_train_steps:
                    break
            if max_train_steps > 0 and global_step >= max_train_steps:
                break

        if accelerator.is_main_process:
            logger.info("RL update phase finished in %.1fs", time.time() - update_start_time)
        epoch += 1

    logger.info("Training finished after %s epochs.", epoch)


if __name__ == "__main__":
    app.run(main)