tennyyyin/flux2-klein-multiview-completion

Multiview image completion — FLUX.2-klein-4B full fine-tune

Fills in a masked region of one camera view using a second, unmasked view of the same scene at the same instant. Trained on robot manipulation data with paired side-view and wrist-camera images.

This repo holds the transformer only (step 9000, weights without optimizer state). The VAE and text encoder are unchanged from black-forest-labs/FLUX.2-klein-4B.

Task

Input	Meaning
cond_image_1	reference view, unmasked
cond_image_2	target view, masked (removed pixels are WHITE)
output	target view, completed

Prompts are per-row templated instructions describing which camera is which and what to fill in.

Training

Full fine-tune (not LoRA) of all 169 transformer tensors, bf16.

Base	black-forest-labs/FLUX.2-klein-4B
Data	28,708 train rows / 758 held-out test rows
Resolution	320x640, single aspect-ratio bucket
Steps	9,000
Batch size	16
LR	1e-5, constant with 200-step warmup
Optimizer	8-bit AdamW
Hardware	8x H200 (ml.p5en.48xlarge), 9.76 h
Max sequence length	176

Results

Held-out test set (all 758 rows, fixed seed), latent-space metrics:

Step	test_loss	latent_x0_mse	latent_x0_psnr
225	0.4099	0.1167	9.33
2250	0.3792	0.1083	9.65
4500	0.3735	0.1064	9.73
6750	0.3708	0.1052	9.78
9000	0.3692	0.1046	9.80

Improvement is monotone through step 9000, so this is the best checkpoint — later steps would likely still help. Gains flatten considerably after ~4,000 steps.

Perceptual metrics (LPIPS 0.41-0.42, PSNR ~15) were measured on only 3 fixed validation images, so treat them as a smoke signal rather than an evaluation; they fluctuate between consecutive evals by more than the trend.

Samples

samples/validation_{0,1,2}.png are this checkpoint's renders of the three fixed validation triplets, at step 9000.

Usage

import torch
from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel

transformer = Flux2Transformer2DModel.from_pretrained(
    "tennyyyin/flux2-klein-multiview-completion",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-4B",
    transformer=transformer,
    torch_dtype=torch.bfloat16,
).to("cuda")

Pass the reference view and the white-masked target view as the two conditioning images, with an instruction prompt describing the camera geometry and what to complete.

Limitations

• Trained at a single 320x640 aspect ratio; other shapes are untested.
• Domain is robot tabletop manipulation (side-view + wrist camera). Generalization to other multiview settings is unmeasured.
• No pixel-space or perceptual evaluation on a meaningful sample of held-out data.