RawGen: Learning Camera Raw Image Generation

Dongyoung Kim¹, Junyong Lee^1*, Abhijith Punnappurath^1*, Mahmoud Afifi^1*, Sangmin Han², Alex Levinshtein¹, Michael S. Brown¹

¹AI Center – Toronto, Samsung Electronics    ²Yonsei University    ^*Equal contribution

[Paper] [Project page]

Official implementation of the ECCV 2026 paper.

Overview

RawGen maps a photo-finished sRGB image or a text prompt back to a standardized scene-referred CIE XYZ representation, then renders that XYZ as camera RAW. The XYZ intermediate is camera-agnostic, so one prediction can be re-rendered for any camera whose colour calibration is known.

sRGB image ─┐
            ├─► VAE encode ─► Kontext LoRA denoise ─► fine-tuned VAE decoder ─► XYZ ─┐
text prompt ┘        (variant → anchor, in latent space)                             │
                                                                                     ▼
                                              9 cross-camera RAW PNGs + Samsung S20 / S25U DNGs

Two fine-tuned components sit on top of frozen FLUX.1-Kontext-dev:

• a LoRA adapter on the Kontext transformer, learning the variant→anchor mapping in latent space (rank 64, on to_q / to_k / to_v / to_out.0);
• a fully fine-tuned VAE decoder that decodes anchor latents to CIE XYZ instead of sRGB.

Requirements

A single CUDA GPU. Inference loads the FLUX transformer, VAE and T5/CLIP text encoders together and peaks at ~37 GB of VRAM (measured on an H100 for image-to-RAW at 1024², 30 steps), so a ≥40 GB GPU is recommended. bfloat16 throughout.

pip install -r requirements.txt

# system dependency: ExifTool (PyExifTool shells out to it to read DNG metadata)
#   Ubuntu:  sudo apt-get install -y libimage-exiftool-perl
#   macOS:   brew install exiftool

requirements.txt pins diffusers exactly because the CLIs depend on internals of the FLUX pipeline, so do not upgrade past the pin. The bottom section of the file lists training-only extras that inference does not need.

The base models black-forest-labs/FLUX.1-Kontext-dev and (for text-to-RAW) black-forest-labs/FLUX.1-dev are gated. Accept each model's license on its model page, then authenticate (hf auth login, or export HF_TOKEN) before the first run; diffusers fetches them automatically from there.

Pretrained models

The two fine-tuned checkpoints are not committed to this repository. Download them from Google Drive and place them under checkpoints/, which is where the CLIs look by default:

checkpoints/
├── kontext_lora.pt
└── vae_decoder_xyz.pt

File	Size	Contents
checkpoints/kontext_lora.pt	143 MB	LoRA adapter for the Kontext transformer
checkpoints/vae_decoder_xyz.pt	95 MB	fine-tuned VAE decoder

Pass --transformer-ckpt / --vae-decoder-ckpt to load them from elsewhere.

The LoRA file holds only the adapter; FLUX base weights are fetched at run time and are not redistributed here. Both checkpoints are governed by the FLUX.1 [dev] Non-Commercial License — see LICENSE.md and NOTICE.

Usage

Image-to-RAW

Run both CLIs from the repository root; their default paths (checkpoints/, samples/) are relative to it.

samples/inputs/ ships empty — drop your own sRGB .jpg / .png images there, or point --input at any file or directory. Inputs are resized to a square --size (default 1024²) because that is the training resolution, so pre-crop non-square photos if you want the geometry preserved.

python image_to_raw.py \
    --input samples/inputs/ \
    --output-dir out/image_to_raw/ \
    --num-illum 2 \
    --num-steps 30 \
    --seed 42

Text-to-RAW

python text_to_raw.py \
    --prompts-file samples/prompts.txt \
    --output-dir out/text_to_raw/

Text-to-RAW runs two diffusion stages, so its step and guidance flags are split per stage; there is no --num-steps here.

Options

Shared by both CLIs:

Flag	Default	Effect
--num-illum	2	How many illuminations to sample per image; each one produces its own set of camera outputs. Illuminants are drawn at random from samples/illuminations.json, so each output carries a different colour cast from the input by design — the one used is recorded in metadata.json
--kontext-guidance	3.5	Guidance for the Kontext denoise stage. 3.5 is the trained value — changing it moves you off the training distribution
--size	1024	Square resolution. 1024 is the training resolution
--seed	42	Seeds both the illumination sampling and the initial denoise noise, so a rerun reproduces the outputs
--transformer-ckpt / --vae-decoder-ckpt	checkpoints/…	Load the two checkpoints from elsewhere
--illum-json / --dng-dir	samples/…	Illumination set and camera-profile DNGs

Image-to-RAW only: --num-steps (30) denoise steps, --save-jpg additionally writes 8-bit JPEG previews next to the RAW PNGs.

Text-to-RAW only: --t2i-steps (30) for the FLUX.1-dev text-to-image stage and --kontext-steps (30) for the Kontext stage; --guidance-scale (3.5) for the text-to-image stage; --batch-size (2) prompts per forward pass.

