Restore public repo before mistaken code upload

Restores files to revision 1a8e9e3ad2130f55c880af8ace85b9af0d0c329f and removes files mistakenly uploaded from unirl_opensource.

.gitignore

@@ -145,11 +145,5 @@ outputs/
 wandb/
 
 assets/large_rl_datasets/
-assets/rl_datasets/*.parquet
-assets/rl_datasets/*.jsonl
-!assets/rl_datasets/README.md
 
 utils/parquet_cache/
-
-rewards_services/api_services/editreward_scorer_service/.venv/
-rewards_services/api_services/editreward_scorer_service/EditReward/

README.md

@@ -1,17 +1,3 @@
----
-license: apache-2.0
-library_name: diffusers
-tags:
-  - reinforcement-learning
-  - image-generation
-  - image-editing
-  - prompt-optimization
-  - flux
-  - qwen
-datasets:
-  - wangfuyun/PrompRL
----
-
 <p align="center">
   <img src="assets/logo.png" width="30%"><br>
   PromptRL
@@ -41,109 +27,6 @@ pip install flash-attn==2.7.4.post1 --no-build-isolation
 # bash gen.sh
 ```
 
-<details>
-<summary><b>Training The Edit Model And Running EditReward</b></summary>
-
-<br>
-
-**Scope**
-
-This release keeps only the edit RL path. The trainer is `unirl/trainer/edit_grpo_trainer.py`, which jointly optimizes the Qwen-VL prompt refiner and the FLUX.1-Kontext transformer. The VAE, text encoders, and vision encoder stay frozen.
-
-The partial-refinement setting is preserved: with the default `NUM_GENERATIONS=8` and `NUM_SKIP_REFINEMENT=2`, each source image produces six edits from Qwen-refined prompts and two edits from the original prompt.
-
-Relevant files:
-
-- `unirl/train_edit.py`: CLI entry point for Qwen-Kontext edit GRPO.
-- `unirl/reward_evaluator/reward_evaluator.py`: EditReward HTTP client used by training.
-- `rewards_services/api_services/editreward_scorer_service`: EditReward service wrapper.
-- `scripts/train/edit_grpo.sh`: launch script with environment-variable configuration.
-
-**Dataset**
-
-By default, `scripts/train/edit_grpo.sh` loads:
-
-```text
-https://huggingface.co/wangfuyun/PrompRL/resolve/main/data/omni_edit_train_50k.parquet
-```
-
-You can override it with `PROMPTS_FILE`. The loader also accepts the Hugging Face web URL form with `/blob/main/`; it is converted to the downloadable `/resolve/main/` URL automatically.
-
-The dataset should be a `.parquet` or `.jsonl` file with:
-
-| Column | Description |
-| --- | --- |
-| `image` | Source image before editing. For jsonl this can be an image path. |
-| `prompt` | Edit instruction. |
-
-Optional columns are `caption` and `target_caption`. For other column names, set `IMAGE_COLUMN` and `PROMPT_COLUMN`.
-
-**1. Start EditReward**
-
-```bash
-cd rewards_services/api_services/editreward_scorer_service
-python -m venv .venv
-source .venv/bin/activate
-pip install --upgrade pip
-pip install torch torchvision torchaudio
-pip install -r requirements.txt
-
-git clone https://github.com/TIGER-AI-Lab/EditReward.git
-huggingface-cli download TIGER-Lab/EditReward-MiMo-VL-7B-SFT-2508 \
-  --local-dir EditReward/EditReward-MiMo-VL-7B-SFT-2508
-
-export EDITREWARD_CUDA_DEVICES=0,1
-export EDITREWARD_WORKERS=2
-export EDITREWARD_PORT=18088
-bash run.sh
-```
-
-If the EditReward repo or checkpoint is stored elsewhere:
-
-```bash
-export EDITREWARD_REPO_DIR=/path/to/EditReward
-export EDITREWARD_CHECKPOINT_PATH=/path/to/EditReward-MiMo-VL-7B-SFT-2508
-```
-
-**2. Launch Training**
-
-From the repository root:
-
-```bash
-export MODEL_NAME_OR_PATH=/path/to/qwenkontext/checkpoint
-# Optional. Defaults to the PromptRL OmniEdit 50k parquet on Hugging Face.
-export PROMPTS_FILE=https://huggingface.co/wangfuyun/PrompRL/blob/main/data/omni_edit_train_50k.parquet
-export EDITREWARD_URL=http://127.0.0.1:18088/
-export CUDA_VISIBLE_DEVICES=2,3,4,5,6,7
-export NPROC_PER_NODE=6
-export RUN_NAME=qwenkontext-editreward
-
-bash scripts/train/edit_grpo.sh
-```
-
-Common options:
-
-```bash
-export NUM_GENERATIONS=8
-export NUM_SKIP_REFINEMENT=2
-export NUM_SDE=4
-export PER_DEVICE_TRAIN_BATCH_SIZE=1
-export DIT_LEARNING_RATE=2e-7
-export LLM_LEARNING_RATE=3e-7
-export BETA=1e-2
-export IMAGE_COLUMN=image
-export PROMPT_COLUMN=prompt
-export REPORT_TO=wandb
-```
-
-Training logs sample source/edited images under:
-
-```text
-outputs/rl/kontext/$RUN_NAME/training_samples/
-```
-
-</details>
-
 ## Qualitative Results
 
 ### Text-to-Image Generation

assets/rl_datasets/README.md

@@ -1,23 +0,0 @@
-# Edit Training Dataset Schema
-
-The training script defaults to:
-
-```text
-https://huggingface.co/wangfuyun/PrompRL/resolve/main/data/omni_edit_train_50k.parquet
-```
-
-Use a `.parquet` or `.jsonl` file with at least:
-
-| Column | Type | Description |
-| --- | --- | --- |
-| `image` | PIL image, image bytes, or image path | Source image before editing. |
-| `prompt` | string | Edit instruction used by FLUX.1-Kontext and EditReward. |
-
-Optional columns:
-
-| Column | Type | Description |
-| --- | --- | --- |
-| `caption` | string | Source-image caption, kept for logging or downstream reward extensions. |
-| `target_caption` | string | Target edited-image caption, kept for logging or downstream reward extensions. |
-
-If your dataset uses different column names, pass `IMAGE_COLUMN=...` and `PROMPT_COLUMN=...` to `scripts/train/edit_grpo.sh`.

rewards_services/api_services/editreward_scorer_service/README.md

@@ -1,35 +0,0 @@
-# EditReward Scorer Service
-
-This service exposes EditReward over HTTP for edit GRPO training. It accepts a pickled payload with source images, edited images, and edit instructions, then returns `{"scores": [...]}`.
-
-## Setup
-
-```bash
-cd rewards_services/api_services/editreward_scorer_service
-python -m venv .venv
-source .venv/bin/activate
-pip install --upgrade pip
-pip install torch torchvision torchaudio
-pip install -r requirements.txt
-pip install flash-attn --no-build-isolation  # optional, recommended when your CUDA/PyTorch build supports it
-
-git clone https://github.com/TIGER-AI-Lab/EditReward.git
-huggingface-cli download TIGER-Lab/EditReward-MiMo-VL-7B-SFT-2508 \
-  --local-dir EditReward/EditReward-MiMo-VL-7B-SFT-2508
-```
-
-If the repository or checkpoint lives elsewhere, set:
-
-```bash
-export EDITREWARD_REPO_DIR=/path/to/EditReward
-export EDITREWARD_CHECKPOINT_PATH=/path/to/EditReward-MiMo-VL-7B-SFT-2508
-```
-
-## Run
-
-```bash
-export EDITREWARD_PORT=18088
-export EDITREWARD_CUDA_DEVICES=0,1
-export EDITREWARD_WORKERS=2
-bash run.sh
-```

rewards_services/api_services/editreward_scorer_service/app.py

@@ -1,94 +0,0 @@
-import os
-import pickle
-import traceback
-from io import BytesIO
-from typing import Any, Dict, List
-
-import torch
-from flask import Blueprint, Flask, request
-from PIL import Image
-
-from editreward_scorer import EditRewardScorer
-
-
-INFERENCE_FN = None
-root = Blueprint("root", __name__)
-
-
-def _deserialize_images(images_bytes: List[bytes]) -> List[Image.Image]:
-    return [Image.open(BytesIO(data)).convert("RGB") for data in images_bytes]
-
-
-def _service_config() -> Dict[str, Any]:
-    repo_dir = os.getenv("EDITREWARD_REPO_DIR", os.path.join(os.path.dirname(__file__), "EditReward"))
-    return {
-        "repo_dir": repo_dir,
-        "config_path": os.getenv(
-            "EDITREWARD_CONFIG_PATH",
-            os.path.join(repo_dir, "EditReward", "config", "EditReward-MiMo-VL-7B-SFT-2508.yaml"),
-        ),
-        "checkpoint_path": os.getenv(
-            "EDITREWARD_CHECKPOINT_PATH",
-            os.path.join(repo_dir, "EditReward-MiMo-VL-7B-SFT-2508"),
-        ),
-        "reward_dim": os.getenv("EDITREWARD_DIM", "overall_detail"),
-        "rm_head_type": os.getenv("EDITREWARD_HEAD_TYPE", "ranknet_multi_head"),
-    }
-
-
-def create_app():
-    global INFERENCE_FN
-    config = _service_config()
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Loading EditReward scorer on {device} from {config['checkpoint_path']}...")
-    INFERENCE_FN = EditRewardScorer(
-        repo_dir=config["repo_dir"],
-        config_path=config["config_path"],
-        checkpoint_path=config["checkpoint_path"],
-        reward_dim=config["reward_dim"],
-        rm_head_type=config["rm_head_type"],
-        device=device,
-    )
-    INFERENCE_FN.eval()
-    print("EditReward scorer loaded.")
-
-    app = Flask(__name__)
-    app.register_blueprint(root)
-    return app
-
-
-@root.route("/", methods=["GET"])
-def healthcheck():
-    return {"status": "ok", "service": "editreward"}, 200
-
-
-@root.route("/", methods=["POST"])
-def inference():
-    try:
-        payload = pickle.loads(request.get_data())
-        images = payload["images"]
-        prompts = payload.get("prompts", [])
-        source_images = _deserialize_images(images.get("source", []))
-        edited_images = _deserialize_images(images.get("edited", []))
-
-        if len(source_images) != len(edited_images) or len(source_images) != len(prompts):
-            raise ValueError(
-                "Mismatched EditReward inputs: "
-                f"{len(source_images)} source images, {len(edited_images)} edited images, {len(prompts)} prompts."
-            )
-
-        with torch.no_grad():
-            scores = INFERENCE_FN(prompts, source_images, edited_images)
-        return pickle.dumps({"scores": [float(score) for score in scores]}), 200
-    except Exception:
-        error_message = traceback.format_exc()
-        print(f"EditReward service error:\n{error_message}")
-        return pickle.dumps({"error": error_message}), 500
-
-
-if __name__ == "__main__":
-    port = int(os.getenv("EDITREWARD_PORT", "18088"))
-    host = os.getenv("EDITREWARD_HOST", "127.0.0.1")
-    app = create_app()
-    app.run(host=host, port=port, debug=False)
-

rewards_services/api_services/editreward_scorer_service/editreward_scorer.py

@@ -1,65 +0,0 @@
-import os
-import shutil
-import sys
-import tempfile
-from typing import List
-
-import torch
-from PIL import Image
-
-
-class EditRewardScorer(torch.nn.Module):
-    def __init__(
-        self,
-        repo_dir: str,
-        config_path: str,
-        checkpoint_path: str,
-        reward_dim: str = "overall_detail",
-        rm_head_type: str = "ranknet_multi_head",
-        device: str = "cuda",
-    ):
-        super().__init__()
-        if not os.path.isdir(repo_dir):
-            raise FileNotFoundError(
-                f"EditReward repository not found at {repo_dir}. "
-                "Clone https://github.com/TIGER-AI-Lab/EditReward.git or set EDITREWARD_REPO_DIR."
-            )
-        sys.path.insert(0, repo_dir)
-        from EditReward import EditRewardInferencer
-
-        self.inferencer = EditRewardInferencer(
-            config_path=config_path,
-            checkpoint_path=checkpoint_path,
-            device=device,
-            reward_dim=reward_dim,
-            rm_head_type=rm_head_type,
-        )
-        self.device = device
-        self.eval()
-
-    @torch.no_grad()
-    def __call__(self, prompts: List[str], source_images: List[Image.Image], edited_images: List[Image.Image]):
-        if not (len(prompts) == len(source_images) == len(edited_images)):
-            raise ValueError("prompts, source_images, and edited_images must have the same length.")
-
-        temp_dir = tempfile.mkdtemp(prefix="editreward_")
-        try:
-            source_paths = []
-            edited_paths = []
-            for index, (source_image, edited_image) in enumerate(zip(source_images, edited_images)):
-                source_path = os.path.join(temp_dir, f"source_{index}.png")
-                edited_path = os.path.join(temp_dir, f"edited_{index}.png")
-                source_image.convert("RGB").save(source_path)
-                edited_image.convert("RGB").save(edited_path)
-                source_paths.append(source_path)
-                edited_paths.append(edited_path)
-
-            rewards = self.inferencer.reward(
-                prompts=prompts,
-                image_src=source_paths,
-                image_paths=edited_paths,
-            )
-            return [reward[0].item() if hasattr(reward[0], "item") else float(reward[0]) for reward in rewards]
-        finally:
-            shutil.rmtree(temp_dir, ignore_errors=True)
-

rewards_services/api_services/editreward_scorer_service/gunicorn.conf.py

@@ -1,34 +0,0 @@
-import os
-import sys
-
-
-bind = f"{os.getenv('EDITREWARD_HOST', '127.0.0.1')}:{os.getenv('EDITREWARD_PORT', '18088')}"
-workers = int(os.getenv("EDITREWARD_WORKERS", os.getenv("EDITREWARD_NUM_DEVICES", "1")))
-worker_class = "sync"
-timeout = int(os.getenv("EDITREWARD_TIMEOUT", "600"))
-
-_raw_devices = os.getenv("EDITREWARD_CUDA_DEVICES") or os.getenv("CUDA_VISIBLE_DEVICES") or ""
-CUDA_DEVICES = [device.strip() for device in _raw_devices.split(",") if device.strip()]
-USED_DEVICES = set()
-
-
-def pre_fork(server, worker):
-    if not CUDA_DEVICES:
-        return
-    available = [device for device in CUDA_DEVICES if device not in USED_DEVICES]
-    worker.cuda_device = available[0] if available else CUDA_DEVICES[len(USED_DEVICES) % len(CUDA_DEVICES)]
-    USED_DEVICES.add(worker.cuda_device)
-    print(f"Worker {worker.pid} assigned CUDA_VISIBLE_DEVICES={worker.cuda_device}", file=sys.stderr)
-
-
-def post_fork(server, worker):
-    cuda_device = getattr(worker, "cuda_device", None)
-    if cuda_device is not None:
-        os.environ["CUDA_VISIBLE_DEVICES"] = cuda_device
-
-
-def child_exit(server, worker):
-    cuda_device = getattr(worker, "cuda_device", None)
-    if cuda_device is not None:
-        USED_DEVICES.discard(cuda_device)
-

rewards_services/api_services/editreward_scorer_service/requirements.txt

@@ -1,18 +0,0 @@
-flask
-gunicorn
-datasets
-huggingface_hub
-pillow
-openai
-megfile
-sentencepiece
-deepspeed
-fire
-omegaconf
-matplotlib
-peft
-trl==0.8.6
-tensorboard
-scipy
-transformers==4.56.1
-accelerate

rewards_services/api_services/editreward_scorer_service/run.sh

@@ -1,13 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-cd "$SCRIPT_DIR"
-
-export EDITREWARD_REPO_DIR="${EDITREWARD_REPO_DIR:-$SCRIPT_DIR/EditReward}"
-export EDITREWARD_PORT="${EDITREWARD_PORT:-18088}"
-export EDITREWARD_HOST="${EDITREWARD_HOST:-127.0.0.1}"
-export EDITREWARD_CUDA_DEVICES="${EDITREWARD_CUDA_DEVICES:-0,1}"
-export EDITREWARD_WORKERS="${EDITREWARD_WORKERS:-${EDITREWARD_NUM_DEVICES:-2}}"
-
-python -m gunicorn -c gunicorn.conf.py "app:create_app()"

scripts/train/deepspeed/zero3.json

@@ -1,39 +0,0 @@
-{
-  "fp16": {
-    "enabled": false,
-    "loss_scale": 0,
-    "loss_scale_window": 1000,
-    "initial_scale_power": 16,
-    "hysteresis": 2,
-    "min_loss_scale": 1
-  },
-  "bf16": {
-    "enabled": true
-  },
-  "zero_optimization": {
-    "stage": 3,
-    "offload_optimizer": {
-      "device": "none",
-      "pin_memory": true
-    },
-    "offload_param": {
-      "device": "none",
-      "pin_memory": true
-    },
-    "overlap_comm": true,
-    "contiguous_gradients": true,
-    "sub_group_size": 1000000000.0,
-    "reduce_bucket_size": "auto",
-    "stage3_prefetch_bucket_size": "auto",
-    "stage3_param_persistence_threshold": "auto",
-    "stage3_max_live_parameters": 1000000000.0,
-    "stage3_max_reuse_distance": 1000000000.0,
-    "stage3_gather_16bit_weights_on_model_save": true
-  },
-  "gradient_accumulation_steps": "auto",
-  "steps_per_print": 100,
-  "train_batch_size": "auto",
-  "train_micro_batch_size_per_gpu": "auto",
-  "wall_clock_breakdown": false
-}
-

scripts/train/edit_grpo.sh

@@ -1,77 +0,0 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
-REPO_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
-cd "$REPO_DIR"
-
-: "${MODEL_NAME_OR_PATH:?Set MODEL_NAME_OR_PATH to a pretrained Qwen-Kontext checkpoint.}"
-
-export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-2,3,4,5,6,7}"
-NPROC_PER_NODE="${NPROC_PER_NODE:-6}"
-MASTER_ADDR="${MASTER_ADDR:-localhost}"
-MASTER_PORT="${MASTER_PORT:-25000}"
-RUN_NAME="${RUN_NAME:-qwenkontext-edit-grpo}"
-OUTPUT_DIR="${OUTPUT_DIR:-outputs/rl/kontext/$RUN_NAME}"
-DEEPSPEED_CONFIG="${DEEPSPEED_CONFIG:-scripts/train/deepspeed/zero3.json}"
-EDITREWARD_URL="${EDITREWARD_URL:-http://127.0.0.1:18088/}"
-PROMPTS_FILE="${PROMPTS_FILE:-https://huggingface.co/wangfuyun/PrompRL/resolve/main/data/omni_edit_train_50k.parquet}"
-REPORT_TO="${REPORT_TO:-none}"
-if [[ -n "${WANDB_PROJECT:-}" && "$REPORT_TO" == "none" ]]; then
-  REPORT_TO="wandb"
-fi
-
-TORCHRUN_ARGS=(
-  --nproc_per_node="$NPROC_PER_NODE"
-  --nnodes="${NNODES:-1}"
-  --node_rank="${NODE_RANK:-0}"
-  --master_addr="$MASTER_ADDR"
-  --master_port="$MASTER_PORT"
-)
-
-TRAIN_ARGS=(
-  -m unirl.train_edit
-  --reward_funcs editreward format
-  --deepspeed "$DEEPSPEED_CONFIG"
-  --output_dir "$OUTPUT_DIR"
-  --model_name_or_path "$MODEL_NAME_OR_PATH"
-  --prompts_file "$PROMPTS_FILE"
-  --image_column "${IMAGE_COLUMN:-image}"
-  --prompt_column "${PROMPT_COLUMN:-prompt}"
-  --editreward_url "$EDITREWARD_URL"
-  --max_prompt_length "${MAX_PROMPT_LENGTH:-8192}"
-  --max_completion_length "${MAX_COMPLETION_LENGTH:-512}"
-  --num_generations "${NUM_GENERATIONS:-8}"
-  --num_skip_refinement "${NUM_SKIP_REFINEMENT:-2}"
-  --num_sde "${NUM_SDE:-4}"
-  --per_device_train_batch_size "${PER_DEVICE_TRAIN_BATCH_SIZE:-1}"
-  --gradient_accumulation_steps "${GRADIENT_ACCUMULATION_STEPS:-1}"
-  --logging_steps "${LOGGING_STEPS:-1}"
-  --learning_rate "${LEARNING_RATE:-3e-7}"
-  --bf16 "${BF16:-true}"
-  --report_to "$REPORT_TO"
-  --gradient_checkpointing "${GRADIENT_CHECKPOINTING:-true}"
-  --attn_implementation "${ATTN_IMPLEMENTATION:-flash_attention_2}"
-  --max_pixels "${MAX_PIXELS:-200704}"
-  --min_pixels "${MIN_PIXELS:-200704}"
-  --image_size "${IMAGE_SIZE:-512}"
-  --save_total_limit "${SAVE_TOTAL_LIMIT:-4}"
-  --save_strategy "${SAVE_STRATEGY:-steps}"
-  --save_steps "${SAVE_STEPS:-100}"
-  --beta "${BETA:-1e-2}"
-  --num_train_epochs "${NUM_TRAIN_EPOCHS:-10}"
-  --run_name "$RUN_NAME"
-)
-
-if [[ -n "${DATASET_CACHE_DIR:-}" ]]; then
-  TRAIN_ARGS+=(--dataset_cache_dir "$DATASET_CACHE_DIR")
-fi
-
-export PROMPTRL_EDIT_GUIDANCE_SCALE="${PROMPTRL_EDIT_GUIDANCE_SCALE:-${EDIT_GUIDANCE_SCALE:-2.5}}"
-export PROMPTRL_EDIT_NUM_INFERENCE_STEPS="${PROMPTRL_EDIT_NUM_INFERENCE_STEPS:-${EDIT_NUM_INFERENCE_STEPS:-8}}"
-export PROMPTRL_EDIT_HEIGHT="${PROMPTRL_EDIT_HEIGHT:-${EDIT_HEIGHT:-1024}}"
-export PROMPTRL_EDIT_WIDTH="${PROMPTRL_EDIT_WIDTH:-${EDIT_WIDTH:-1024}}"
-export DIT_LEARNING_RATE="${DIT_LEARNING_RATE:-2e-7}"
-export LLM_LEARNING_RATE="${LLM_LEARNING_RATE:-3e-7}"
-
-torchrun "${TORCHRUN_ARGS[@]}" "${TRAIN_ARGS[@]}"

unimodel/qwenkontext/fluxkontext_pipeline.py

@@ -14,10 +14,9 @@
 # limitations under the License.
 
 import inspect
-from typing import Any, Callable, Dict, List, Optional, Union, Tuple
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import numpy as np
-import math
 import torch
 from transformers import (
     CLIPImageProcessor,
@@ -1160,566 +1159,3 @@ class FluxKontextPipeline(
             return (image,)
 
         return FluxPipelineOutput(images=image)
-    
-    
-    
-    
-    # This method should be added to the FluxKontextPipeline class
-    def sde_sampling(
-        self,
-        image: Optional[PipelineImageInput] = None,
-        prompt: Union[str, List[str]] = None,
-        prompt_2: Optional[Union[str, List[str]]] = None,
-        negative_prompt: Union[str, List[str]] = None,
-        negative_prompt_2: Optional[Union[str, List[str]]] = None,
-        true_cfg_scale: float = 1.0,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 28,
-        sigmas: Optional[List[float]] = None,
-        guidance_scale: float = 3.5,
-        num_images_per_prompt: Optional[int] = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.FloatTensor] = None,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        ip_adapter_image: Optional[PipelineImageInput] = None,
-        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
-        negative_ip_adapter_image: Optional[PipelineImageInput] = None,
-        negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
-        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
-        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 512,
-        max_area: int = 1024**2,
-        num_sde: int = None,
-        _auto_resize: bool = True,
-    ):
-        r"""
-        SDE sampling function for FLUX Kontext pipeline with log probability tracking.
-        
-        This method performs stochastic differential equation (SDE) based sampling while
-        tracking log probabilities at each step. Useful for training and analysis purposes.
-        
-        Args:
-            image: Input image for image-to-image generation
-            prompt: Text prompt(s) to guide generation
-            prompt_2: Secondary text prompt for text_encoder_2
-            negative_prompt: Negative text prompt(s)
-            negative_prompt_2: Secondary negative prompt
-            true_cfg_scale: Classifier-free guidance scale (when > 1.0)
-            height: Output height in pixels
-            width: Output width in pixels
-            num_inference_steps: Number of denoising steps
-            sigmas: Custom sigma schedule
-            guidance_scale: Embedded guidance scale
-            num_images_per_prompt: Number of images per prompt
-            generator: Random number generator(s)
-            latents: Pre-generated latents
-            prompt_embeds: Pre-generated prompt embeddings
-            pooled_prompt_embeds: Pre-generated pooled embeddings
-            ip_adapter_image: IP-Adapter input image(s)
-            ip_adapter_image_embeds: Pre-generated IP-Adapter embeddings
-            negative_ip_adapter_image: Negative IP-Adapter image(s)
-            negative_ip_adapter_image_embeds: Negative IP-Adapter embeddings
-            negative_prompt_embeds: Pre-generated negative embeddings
-            negative_pooled_prompt_embeds: Pre-generated negative pooled embeddings
-            output_type: Output format ("pil" or "latent")
-            return_dict: Whether to return dict or tuple
-            joint_attention_kwargs: Additional attention parameters
-            callback_on_step_end: Callback function after each step
-            callback_on_step_end_tensor_inputs: Tensors to pass to callback
-            max_sequence_length: Maximum prompt sequence length
-            max_area: Maximum output area in pixels
-            _auto_resize: Whether to auto-resize to preferred resolutions
-            
-        Returns:
-            Tuple of (images, prev_latents, log_probs, pred_latents, timesteps, batched_states)
-        """
-        
-        height = height or self.default_sample_size * self.vae_scale_factor
-        width = width or self.default_sample_size * self.vae_scale_factor
-
-        original_height, original_width = height, width
-        aspect_ratio = width / height
-        
-        width = round((max_area * aspect_ratio) ** 0.5)
-        height = round((max_area / aspect_ratio) ** 0.5)
-
-        multiple_of = self.vae_scale_factor * 2
-        width = width // multiple_of * multiple_of
-        height = height // multiple_of * multiple_of
-
-        if height != original_height or width != original_width:
-            logger.warning(
-                f"Generation `height` and `width` have been adjusted to {height} and {width} to fit the model requirements."
-            )
-
-        # 1. Check inputs
-        self.check_inputs(
-            prompt,
-            prompt_2,
-            height,
-            width,
-            negative_prompt=negative_prompt,
-            negative_prompt_2=negative_prompt_2,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            max_sequence_length=max_sequence_length,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._joint_attention_kwargs = joint_attention_kwargs
-        self._current_timestep = None
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-
-        lora_scale = (
-            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
-        )
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
-        )
-        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
-        
-        # Encode prompts
-        (
-            prompt_embeds,
-            pooled_prompt_embeds,
-            text_ids,
-        ) = self.encode_prompt(
-            prompt=prompt,
-            prompt_2=prompt_2,
-            prompt_embeds=prompt_embeds,
-            pooled_prompt_embeds=pooled_prompt_embeds,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            max_sequence_length=max_sequence_length,
-            lora_scale=lora_scale,
-        )
-        
-        if do_true_cfg:
-            (
-                negative_prompt_embeds,
-                negative_pooled_prompt_embeds,
-                negative_text_ids,
-            ) = self.encode_prompt(
-                prompt=negative_prompt,
-                prompt_2=negative_prompt_2,
-                prompt_embeds=negative_prompt_embeds,
-                pooled_prompt_embeds=negative_pooled_prompt_embeds,
-                device=device,
-                num_images_per_prompt=num_images_per_prompt,
-                max_sequence_length=max_sequence_length,
-                lora_scale=lora_scale,
-            )
-
-        # 3. Preprocess image
-        if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-            from diffusers.pipelines.flux.pipeline_flux_kontext import PREFERRED_KONTEXT_RESOLUTIONS
-            
-            img = image[0] if isinstance(image, list) else image
-            image_height, image_width = self.image_processor.get_default_height_width(img)
-            aspect_ratio = image_width / image_height
-            if _auto_resize:
-                _, image_width, image_height = min(
-                    (abs(aspect_ratio - w / h), w, h) for w, h in PREFERRED_KONTEXT_RESOLUTIONS
-                )
-            image_width = image_width // multiple_of * multiple_of
-            image_height = image_height // multiple_of * multiple_of
-            image = self.image_processor.resize(image, image_height, image_width)
-            image = self.image_processor.preprocess(image, image_height, image_width)
-
-        # 4. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels // 4
-        latents, image_latents, latent_ids, image_ids = self.prepare_latents(
-            image,
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-        
-        if image_ids is not None:
-            latent_ids = torch.cat([latent_ids, image_ids], dim=0)
-
-        # 5. Prepare timesteps
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
-        image_seq_len = latents.shape[1]
-        
-        from diffusers.pipelines.flux.pipeline_flux_kontext import calculate_shift, retrieve_timesteps
-        
-        mu = calculate_shift(
-            image_seq_len,
-            self.scheduler.config.get("base_image_seq_len", 256),
-            self.scheduler.config.get("max_image_seq_len", 4096),
-            self.scheduler.config.get("base_shift", 0.5),
-            self.scheduler.config.get("max_shift", 1.15),
-        )
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            sigmas=sigmas,
-            mu=mu,
-        )
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        # Handle guidance
-        if self.transformer.config.guidance_embeds:
-            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
-            guidance = guidance.expand(latents.shape[0])
-        else:
-            guidance = None
-
-        # Handle IP-Adapter images
-        if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
-            negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
-        ):
-            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
-            negative_ip_adapter_image = [negative_ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
-
-        elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
-            negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
-        ):
-            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
-            ip_adapter_image = [ip_adapter_image] * self.transformer.encoder_hid_proj.num_ip_adapters
-
-        if self.joint_attention_kwargs is None:
-            self._joint_attention_kwargs = {}
-
-        image_embeds = None
-        negative_image_embeds = None
-        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
-            image_embeds = self.prepare_ip_adapter_image_embeds(
-                ip_adapter_image,
-                ip_adapter_image_embeds,
-                device,
-                batch_size * num_images_per_prompt,
-            )
-        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
-            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
-                negative_ip_adapter_image,
-                negative_ip_adapter_image_embeds,
-                device,
-                batch_size * num_images_per_prompt,
-            )
-
-        # 6. SDE Denoising loop with state tracking
-        prev_latents = []
-        pred_latents = []
-        states = {
-            "timestep": [],
-            "guidance": [],
-            "pooled_projections": [],
-            "encoder_hidden_states": [],
-            "txt_ids": None,
-            "img_ids": None,
-        }
-        log_probs = []
-        ts = []
-        
-        states["txt_ids"] = text_ids if text_ids is not None else None  
-        states["img_ids"] = latent_ids if latent_ids is not None else None
-
-        if num_sde is None:
-            num_sde = num_inference_steps
-        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
-                if image_embeds is not None:
-                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
-                
-                # Prepare model input
-                latent_model_input = latents
-                if image_latents is not None:
-                    latent_model_input = torch.cat([latents, image_latents], dim=1)
-                
-                timestep = (t.expand(latents.shape[0]) / 1000.).to(latents.dtype)
-                
-                
-                if i < num_sde:
-                    # Store states
-                    states["timestep"].append(timestep.unsqueeze(1))
-                    states["guidance"].append(guidance.unsqueeze(1) if torch.is_tensor(guidance) else guidance)
-                    states["pooled_projections"].append(pooled_prompt_embeds.unsqueeze(1) if pooled_prompt_embeds is not None else None)
-                    states["encoder_hidden_states"].append(prompt_embeds.unsqueeze(1) if prompt_embeds is not None else None)
-                    
-                    ts.append(t.expand(latents.shape[0]).unsqueeze(1))
-                    # prev_latents.append(latents.detach().clone().unsqueeze(1))
-                    prev_latents.append(latent_model_input.detach().clone().unsqueeze(1))
-
-                # Forward pass
-                noise_pred = self.transformer(
-                    hidden_states=latent_model_input,
-                    timestep=timestep,
-                    guidance=guidance,
-                    pooled_projections=pooled_prompt_embeds,
-                    encoder_hidden_states=prompt_embeds,
-                    txt_ids=text_ids,
-                    img_ids=latent_ids,
-                    joint_attention_kwargs=self.joint_attention_kwargs,
-                    return_dict=False,
-                )[0]
-                noise_pred = noise_pred[:, :latents.size(1)]
-
-                # Apply true CFG if needed
-                if do_true_cfg:
-                    if negative_image_embeds is not None:
-                        self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
-                    
-                    neg_latent_model_input = latents
-                    if image_latents is not None:
-                        neg_latent_model_input = torch.cat([latents, image_latents], dim=1)
-                        
-                    neg_noise_pred = self.transformer(
-                        hidden_states=neg_latent_model_input,
-                        timestep=timestep,
-                        guidance=guidance,
-                        pooled_projections=negative_pooled_prompt_embeds,
-                        encoder_hidden_states=negative_prompt_embeds,
-                        txt_ids=negative_text_ids,
-                        img_ids=latent_ids,
-                        joint_attention_kwargs=self.joint_attention_kwargs,
-                        return_dict=False,
-                    )[0]
-                    neg_noise_pred = neg_noise_pred[:, :latents.size(1)]
-                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
-
-                if i < num_sde:
-                    # SDE step with log probability
-                    latents_dtype = latents.dtype
-                    latents, log_prob, prev_latents_mean, std_dev = sde_step_with_logprob(
-                        self.scheduler,
-                        noise_pred.float(),
-                        t.expand(latents.shape[0]),
-                        latents.float()
-                    )
-                    
-                    log_probs.append(log_prob.detach().clone().unsqueeze(1))
-                    pred_latents.append(latents.detach().clone().unsqueeze(1))
-                
-                else:
-                    # Standard scheduler step
-                    latents_dtype = latents.dtype
-                    latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-                    
-                    
-                if latents.dtype != latents_dtype:
-                    latents = latents.to(latents_dtype)
-
-                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)
-
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        self._current_timestep = None
-
-        # Decode latents to images
-        if output_type == "latent":
-            image = latents
-        else:
-            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
-            image = self.vae.decode(latents, return_dict=False)[0]
-            image = self.image_processor.postprocess(image, output_type=output_type)
-
-        # Batch states for output
-        batched_states = {}
-        for key, value_list in states.items():
-            if value_list is None or len(value_list) == 0:
-                batched_states[key] = None
-                continue
-            if isinstance(value_list, list) and value_list[0] is None:
-                batched_states[key] = None
-                continue
-            if isinstance(value_list, list):
-                concatenated = torch.cat(value_list, dim=1)
-                if len(concatenated.shape) <= 2:
-                    batched_states[key] = concatenated.view(-1)
-                else:
-                    batched_states[key] = concatenated.view(-1, *concatenated.shape[2:])
-            else:
-                batched_states[key] = value_list
-
-        # Reshape outputs
-        prev_latents = torch.cat(prev_latents, dim=1)
-        log_probs = torch.cat(log_probs, dim=1)
-        pred_latents = torch.cat(pred_latents, dim=1)
-        ts = torch.cat(ts, dim=1)
-
-        prev_latents = prev_latents.view(prev_latents.shape[0] * prev_latents.shape[1], *prev_latents.shape[2:])
-        log_probs = log_probs.view(log_probs.shape[0] * log_probs.shape[1], *log_probs.shape[2:])
-        pred_latents = pred_latents.view(pred_latents.shape[0] * pred_latents.shape[1], *pred_latents.shape[2:])
-        ts = ts.view(-1)
-
-        # Offload models
-        self.maybe_free_model_hooks()
-
-        return (image, prev_latents, log_probs, pred_latents, ts, batched_states)
-    
-    
-def sde_step_with_logprob(
-    scheduler: FlowMatchEulerDiscreteScheduler,
-    model_output: torch.FloatTensor,
-    timestep: Union[float, torch.FloatTensor],
-    sample: torch.FloatTensor,
-    prev_sample: Optional[torch.FloatTensor] = None,
-    generator: Optional[torch.Generator] = None,
-) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
-    """
-    Predict the sample from the previous timestep by reversing the SDE with log probability tracking.
-    
-    Args:
-        scheduler: The FlowMatchEulerDiscreteScheduler instance
-        model_output: The direct output from learned flow model
-        timestep: The current discrete timestep in the diffusion chain
-        sample: A current instance of a sample created by the diffusion process
-        prev_sample: Optional pre-computed previous sample
-        generator: A random number generator
-        
-    Returns:
-        Tuple of (prev_sample, log_prob, prev_sample_mean, std_dev)
-    """
-    step_index = [scheduler.index_for_timestep(t) for t in timestep]
-    prev_step_index = [step + 1 for step in step_index]
-    sigma = scheduler.sigmas[step_index].view(-1, 1, 1).to(model_output.device)
-    sigma_prev = scheduler.sigmas[prev_step_index].view(-1, 1, 1).to(model_output.device)
-    sigma_max = scheduler.sigmas[1].item()
-    dt = sigma_prev - sigma
-
-    std_dev_t = torch.sqrt(sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))) * 0.8
-    
-    # SDE formulation
-    prev_sample_mean = (
-        sample * (1 + std_dev_t**2 / (2 * sigma) * dt) + 
-        model_output * (1 + std_dev_t**2 * (1 - sigma) / (2 * sigma)) * dt
-    )
-    
-    if prev_sample is not None and generator is not None:
-        raise ValueError(
-            "Cannot pass both generator and prev_sample. Please make sure that either `generator` or"
-            " `prev_sample` stays `None`."
-        )
-
-    if prev_sample is None:
-        variance_noise = randn_tensor(
-            model_output.shape,
-            generator=generator,
-            device=model_output.device,
-            dtype=model_output.dtype,
-        )
-        prev_sample = prev_sample_mean + std_dev_t * torch.sqrt(-1 * dt) * variance_noise
-    
-    # Calculate log probability
-    variance = (std_dev_t * torch.sqrt(-1 * dt)) ** 2
-    log_prob = (
-        -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * variance)
-        - torch.log(torch.sqrt(variance))
-        - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
-    )
-
-    # Mean along all but batch dimension
-    log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
-    
-    return prev_sample, log_prob, prev_sample_mean, std_dev_t * torch.sqrt(-1 * dt)
-
-
-def sde_step_with_logprob_simple(
-    scheduler: FlowMatchEulerDiscreteScheduler,
-    model_output: torch.FloatTensor,
-    timestep: Union[float, torch.FloatTensor],
-    sample: torch.FloatTensor,
-    prev_sample: Optional[torch.FloatTensor] = None,
-    generator: Optional[torch.Generator] = None,
-) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]:
-    """
-    Simplified SDE step with log probability tracking using eta parameter.
-    
-    Args:
-        scheduler: The FlowMatchEulerDiscreteScheduler instance
-        model_output: The direct output from learned flow model
-        timestep: The current discrete timestep in the diffusion chain
-        sample: A current instance of a sample created by the diffusion process
-        prev_sample: Optional pre-computed previous sample
-        generator: A random number generator
-        
-    Returns:
-        Tuple of (prev_sample, log_prob, prev_sample_mean, std_dev)
-    """
-    step_index = [scheduler.index_for_timestep(t) for t in timestep]
-    prev_step_index = [step + 1 for step in step_index]
-    sigma = scheduler.sigmas[step_index].view(-1, 1, 1).to(model_output.device)
-    sigma_prev = scheduler.sigmas[prev_step_index].view(-1, 1, 1).to(model_output.device)
-    sigma_max = scheduler.sigmas[1].item()
-    dt = sigma_prev - sigma
-    
-    eta = 0.5
-    Dt = -dt * eta
-    
-    prev_sample_mean = (
-        sample * (1 - Dt / (1 - torch.where(sigma == 1, sigma_max, sigma))) + 
-        model_output * (dt - Dt)
-    )
-    
-    std_dev_t = torch.sqrt(2 * Dt * (sigma / (1 - torch.where(sigma == 1, sigma_max, sigma))))
-    
-    if prev_sample is not None and generator is not None:
-        raise ValueError(
-            "Cannot pass both generator and prev_sample. Please make sure that either `generator` or"
-            " `prev_sample` stays `None`."
-        )
-
-    if prev_sample is None:
-        variance_noise = randn_tensor(
-            model_output.shape,
-            generator=generator,
-            device=model_output.device,
-            dtype=model_output.dtype,
-        )
-        prev_sample = prev_sample_mean + std_dev_t * variance_noise
-
-    # Calculate log probability
-    log_prob = (
-        -((prev_sample.detach() - prev_sample_mean) ** 2) / (2 * (std_dev_t**2))
-        - torch.log(std_dev_t)
-        - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi)))
-    )
-
-    # Mean along all but batch dimension
-    log_prob = log_prob.mean(dim=tuple(range(1, log_prob.ndim)))
-    
-    return prev_sample, log_prob, prev_sample_mean, std_dev_t

unirl/__init__.py

@@ -1,2 +0,0 @@
-"""PromptRL training package."""
-

unirl/reward_evaluator/__init__.py

@@ -1,4 +0,0 @@
-from .reward_evaluator import RewardEvaluatorClient
-
-__all__ = ["RewardEvaluatorClient"]
-

unirl/reward_evaluator/reward_evaluator.py

@@ -1,71 +0,0 @@
-import pickle
-from io import BytesIO
-from typing import Any, Dict, List, Mapping, Optional, Union
-
-import requests
-from PIL import Image
-
-
-DEFAULT_EDITREWARD_URL = "http://127.0.0.1:18088/"
-
-
-def _serialize_image(image: Image.Image) -> bytes:
-    buffer = BytesIO()
-    if image.mode != "RGB":
-        image = image.convert("RGB")
-    image.save(buffer, format="PNG")
-    return buffer.getvalue()
-
-
-def _serialize_images(
-    images: Union[List[Image.Image], Mapping[str, List[Image.Image]]],
-) -> Union[List[bytes], Dict[str, List[bytes]]]:
-    if isinstance(images, Mapping):
-        return {key: [_serialize_image(image) for image in value] for key, value in images.items()}
-    return [_serialize_image(image) for image in images]
-
-
-def _create_payload(
-    images: Union[List[Image.Image], Mapping[str, List[Image.Image]]],
-    prompts: List[str],
-    metadata: Optional[Dict[str, Any]] = None,
-) -> bytes:
-    return pickle.dumps(
-        {
-            "images": _serialize_images(images),
-            "prompts": prompts,
-            "metadata": metadata or {},
-        }
-    )
-
-
-class RewardEvaluatorClient:
-    """HTTP client for the EditReward scorer service."""
-
-    def __init__(self, editreward_url: str = DEFAULT_EDITREWARD_URL, timeout: int = 600):
-        self.editreward_url = editreward_url
-        self.timeout = timeout
-
-    def evaluate_editreward(
-        self,
-        source_images: List[Image.Image],
-        edited_images: List[Image.Image],
-        prompts: List[str],
-    ) -> Dict[str, Any]:
-        if not (len(source_images) == len(edited_images) == len(prompts)):
-            raise ValueError(
-                "EditReward inputs must have equal lengths: "
-                f"{len(source_images)} source images, {len(edited_images)} edited images, {len(prompts)} prompts."
-            )
-
-        payload = _create_payload(
-            {"source": source_images, "edited": edited_images},
-            prompts,
-        )
-        response = requests.post(self.editreward_url, data=payload, timeout=self.timeout)
-        response.raise_for_status()
-        result = pickle.loads(response.content)
-        if isinstance(result, dict) and "error" in result:
-            raise RuntimeError(f"EditReward service returned an error: {result['error']}")
-        return result
-

unirl/train_edit.py

@@ -1,265 +0,0 @@
-import json
-import os
-import re
-from dataclasses import dataclass, field
-from io import BytesIO
-from typing import Any, Dict, List, Optional
-
-from datasets import load_dataset
-from PIL import Image, ImageOps
-from torch.utils.data import Dataset
-from transformers.trainer_utils import get_last_checkpoint
-from trl import GRPOConfig, ModelConfig, ScriptArguments, TrlParser, get_peft_config
-
-from .reward_evaluator import RewardEvaluatorClient
-from .trainer import QwenKontextEditGRPOTrainer
-
-
-DEFAULT_EDIT_DATASET = "https://huggingface.co/wangfuyun/PrompRL/resolve/main/data/omni_edit_train_50k.parquet"
-
-EDIT_QUESTION_TEMPLATE = """Please provide an enhanced prompt for the following image editing prompt.
-Ensure the revised prompt is clear, specific, and includes detailed instructions to achieve the desired outcome while maintaining the original intent.
-Original prompt: {Question}. Directly provide the improved prompt in <answer> </answer> tags."""
-
-
-@dataclass
-class EditGRPOScriptArguments(ScriptArguments):
-    reward_funcs: List[str] = field(
-        default_factory=lambda: ["editreward", "format"],
-        metadata={"help": "Reward functions to use. Edit training supports only: editreward, format."},
-    )
-    prompts_file: str = field(
-        default=DEFAULT_EDIT_DATASET,
-        metadata={"help": "Path or URL to a .parquet or .jsonl edit-training dataset."},
-    )
-    image_column: str = field(default="image", metadata={"help": "Dataset column containing the source image."})
-    prompt_column: str = field(default="prompt", metadata={"help": "Dataset column containing the edit instruction."})
-    caption_column: Optional[str] = field(default="caption", metadata={"help": "Optional source-caption column."})
-    target_caption_column: Optional[str] = field(
-        default="target_caption",
-        metadata={"help": "Optional target-caption column."},
-    )
-    image_size: int = field(default=512, metadata={"help": "Center-cropped source image size used for editing."})
-    dataset_cache_dir: Optional[str] = field(
-        default=None,
-        metadata={"help": "Optional Hugging Face datasets cache dir. Defaults to HF_DATASETS_CACHE."},
-    )
-    editreward_url: str = field(
-        default="http://127.0.0.1:18088/",
-        metadata={"help": "HTTP URL of the EditReward scorer service."},
-    )
-    processor_name_or_path: str = field(
-        default="Qwen/Qwen2.5-VL-3B-Instruct",
-        metadata={"help": "Processor used for Qwen2.5-VL chat formatting and image preprocessing."},
-    )
-    max_pixels: int = field(default=200704, metadata={"help": "Maximum pixels passed to the Qwen-VL processor."})
-    min_pixels: int = field(default=200704, metadata={"help": "Minimum pixels passed to the Qwen-VL processor."})
-    num_skip_refinement: int = field(
-        default=2,
-        metadata={"help": "Generations per input that use the original edit prompt instead of a Qwen-refined prompt."},
-    )
-    num_sde: int = field(
-        default=4,
-        metadata={"help": "Number of FLUX denoising steps sampled with SDE log-prob tracking for diffusion GRPO."},
-    )
-
-
-class EditPromptDataset(Dataset):
-    """Loads image-edit instructions from parquet or jsonl files."""
-
-    def __init__(
-        self,
-        prompts_file: str,
-        question_template: str,
-        image_column: str = "image",
-        prompt_column: str = "prompt",
-        caption_column: Optional[str] = "caption",
-        target_caption_column: Optional[str] = "target_caption",
-        image_size: int = 512,
-        cache_dir: Optional[str] = None,
-    ):
-        self.prompts_file = normalize_data_path(prompts_file)
-        self.question_template = question_template
-        self.image_column = image_column
-        self.prompt_column = prompt_column
-        self.caption_column = caption_column
-        self.target_caption_column = target_caption_column
-        self.image_size = image_size
-        self.base_dir = (
-            os.path.dirname(os.path.abspath(self.prompts_file))
-            if not is_remote_path(self.prompts_file)
-            else os.getcwd()
-        )
-
-        if self.prompts_file.endswith(".parquet"):
-            self.records = load_dataset(
-                "parquet",
-                data_files={"train": self.prompts_file},
-                split="train",
-                cache_dir=cache_dir or os.getenv("HF_DATASETS_CACHE"),
-            )
-        elif self.prompts_file.endswith(".jsonl") or self.prompts_file.endswith(".json"):
-            if is_remote_path(self.prompts_file):
-                self.records = load_dataset(
-                    "json",
-                    data_files={"train": self.prompts_file},
-                    split="train",
-                    cache_dir=cache_dir or os.getenv("HF_DATASETS_CACHE"),
-                )
-            else:
-                with open(self.prompts_file, "r", encoding="utf-8") as file:
-                    self.records = [json.loads(line) for line in file if line.strip()]
-        else:
-            raise ValueError("Edit training datasets must be .parquet or .jsonl files.")
-
-        if len(self.records) == 0:
-            raise ValueError(f"No training records found in {prompts_file}.")
-
-    def __len__(self) -> int:
-        return len(self.records)
-
-    def __getitem__(self, index: int) -> Dict[str, Any]:
-        item = self.records[index]
-        instruction = self._read_text(item, self.prompt_column)
-        image = self._read_image(item, self.image_column)
-        formatted_prompt = self.question_template.format(Question=instruction)
-
-        return {
-            "image": image,
-            "caption": self._read_optional_text(item, self.caption_column),
-            "target_caption": self._read_optional_text(item, self.target_caption_column),
-            "editing_instruction": instruction,
-            "prompt": [
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "image"},
-                        {"type": "text", "text": formatted_prompt},
-                    ],
-                }
-            ],
-        }
-
-    def _read_text(self, item: Dict[str, Any], column: str) -> str:
-        if column not in item:
-            raise KeyError(f"Missing required column '{column}' in {self.prompts_file}.")
-        value = item[column]
-        if isinstance(value, list):
-            value = value[-1] if value else ""
-        if value is None or not str(value).strip():
-            raise ValueError(f"Empty edit instruction in column '{column}'.")
-        return str(value).strip()
-
-    def _read_optional_text(self, item: Dict[str, Any], column: Optional[str]) -> str:
-        if not column or column not in item or item[column] is None:
-            return ""
-        value = item[column]
-        if isinstance(value, list):
-            value = value[-1] if value else ""
-        return str(value).strip()
-
-    def _read_image(self, item: Dict[str, Any], column: str) -> Image.Image:
-        if column not in item:
-            raise KeyError(f"Missing required image column '{column}' in {self.prompts_file}.")
-        image = self._coerce_image(item[column])
-        if self.image_size > 0:
-            image = ImageOps.fit(image, (self.image_size, self.image_size), method=Image.Resampling.BICUBIC)
-        return image.convert("RGB")
-
-    def _coerce_image(self, value: Any) -> Image.Image:
-        if isinstance(value, Image.Image):
-            return value.convert("RGB")
-        if isinstance(value, str):
-            image_path = value if os.path.isabs(value) else os.path.join(self.base_dir, value)
-            return Image.open(image_path).convert("RGB")
-        if isinstance(value, bytes):
-            return Image.open(BytesIO(value)).convert("RGB")
-        if isinstance(value, dict):
-            if value.get("bytes") is not None:
-                return Image.open(BytesIO(value["bytes"])).convert("RGB")
-            if value.get("path") is not None:
-                image_path = value["path"] if os.path.isabs(value["path"]) else os.path.join(self.base_dir, value["path"])
-                return Image.open(image_path).convert("RGB")
-        raise TypeError(f"Unsupported image value type: {type(value)!r}")
-
-
-def is_remote_path(path: str) -> bool:
-    return path.startswith(("http://", "https://", "hf://"))
-
-
-def normalize_data_path(path: str) -> str:
-    if path.startswith("hf://"):
-        parts = path[len("hf://") :].split("/", 2)
-        if len(parts) != 3:
-            raise ValueError("hf:// dataset paths must look like hf://owner/repo/path/to/file.parquet")
-        repo_id = f"{parts[0]}/{parts[1]}"
-        file_path = parts[2]
-        return f"https://huggingface.co/{repo_id}/resolve/main/{file_path}"
-
-    if "huggingface.co/" in path and "/blob/" in path:
-        return path.replace("/blob/", "/resolve/", 1)
-    return path
-
-
-def format_reward(completions: List[str]) -> List[float]:
-    pattern = re.compile(r"<answer>.*?</answer>", re.DOTALL)
-    return [1.0 if pattern.search(completion) else 0.0 for completion in completions]
-
-
-def build_editreward_func(editreward_url: str):
-    reward_client = RewardEvaluatorClient(editreward_url=editreward_url)
-
-    def editreward(source_images, edited_images, prompts):
-        return reward_client.evaluate_editreward(source_images, edited_images, prompts)
-
-    return editreward
-
-
-def main(script_args: EditGRPOScriptArguments, training_args: GRPOConfig, model_args: ModelConfig) -> None:
-    supported_rewards = {"editreward", "format"}
-    unsupported_rewards = sorted(set(script_args.reward_funcs) - supported_rewards)
-    if unsupported_rewards:
-        raise ValueError(f"Edit training supports only {sorted(supported_rewards)}, got {unsupported_rewards}.")
-
-    reward_registry = {
-        "editreward": build_editreward_func(script_args.editreward_url),
-        "format": format_reward,
-    }
-    reward_funcs = [(name, None, reward_registry[name]) for name in script_args.reward_funcs]
-
-    train_dataset = EditPromptDataset(
-        prompts_file=script_args.prompts_file,
-        question_template=EDIT_QUESTION_TEMPLATE,
-        image_column=script_args.image_column,
-        prompt_column=script_args.prompt_column,
-        caption_column=script_args.caption_column,
-        target_caption_column=script_args.target_caption_column,
-        image_size=script_args.image_size,
-        cache_dir=script_args.dataset_cache_dir,
-    )
-
-    trainer = QwenKontextEditGRPOTrainer(
-        model=model_args.model_name_or_path,
-        reward_funcs=reward_funcs,
-        args=training_args,
-        train_dataset=train_dataset,
-        peft_config=get_peft_config(model_args),
-        max_pixels=script_args.max_pixels,
-        min_pixels=script_args.min_pixels,
-        processor_name_or_path=script_args.processor_name_or_path,
-        attn_implementation=model_args.attn_implementation,
-        num_skip_refinement=script_args.num_skip_refinement,
-        num_sde=script_args.num_sde,
-    )
-
-    checkpoint = get_last_checkpoint(training_args.output_dir) if os.path.isdir(training_args.output_dir) else None
-    trainer.train(resume_from_checkpoint=checkpoint)
-    trainer.save_model(training_args.output_dir)
-    if training_args.push_to_hub:
-        trainer.push_to_hub(dataset_name=script_args.dataset_name)
-
-
-if __name__ == "__main__":
-    parser = TrlParser((EditGRPOScriptArguments, GRPOConfig, ModelConfig))
-    parsed_script_args, parsed_training_args, parsed_model_args = parser.parse_args_and_config()
-    main(parsed_script_args, parsed_training_args, parsed_model_args)

unirl/trainer/__init__.py

@@ -1,4 +0,0 @@
-from .edit_grpo_trainer import QwenKontextEditGRPOTrainer
-
-__all__ = ["QwenKontextEditGRPOTrainer"]
-

unirl/trainer/edit_grpo_trainer.py

@@ -1,623 +0,0 @@
-import os
-from collections import defaultdict
-from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import transformers
-from accelerate.utils import DistributedType
-from datasets import Dataset, IterableDataset
-from packaging import version
-from PIL import Image
-from transformers import AutoProcessor, GenerationConfig, PreTrainedModel, Trainer, TrainerCallback
-from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled
-from transformers.utils import is_peft_available
-from trl.data_utils import maybe_apply_chat_template
-from trl.models import create_reference_model, prepare_deepspeed, unwrap_model_for_generation
-from trl.trainer.grpo_config import GRPOConfig
-
-from unimodel.qwenkontext.fluxkontext_pipeline import sde_step_with_logprob
-from unimodel.qwenkontext.qwenkontext_inference import QwenKontextForInferenceLM
-
-if is_peft_available():
-    from peft import PeftConfig, get_peft_model
-
-
-RewardFunc = Callable[..., Union[List[float], Dict[str, Any]]]
-
-
-def compute_log_prob(
-    model_pred: torch.Tensor,
-    scheduler,
-    prev_latents: torch.Tensor,
-    pred_latents: torch.Tensor,
-    timesteps: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-    return sde_step_with_logprob(
-        scheduler,
-        model_pred.float(),
-        timesteps,
-        prev_latents.float(),
-        pred_latents.float(),
-    )
-
-
-class QwenKontextEditGRPOTrainer(Trainer):
-    """Joint GRPO trainer for Qwen prompt refinement and FLUX.1-Kontext edit generation."""
-
-    def __init__(
-        self,
-        model: Union[str, PreTrainedModel],
-        reward_funcs: List[Tuple[str, Optional[Any], RewardFunc]],
-        args: Optional[GRPOConfig] = None,
-        train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
-        eval_dataset: Optional[Union[Dataset, IterableDataset, Dict[str, Union[Dataset, IterableDataset]]]] = None,
-        processing_class: Optional[Any] = None,
-        callbacks: Optional[List[TrainerCallback]] = None,
-        optimizers: Tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None),
-        peft_config: Optional["PeftConfig"] = None,
-        max_pixels: int = 200704,
-        min_pixels: int = 200704,
-        processor_name_or_path: str = "Qwen/Qwen2.5-VL-3B-Instruct",
-        attn_implementation: str = "flash_attention_2",
-        num_skip_refinement: int = 2,
-        num_sde: int = 4,
-    ):
-        if args is None:
-            model_name = model if isinstance(model, str) else model.config._name_or_path
-            args = GRPOConfig(f"{os.path.basename(model_name)}-edit-joint-grpo")
-
-        model_init_kwargs = args.model_init_kwargs or {}
-        model_init_kwargs["attn_implementation"] = attn_implementation
-        model_init_kwargs["use_cache"] = False if args.gradient_checkpointing else model_init_kwargs.get("use_cache")
-
-        if isinstance(model, str):
-            self.model_id = model
-            model = self._load_model(model, model_init_kwargs)
-        else:
-            self.model_id = model.config._name_or_path
-            if args.model_init_kwargs is not None:
-                raise ValueError("model_init_kwargs can only be used when model is a path.")
-
-        if peft_config is not None:
-            model = get_peft_model(model, peft_config)
-
-        self._configure_trainable_parameters(model)
-        self.ref_model = self._create_reference_model(model, model_init_kwargs)
-        self.scheduler = model.get_model().diffusion_expert.scheduler
-
-        if processing_class is None:
-            processing_class = self._create_processor(processor_name_or_path, max_pixels, min_pixels)
-        self.processing_class = processing_class
-        self.reward_funcs = reward_funcs
-        self.max_prompt_length = args.max_prompt_length
-        self.num_generations = args.num_generations
-        self.beta = args.beta
-        self.num_sde = num_sde
-
-        if not 0 <= num_skip_refinement < self.num_generations:
-            raise ValueError(
-                f"num_skip_refinement must be in [0, num_generations), got {num_skip_refinement} "
-                f"for num_generations={self.num_generations}."
-            )
-        self.num_skip_refinement = num_skip_refinement
-        self.num_refined = self.num_generations - num_skip_refinement
-
-        self.generation_config = GenerationConfig(
-            max_new_tokens=args.max_completion_length or 256,
-            do_sample=True,
-            temperature=1.0,
-            num_return_sequences=1,
-            pad_token_id=processing_class.pad_token_id,
-            eos_token_id=processing_class.eos_token_id,
-        )
-        model.generation_config = self.generation_config
-        self.ref_model.generation_config = self.generation_config
-        if hasattr(model, "warnings_issued"):
-            model.warnings_issued["estimate_tokens"] = True
-
-        self._metrics = defaultdict(list)
-
-        def data_collator(features):
-            return features
-
-        super().__init__(
-            model=model,
-            args=args,
-            data_collator=data_collator,
-            train_dataset=train_dataset,
-            eval_dataset=eval_dataset,
-            processing_class=processing_class,
-            callbacks=callbacks,
-            optimizers=optimizers,
-        )
-        self.model_accepts_loss_kwargs = False
-
-        if self.is_deepspeed_enabled and is_deepspeed_zero3_enabled():
-            self.ref_model = prepare_deepspeed(self.ref_model, self.accelerator)
-        else:
-            self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
-
-        self.diffusion_generation_config = self._get_diffusion_config()
-        self.start_time = datetime.now().strftime("%Y-%m-%d_%H-%M")
-        self.log_dir = os.path.join(args.output_dir, "training_samples", self.start_time)
-        os.makedirs(self.log_dir, exist_ok=True)
-
-    def _load_model(self, model_id: str, model_init_kwargs: Dict[str, Any]) -> PreTrainedModel:
-        torch_dtype = model_init_kwargs.get("torch_dtype")
-        if isinstance(torch_dtype, str) and torch_dtype != "auto":
-            model_init_kwargs["torch_dtype"] = getattr(torch, torch_dtype)
-        if "qwenkontext" not in model_id.lower():
-            raise ValueError("Edit joint training expects a Qwen-Kontext checkpoint path.")
-        return QwenKontextForInferenceLM.from_pretrained(model_id, **model_init_kwargs)
-
-    def _create_reference_model(self, model: PreTrainedModel, model_init_kwargs: Dict[str, Any]) -> PreTrainedModel:
-        if is_deepspeed_zero3_enabled():
-            ref_model = self._load_model(self.model_id, model_init_kwargs)
-        else:
-            ref_model = create_reference_model(model)
-        for parameter in ref_model.parameters():
-            parameter.requires_grad = False
-        return ref_model
-
-    def _configure_trainable_parameters(self, model: PreTrainedModel) -> None:
-        try:
-            model.get_model().diffusion_expert.enable_vae_slicing()
-        except AttributeError:
-            try:
-                model.get_model().diffusion_expert.vae.enable_slicing()
-            except AttributeError:
-                pass
-
-        for parameter in model.parameters():
-            parameter.requires_grad = False
-        for parameter in model.get_model().parameters():
-            parameter.requires_grad = True
-        for parameter in model.lm_head.parameters():
-            parameter.requires_grad = True
-
-        if hasattr(model, "visual"):
-            for parameter in model.visual.parameters():
-                parameter.requires_grad = False
-
-        for component_name in ("visual", "vae", "text_encoder", "text_encoder_2", "text_encoder_3"):
-            component = getattr(model.get_model(), component_name, None)
-            if component is not None:
-                for parameter in component.parameters():
-                    parameter.requires_grad = False
-
-        transformer = getattr(model.get_model(), "transformer", None)
-        if transformer is None:
-            raise ValueError("Qwen-Kontext model does not expose a FLUX transformer.")
-        for parameter in transformer.parameters():
-            parameter.requires_grad = True
-
-    def _create_processor(self, processor_name_or_path: str, max_pixels: int, min_pixels: int) -> AutoProcessor:
-        processor = AutoProcessor.from_pretrained(processor_name_or_path)
-        processor.pad_token_id = processor.tokenizer.pad_token_id
-        processor.eos_token_id = processor.tokenizer.eos_token_id
-        processor.image_processor.max_pixels = max_pixels
-        processor.image_processor.min_pixels = min_pixels
-        return processor
-
-    def _get_diffusion_config(self) -> Dict[str, Any]:
-        device_text = str(self.accelerator.device)
-        device_id = int(device_text.split(":")[-1]) if ":" in device_text else 0
-        return {
-            "guidance_scale": float(os.getenv("PROMPTRL_EDIT_GUIDANCE_SCALE", "2.5")),
-            "num_inference_steps": int(os.getenv("PROMPTRL_EDIT_NUM_INFERENCE_STEPS", "8")),
-            "num_images_per_prompt": 1,
-            "generator": torch.manual_seed(42 + device_id),
-            "height": int(os.getenv("PROMPTRL_EDIT_HEIGHT", "1024")),
-            "width": int(os.getenv("PROMPTRL_EDIT_WIDTH", "1024")),
-            "num_sde": self.num_sde,
-        }
-
-    def _set_signature_columns_if_needed(self):
-        if self._signature_columns is None:
-            self._signature_columns = ["prompt"]
-
-    def create_optimizer(self):
-        if self.optimizer is not None:
-            return self.optimizer
-
-        optimizer_kwargs = {
-            "betas": (self.args.adam_beta1, self.args.adam_beta2),
-            "eps": self.args.adam_epsilon,
-            "weight_decay": self.args.weight_decay,
-        }
-        dit_lr = float(os.getenv("DIT_LEARNING_RATE", os.getenv("PROMPTRL_DIT_LR", "2e-7")))
-        llm_lr = float(os.getenv("LLM_LEARNING_RATE", os.getenv("PROMPTRL_LLM_LR", "3e-7")))
-
-        dit_params = [
-            parameter for parameter in self.model.get_model().transformer.parameters() if parameter.requires_grad
-        ]
-        dit_param_ids = {id(parameter) for parameter in dit_params}
-        llm_params = [
-            parameter
-            for parameter in self.model.parameters()
-            if parameter.requires_grad and id(parameter) not in dit_param_ids
-        ]
-
-        param_groups = []
-        if dit_params:
-            param_groups.append({"params": dit_params, "lr": dit_lr})
-        if llm_params:
-            param_groups.append({"params": llm_params, "lr": llm_lr})
-        if not param_groups:
-            raise ValueError("No trainable parameters were found for edit joint GRPO training.")
-
-        self.optimizer = torch.optim.AdamW(param_groups, **optimizer_kwargs)
-        return self.optimizer
-
-    def training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch=None):
-        model.eval()
-        self.ref_model.eval()
-        if hasattr(self.optimizer, "train") and callable(self.optimizer.train):
-            self.optimizer.train()
-
-        inputs = self._prepare_inputs(inputs)
-
-        def loss_update(loss: torch.Tensor, scale_factor: float = 1.0) -> None:
-            if self.args.n_gpu > 1:
-                loss = loss.mean()
-            if not self.model_accepts_loss_kwargs and self.compute_loss_func is None:
-                loss = loss / self.args.gradient_accumulation_steps
-            if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
-                loss = loss / scale_factor
-                model.backward(loss)
-            else:
-                self.accelerator.backward(loss / scale_factor)
-
-        with self.compute_loss_context_manager():
-            generations = self.generate_samples(model, inputs)
-            torch.cuda.empty_cache()
-
-            if self.num_refined > 0:
-                cot_loss = self.cot_loss_computation(
-                    model,
-                    generations["prompt_completion_ids"],
-                    generations["completion_ids"],
-                    generations["prompt_length"],
-                    generations["advantages_refined"],
-                    generations["prompt_inputs"],
-                )
-                loss_update(cot_loss, 1.0)
-            else:
-                cot_loss = torch.tensor(0.0, device=self.accelerator.device)
-
-            diff_advantages = generations["advantages"].repeat_interleave(self.num_sde, dim=0)
-            total_len = diff_advantages.shape[0]
-            diff_loss_values = []
-            diff_kl_values = []
-            diffusion_batch_size = int(os.getenv("PROMPTRL_DIFFUSION_LOSS_BATCH_SIZE", "4"))
-
-            for idx in range(0, total_len, diffusion_batch_size):
-                batched_states_slice = {}
-                for key, value in generations["batched_states"].items():
-                    if key in {"img_ids", "txt_ids"}:
-                        batched_states_slice[key] = value
-                    elif value is None:
-                        batched_states_slice[key] = None
-                    else:
-                        batched_states_slice[key] = value[idx : idx + diffusion_batch_size]
-
-                diff_loss, diff_kl = self.diffusion_loss_computation(
-                    generations["prev_latents"][idx : idx + diffusion_batch_size],
-                    generations["diff_sampling_log_probs"][idx : idx + diffusion_batch_size],
-                    generations["pred_latents"][idx : idx + diffusion_batch_size],
-                    generations["ts"][idx : idx + diffusion_batch_size],
-                    batched_states_slice,
-                    diff_advantages[idx : idx + diffusion_batch_size],
-                )
-                loss_update(diff_loss, max(1.0, float(total_len / diffusion_batch_size)))
-                diff_loss_values.append(diff_loss.detach())
-                diff_kl_values.append(diff_kl.detach())
-
-        diff_loss = torch.stack(diff_loss_values).mean()
-        diff_kl = torch.stack(diff_kl_values).mean()
-        loss = diff_loss + cot_loss.detach()
-
-        if self.args.torch_empty_cache_steps is not None and self.state.global_step % self.args.torch_empty_cache_steps == 0:
-            torch.cuda.empty_cache()
-
-        if hasattr(model, "step") and callable(model.step):
-            model.step()
-
-        self._metrics["diff_kl"].append(self.accelerator.gather_for_metrics(diff_kl).mean().item())
-        self._metrics["diff_loss"].append(self.accelerator.gather_for_metrics(diff_loss).mean().item())
-        torch.cuda.empty_cache()
-        return loss.detach()
-
-    def generate_samples(self, model: nn.Module, inputs: List[Dict]) -> Dict[str, Any]:
-        source_images = [example["image"] for example in inputs]
-        batch_size = len(inputs)
-        prompt_inputs = None
-        prompt_completion_ids = None
-        completion_ids = None
-        prompt_length = 0
-        completions_refined: List[str] = []
-        refined_prompts: List[str] = []
-
-        if self.num_refined > 0:
-            prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]
-            prompt_inputs = self.processing_class(
-                images=[image for image in source_images for _ in range(self.num_refined)],
-                text=[prompt for prompt in prompts_text for _ in range(self.num_refined)],
-                return_tensors="pt",
-                padding=True,
-                padding_side="left",
-                add_special_tokens=False,
-            )
-            prompt_inputs = super()._prepare_inputs(prompt_inputs)
-            if self.max_prompt_length is not None:
-                prompt_inputs["input_ids"] = prompt_inputs["input_ids"][:, -self.max_prompt_length :]
-                prompt_inputs["attention_mask"] = prompt_inputs["attention_mask"][:, -self.max_prompt_length :]
-
-            with unwrap_model_for_generation(model, self.accelerator) as unwrapped_model:
-                with torch.no_grad():
-                    prompt_completion_ids = unwrapped_model.generate(
-                        **prompt_inputs,
-                        generation_config=self.generation_config,
-                    )
-
-            prompt_length = prompt_inputs["input_ids"].size(1)
-            completion_ids = prompt_completion_ids[:, prompt_length:]
-            completions_refined = self.processing_class.tokenizer.batch_decode(
-                completion_ids,
-                skip_special_tokens=True,
-            )
-            refined_prompts = [self.model.extract_thinking_content(completion) for completion in completions_refined]
-
-        original_prompts = [
-            example["editing_instruction"]
-            for example in inputs
-            for _ in range(self.num_skip_refinement)
-        ]
-        all_prompts: List[str] = []
-        for batch_idx in range(batch_size):
-            refined_start = batch_idx * self.num_refined
-            refined_end = refined_start + self.num_refined
-            all_prompts.extend(refined_prompts[refined_start:refined_end])
-
-            original_start = batch_idx * self.num_skip_refinement
-            original_end = original_start + self.num_skip_refinement
-            all_prompts.extend(original_prompts[original_start:original_end])
-
-        all_source_images = [image for image in source_images for _ in range(self.num_generations)]
-        with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model:
-            with torch.no_grad():
-                (
-                    edited_images,
-                    prev_latents,
-                    diff_sampling_log_probs,
-                    pred_latents,
-                    timesteps,
-                    batched_states,
-                ) = unwrapped_model.generate_image(
-                    images=all_source_images,
-                    texts=all_prompts,
-                    diffusion_kwargs=self.diffusion_generation_config,
-                    sde_sampling=True,
-                )
-
-        rewards, rewards_per_func = self.compute_rewards(inputs, edited_images, completions_refined)
-        advantages = self.compute_advantages(rewards)
-        advantages_refined = (
-            advantages.view(batch_size, self.num_generations)[:, : self.num_refined].flatten()
-            if self.num_refined > 0
-            else torch.tensor([], device=advantages.device)
-        )
-
-        self._metrics["reward"].append(self.accelerator.gather_for_metrics(rewards).mean().item())
-        for index, (func_name, _, _) in enumerate(self.reward_funcs):
-            self._metrics[f"reward/{func_name}"].append(
-                self.accelerator.gather_for_metrics(rewards_per_func[:, index]).mean().item()
-            )
-        self._log_samples(source_images, edited_images, all_prompts, advantages)
-
-        return {
-            "images": edited_images,
-            "prev_latents": prev_latents,
-            "diff_sampling_log_probs": diff_sampling_log_probs,
-            "pred_latents": pred_latents,
-            "batched_states": batched_states,
-            "prompt_length": prompt_length,
-            "completion_ids": completion_ids,
-            "prompt_completion_ids": prompt_completion_ids,
-            "prompt_inputs": prompt_inputs,
-            "advantages": advantages,
-            "advantages_refined": advantages_refined,
-            "ts": timesteps,
-        }
-
-    def compute_rewards(
-        self,
-        inputs: List[Dict],
-        edited_images: List[Image.Image],
-        completions_refined: List[str],
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        device = self.accelerator.device
-        rewards_per_func = torch.zeros(len(edited_images), len(self.reward_funcs), device=device)
-        batch_size = len(inputs)
-
-        for index, (func_name, _, reward_func) in enumerate(self.reward_funcs):
-            if func_name == "format":
-                refined_scores = torch.tensor(reward_func(completions_refined), device=device, dtype=torch.float32)
-                for batch_idx in range(batch_size):
-                    start = batch_idx * self.num_generations
-                    refined_start = batch_idx * self.num_refined
-                    refined_end = refined_start + self.num_refined
-                    batch_refined_scores = refined_scores[refined_start:refined_end]
-                    rewards_per_func[start : start + self.num_refined, index] = batch_refined_scores
-                    rewards_per_func[start + self.num_refined : start + self.num_generations, index] = (
-                        batch_refined_scores.mean() if len(batch_refined_scores) else 0.0
-                    )
-            elif func_name == "editreward":
-                source_images = [example["image"] for example in inputs for _ in range(self.num_generations)]
-                prompts = [example["editing_instruction"] for example in inputs for _ in range(self.num_generations)]
-                rewards_per_func[:, index] = torch.tensor(
-                    reward_func(source_images, edited_images, prompts)["scores"],
-                    device=device,
-                    dtype=torch.float32,
-                )
-            else:
-                raise ValueError(f"Unsupported reward function for edit joint training: {func_name}")
-
-        return rewards_per_func.sum(dim=1), rewards_per_func
-
-    def compute_advantages(self, rewards: torch.Tensor) -> torch.Tensor:
-        grouped_rewards = rewards.view(-1, self.num_generations)
-        mean = grouped_rewards.mean(dim=1).repeat_interleave(self.num_generations, dim=0)
-        std = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(self.num_generations, dim=0)
-        return torch.clamp((rewards - mean) / (std + 1e-4), -5, 5)
-
-    def cot_loss_computation(
-        self,
-        model: nn.Module,
-        input_ids: torch.Tensor,
-        completion_ids: torch.Tensor,
-        prompt_length: int,
-        advantages: torch.Tensor,
-        prompt_inputs: Dict[str, torch.Tensor],
-    ) -> torch.Tensor:
-        image_kwargs = {
-            key: value for key, value in prompt_inputs.items() if key not in {"input_ids", "attention_mask"}
-        }
-        per_token_logps = self._get_per_token_logps(model, input_ids, image_kwargs)[:, prompt_length - 1 :]
-        with torch.inference_mode():
-            ref_per_token_logps = self._get_per_token_logps(self.ref_model, input_ids, image_kwargs)[:, prompt_length - 1 :]
-
-        per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1
-        completion_mask = self._completion_mask(completion_ids)
-        per_token_loss = torch.exp(per_token_logps - per_token_logps.detach()) * advantages.unsqueeze(1)
-        per_token_loss = -(per_token_loss - 0.01 * per_token_kl)
-        cot_loss = ((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp_min(1)).mean()
-        mean_kl = ((per_token_kl * completion_mask).sum(dim=1) / completion_mask.sum(dim=1).clamp_min(1)).mean()
-
-        self._metrics["completion_length"].append(
-            self.accelerator.gather_for_metrics(completion_mask.sum(1)).float().mean().item()
-        )
-        self._metrics["cot_kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
-        self._metrics["cot_loss"].append(self.accelerator.gather_for_metrics(cot_loss).mean().item())
-        return cot_loss
-
-    def _get_per_token_logps(
-        self,
-        model: nn.Module,
-        input_ids: torch.Tensor,
-        image_kwargs: Dict[str, torch.Tensor],
-    ) -> torch.Tensor:
-        logits = model(input_ids, **image_kwargs).logits[:, :-1, :]
-        target_ids = input_ids[:, 1:]
-        per_token_logps = []
-        for logits_row, target_ids_row in zip(logits, target_ids):
-            log_probs = logits_row.log_softmax(dim=-1)
-            per_token_logps.append(torch.gather(log_probs, dim=1, index=target_ids_row.unsqueeze(1)).squeeze(1))
-        return torch.stack(per_token_logps)
-
-    def _completion_mask(self, completion_ids: torch.Tensor) -> torch.Tensor:
-        is_eos = completion_ids == self.processing_class.eos_token_id
-        device = completion_ids.device
-        eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device)
-        eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
-        sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1)
-        return (sequence_indices <= eos_idx.unsqueeze(1)).int()
-
-    def diffusion_loss_computation(
-        self,
-        prev_latents: torch.Tensor,
-        diff_sampling_log_probs: torch.Tensor,
-        pred_latents: torch.Tensor,
-        timesteps: torch.Tensor,
-        batched_states: Dict[str, torch.Tensor],
-        advantages: torch.Tensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        model_pred = self.model.get_model().transformer(
-            hidden_states=prev_latents.to(self.model.device),
-            **batched_states,
-            joint_attention_kwargs={},
-            return_dict=False,
-        )[0][:, : pred_latents.size(1)]
-
-        with torch.no_grad():
-            ref_model_pred = self.ref_model.get_model().transformer(
-                hidden_states=prev_latents.to(self.model.device),
-                **batched_states,
-                joint_attention_kwargs={},
-                return_dict=False,
-            )[0][:, : pred_latents.size(1)]
-
-        _, log_prob, prev_sample_mean, std_dev_t = compute_log_prob(
-            model_pred,
-            self.scheduler,
-            prev_latents[:, : pred_latents.size(1)],
-            pred_latents,
-            timesteps,
-        )
-        _, _, ref_prev_sample_mean, ref_std_dev_t = compute_log_prob(
-            ref_model_pred,
-            self.scheduler,
-            prev_latents[:, : pred_latents.size(1)],
-            pred_latents,
-            timesteps,
-        )
-        if not torch.equal(std_dev_t, ref_std_dev_t):
-            raise RuntimeError("Current and reference SDE std-dev tensors diverged.")
-
-        kl = ((prev_sample_mean - ref_prev_sample_mean) ** 2 / (2 * std_dev_t**2)).mean(
-            dim=tuple(range(1, prev_sample_mean.ndim))
-        )
-        ratio = torch.exp(log_prob - diff_sampling_log_probs)
-        unclipped_loss = -advantages * ratio
-        clipped_loss = -advantages * torch.clamp(ratio, 1.0 - 1e-4, 1.0 + 1e-4)
-        diff_loss = torch.maximum(unclipped_loss, clipped_loss).mean() + self.beta * kl.mean()
-        return diff_loss, kl
-
-    def _log_samples(
-        self,
-        source_images: List[Image.Image],
-        edited_images: List[Image.Image],
-        prompts: List[str],
-        advantages: torch.Tensor,
-    ) -> None:
-        global_step = self.state.global_step
-        if global_step % 10 != 0 or not edited_images:
-            return
-
-        device_id = str(self.accelerator.device).replace(":", "")
-        text_content = []
-        for batch_idx in range(len(source_images)):
-            for gen_idx in range(self.num_generations):
-                overall_idx = batch_idx * self.num_generations + gen_idx
-                status = "REFINED" if gen_idx < self.num_refined else "ORIGINAL"
-                text_content.append(f"[{status}] Generation {gen_idx}: {prompts[overall_idx]}")
-            text_content.append("")
-
-        txt_path = os.path.join(self.log_dir, f"step_{global_step}_{device_id}.txt")
-        if not os.path.exists(txt_path):
-            with open(txt_path, "w", encoding="utf-8") as file:
-                file.write("\n".join(text_content))
-
-        for batch_idx, source_image in enumerate(source_images):
-            source_image.save(os.path.join(self.log_dir, f"step_{global_step}_{device_id}_batch{batch_idx}_source.jpg"))
-            for gen_idx in range(self.num_generations):
-                overall_idx = batch_idx * self.num_generations + gen_idx
-                prefix = "refined" if gen_idx < self.num_refined else "original"
-                edited_images[overall_idx].save(
-                    os.path.join(
-                        self.log_dir,
-                        f"step_{global_step}_{device_id}_batch{batch_idx}_{prefix}_gen{gen_idx}_{advantages[overall_idx].item():.5f}.jpg",
-                    )
-                )
-
-    def log(self, logs: Dict[str, float], start_time: Optional[float] = None) -> None:
-        metrics = {key: sum(value) / len(value) for key, value in self._metrics.items() if value}
-        logs = {**logs, **metrics}
-        if version.parse(transformers.__version__) >= version.parse("4.47.0.dev0"):
-            super().log(logs, start_time)
-        else:
-            super().log(logs)
-        self._metrics.clear()
-