unpaired-edit-sd35-large-stage2

unpaired-edit-sd35-large-stage2

Stage-2 (NFT + gating) LoRA adapter and gating network for unpaired instruction-based image editing on top of Stable Diffusion 3.5 Large.

This is an intermediate training checkpoint, not a finished model. It was exported at global_step = 14500 (~epoch 906) of a run whose configured schedule is far longer; training had not converged when this snapshot was taken.

Contents

Optimizer, EMA and RNG state are not included, so this snapshot is for inference/evaluation only and cannot be used to resume training.

Requirements

This adapter is not standalone. To use it you also need:

• stabilityai/stable-diffusion-3.5-large (base model)
• a stage-1 SD3.5-Large ControlNet checkpoint, which the gating network modulates

Training configuration

• LoRA: rank 128, alpha 128, dropout 0.0
• Target modules: to_q, to_k, to_v, to_out.0, add_q_proj, add_k_proj, add_v_proj, to_add_out
• Precision: bf16
• Reward model: Qwen/Qwen3-VL-4B-Instruct (bf16)
• Objective: NFT with edit + preservation rewards (lambda_edit = 1.0, lambda_preservation = 1.0), KL beta = 0.01, nft_beta = 0.5, EMA enabled
• Gating: default scale 0.7, trained jointly with the transformer LoRA

Training metrics at export

Values logged by the training loop at the epochs around this checkpoint. These are in-training reward-model scores on sampled batches, not results on any held-out benchmark:

• reward_edit_score ≈ 0.80–0.91
• reward_preservation_score ≈ 0.51–0.62
• gate_mean ≈ 0.717

Loading

import torch
from diffusers import SD3Transformer2DModel
from peft import PeftModel

transformer = SD3Transformer2DModel.from_pretrained(
    "stabilityai/stable-diffusion-3.5-large",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)
transformer = PeftModel.from_pretrained(
    transformer, "phanviethoang1512/unpaired-edit-sd35-large-stage2"
)

# gating network weights (architecture comes from the training repo)
gating_state = torch.load("gating_net.pt", map_location="cpu")

License

Inherits the Stability AI Community License of the base model; review those terms before use.