Output layout

out/<task>/
├── xyz/<name>_xyz.png                    16-bit OETF-encoded XYZ
├── png/<name>_<camera>_illum<i>.png      cross-camera RAW, 9 NUS cameras
├── dng/<name>_S20_illum<i>.dng           Samsung S20 RAW (Bayer)
├── dng/<name>_S25U_illum<i>.dng          Samsung S25 Ultra RAW (ProRAW)
└── metadata.json                         checkpoints, seed, per-camera gt_illum

The xyz/ image is sRGB-OETF-encoded, not linear. Apply the inverse OETF before using it in any linear-domain computation (this is what xyz_to_raw.py does internally before the camera transform).

<camera> is one of the nine NUS bodies (NikonD40, Canon1DsMkIII, …) fixed in the CLIs, each matched to its profile DNG in samples/dng_profiles/; adding a DNG does not add a camera. <i> indexes the sampled illuminations (0 to --num-illum - 1). The png/ files are the camera's RAW RGB rescaled to [0,1] by that camera's black/white level, gamma-encoded for viewing; they are previews, not loadable RAW. The two dng/ files are real RAW containers. metadata.json records the checkpoints, seed, and the illumination vector used for each camera.

Text-to-RAW indexes outputs by prompt rather than by input filename, so <name> above is <NNNN>_<slug>; it also writes the prompt to text_prompt/<NNNN>_<slug>.txt.

S25U output is Samsung ProRAW (JPEG-XL compressed, linear 3-channel). Read it with tifffile + imagecodecs, or with Adobe software; rawpy / LibRaw do not support that compression.

Training

Training takes four steps. All multi-GPU scripts use torchrun. Run them from training/.

The commands below are minimal, runnable examples on a single dataset subset, not the runs that produced the released checkpoints. The settings behind those are listed under Released checkpoint settings.

Dataset manifest

Both data preparation and training read a trainval.json manifest:

{
  "train": [ {"dataset_type": "a5k", "basename": "a0001-..."}, ... ],
  "val":   [ ... ],
  "test":  [ ... ],
  "meta":  { "root": "/abs/path/to/data",
             "srgb_suffix": "_srgb.png", "xyz_suffix": "_xyz.png" }
}

--root-key selects which meta.* root to resolve against. --split all concatenates train+val+test; --datasets a5k raise (or all) selects subsets by dataset_type.

The manifest is not shipped. There is no generator script; write the JSON above yourself from the basenames you have on disk. The released checkpoints used an 8:1:1 train/val/test split.

Getting the source data

Neither dataset is redistributed here. Both are free for research use but require their own registration/download:

dataset_type	Dataset	Ships as	Prepare by
a5k	MIT-Adobe FiveK	.dng	use as-is
raise	RAISE	.nef (Nikon raw)	convert to DNG (e.g. Adobe DNG Converter)

dataset_type is the directory name, and basename is the DNG filename without its extension, so a manifest row {"dataset_type": "a5k", "basename": "a0104-dvf_003"} resolves to <meta.root>/a5k/DNG/a0104-dvf_003.dng. The released manifest keeps each dataset's native naming (FiveK a0104-dvf_003, RAISE rc5ae823dt); any naming works as long as the manifest and the filenames agree.

Directory layout

<meta.root>/
├── a5k/
│   ├── DNG/<basename>.dng          # you provide — the only input that must pre-exist
│   ├── anchor-variants-imgs/       # Step 1 writes here
│   └── latents/                    # Step 2 writes here
└── raise/
    └── ... same three

A manifest entry may instead give explicit absolute "srgb" / "xyz" paths per sample; they still have to follow the <root>/<dataset_type>/<subdir>/ shape, since the DNG path and the dataset type are derived from them.

Keeping the four steps consistent

Steps 2-4 locate the earlier outputs by directory name, so the same name has to be repeated:

What	Step 1	Step 2	Step 3	Step 4
variant/anchor images	--output-dir	--variant-dir-name	—	--gt-xyz-dir-name
latents	—	--output-dir	--latent-dir-name	--latent-dir-name
variants per anchor	--num-variations	--num-variations	--num-variations	—

The example commands below use anchor-variants-imgs and latents throughout.

Step 1 — ISP variants and anchors from DNGs

Synthesizes N sRGB variants plus the XYZ / sRGB anchors from each source DNG, using the vendored mini ISP simulator. The XYZ anchor runs the ISP through the xyz stage and then the sRGB OETF (--xyz-anchor-gamma, on by default); pass --no-xyz-anchor-gamma for linear anchors, which requires retraining both stages and a matching change on the inference side. The paper's reconstruction metrics were computed in this encoded space.

python generate_variants_and_anchors.py \
    --trainval-json /path/to/trainval.json --split all --datasets a5k \
    --output-dir anchor-variants-imgs --num-variations 5 --cpu-workers 24 \
    --generate-xyz-anchor

--output-dir here is a directory name, not a path: images are written to <manifest root>/<dataset_type>/anchor-variants-imgs/. The later steps locate them by the same name, so keep it consistent across all four commands.

Step 2 — VAE-encode to latents

Pre-encodes the Step 1 images so training iterates without rerunning the VAE. Re-runnable with a different VAE without redoing Step 1.

torchrun --nproc_per_node=4 generate_latents_for_variants_and_anchors.py \
    --model-id black-forest-labs/FLUX.1-dev \
    --trainval-json /path/to/trainval.json --split all --datasets a5k \
    --variant-dir-name anchor-variants-imgs \
    --output-dir latents --num-variations 5 --batch-size 8 \
    --generate-xyz-anchor

Reads the Step 1 images from anchor-variants-imgs/ and writes latents to <manifest root>/<dataset_type>/latents/ (only the last path component of --output-dir is used).

Produces {basename}_anchor_srgb.pt, {basename}_anchor_xyz.pt and {basename}_var_{ii:02d}.pt per basename. Encoding uses latent_dist.mode(), so the cache is deterministic.

Step 3 — Kontext diffusion LoRA

torchrun --nproc_per_node=4 train_kontext_lora.py \
    --model-id black-forest-labs/FLUX.1-Kontext-dev \
    --trainval-json /path/to/trainval.json --datasets a5k \
    --latent-dir-name latents --root-key root \
    --batch-size 4 --lr 1e-4 --max-epochs 40 --save-every 5

The saved .pt holds the LoRA adapter (no FLUX base weights) plus optimizer state for resuming, so it is larger than the released file; the CLIs read only the adapter.

Step 4 — VAE decoder

torchrun --nproc_per_node=4 finetune_vae_decoder.py \
    --model-id black-forest-labs/FLUX.1-dev \
    --trainval-json /path/to/trainval.json --datasets a5k \
    --latent-dir-name latents --gt-xyz-dir-name anchor-variants-imgs \
    --amp --batch-size 8 --lr 1e-4 --max-epochs 30 --save-every 10

--amp runs the decode under autocast and is what the released run used. The decoder is trained on the FLUX.1-dev VAE, whose architecture matches FLUX.1-Kontext-dev's, so it loads onto the Kontext VAE at inference, which is what the CLIs do.

Using your own checkpoints

Both trainers write to a timestamped run directory under --results-root (default ./results/): Step 3 to <run>/kontext_lora_ckpt/kontext_lora_epoch{N}.pt (plus _final), Step 4 to <run>/vae_decoder_xyz_ckpt_finetune/vae_decoder_xyz_epoch{N}.pt (plus _best and _final). Pass either file to --transformer-ckpt / --vae-decoder-ckpt, or copy them to checkpoints/kontext_lora.pt and checkpoints/vae_decoder_xyz.pt to use the defaults.

Released checkpoint settings

Both released checkpoints were trained on 8 GPUs over both dataset subsets (--datasets all), with 5 ISP variants per anchor.

kontext_lora.pt — 40 epochs, LoRA rank 64 / alpha 64 on to_q, to_k, to_v, to_out.0, bfloat16:

torchrun --nproc_per_node=8 train_kontext_lora.py \
    --model-id black-forest-labs/FLUX.1-Kontext-dev \
    --trainval-json /path/to/trainval.json --datasets all \
    --latent-dir-name latents --num-variations 5 \
    --amp --gradient-checkpointing \
    --batch-size 1 --grad-accum-steps 2 --lr 1e-4 --warmup-steps 200 \
    --max-epochs 40 --save-every 5

vae_decoder_xyz.pt — 30 epochs, full fine-tune, L1 loss, tighter gradient clipping:

torchrun --nproc_per_node=8 finetune_vae_decoder.py \
    --model-id black-forest-labs/FLUX.1-dev \
    --trainval-json /path/to/trainval.json --datasets all \
    --latent-dir-name latents --gt-xyz-dir-name anchor-variants-imgs \
    --amp --dtype bfloat16 \
    --clip-grad-norm 0.5 --max-epochs 30

Acknowledgements

The low-level DNG and camera-pipeline code is vendored from graphics2raw and simple-camera-pipeline; each vendored file records its origin and license in its header. The nine reference camera profiles come from the NUS Illumination Dataset (Cheng et al.). See LICENSE.md for provenance details.

Citation

@inproceedings{kim2026rawgen,
  title     = {RawGen: Learning Camera Raw Image Generation},
  author    = {Dongyoung Kim and Junyong Lee and Abhijith Punnappurath and
               Mahmoud Afifi and Sangmin Han and Alex Levinshtein and
               Michael S. Brown},
  booktitle = {European Conference on Computer Vision (ECCV)},
  year      = {2026}
}

License

Code written for RawGen is CC BY-NC-SA 4.0. Vendored third-party code keeps its own upstream license, and the model weights are governed by the FLUX.1 [dev] Non-Commercial License. See LICENSE.md for the per-path breakdown and NOTICE for the weight attribution.