Upload folder using huggingface_hub

.gitattributes

@@ -115,3 +115,4 @@ CKPTs_qwen/sample_set_best64_maxnegdraw_fixeddraw_32_lr/tar_cache/lindau_000020_
 CKPTs_qwen/sample_set_best64_maxnegdraw_fixeddraw_32_lr/tar_cache/sa_3512379.jpg filter=lfs diff=lfs merge=lfs -text
 CKPTs_qwen/sample_set_best64_maxnegdraw_fixeddraw_32_lr/tar_cache/sa_8539606.jpg filter=lfs diff=lfs merge=lfs -text
 ovb_prepared_semseg_instanceseg_v4_objects365_adddriveds/semseg/ filter=lfs diff=lfs merge=lfs -text
+CKPTS/training_metrics/metrics.jsonl filter=lfs diff=lfs merge=lfs -text

CKPTS/README.md

@@ -0,0 +1,68 @@
+---
+base_model: black-forest-labs/FLUX.2-klein-base-9B
+library_name: diffusers
+license: other
+pipeline_tag: image-to-image
+instance_prompt: '[TASK: IMAGE_EDITING] Edit the image according to the instruction.'
+widget: []
+tags:
+- image-to-image
+- image-editing
+- diffusers-training
+- diffusers
+- flux2
+- flux2-diffusers
+- partial-finetune
+---
+
+<!-- This model card has been generated automatically according to the information the training script had access to. You
+should probably proofread and complete it, then remove this comment. -->
+
+
+# Flux.2 [Klein] partial transformer fine-tune - ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000
+
+<Gallery />
+
+## Model description
+
+These are ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000 partial transformer fine-tuning weights for `black-forest-labs/FLUX.2-klein-base-9B`.
+
+The weights were trained using an image-to-image Flux2-Klein trainer that freezes the base model except for selected late transformer layers and output layers.
+
+FP8 training? False
+
+## Trigger words
+
+You should use `[TASK: IMAGE_EDITING] Edit the image according to the instruction.` to trigger the image generation or image-editing behavior.
+
+## Use it with diffusers
+
+```py
+from diffusers import Flux2KleinPipeline, Flux2Transformer2DModel
+import torch
+
+transformer = Flux2Transformer2DModel.from_pretrained("arpitjadon/OVB", subfolder="CKPTS/transformer", torch_dtype=torch.bfloat16)
+pipeline = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-base-9B", transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
+image = pipeline("[TASK: IMAGE_EDITING] Edit the image according to the instruction.").images[0]
+```
+
+## License
+
+Please adhere to the licensing terms as described [here](https://huggingface.co/black-forest-labs/FLUX.2/blob/main/LICENSE.md).
+
+
+## Intended uses & limitations
+
+#### How to use
+
+```python
+# TODO: add an example code snippet for running this diffusion pipeline
+```
+
+#### Limitations and bias
+
+[TODO: provide examples of latent issues and potential remediations]
+
+## Training details
+
+[TODO: describe the data used to train the model]

CKPTS/checkpoint-37500/optimizer.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71b33f9a1b6a83cb51904ee455d36173969e089b73c783a7da3040b96f47edba
+size 72628875483

CKPTS/checkpoint-37500/pytorch_model_fsdp.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:17ca0722be9cfa2a8381842809fafa84e3006177349e20dd90b3127cded7c378
+size 36314461302

CKPTS/checkpoint-37500/random_states_0.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb3f4773bffbe54fd13e2905c2778e8ed2eb31d2df761e25317750157b59c473
+size 16513

CKPTS/checkpoint-37500/random_states_1.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:439347e8ec7bc990b4555af0e0cb431b893daf96b1b596e3b038fb2fe1102663
+size 16513

CKPTS/checkpoint-37500/random_states_2.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c21dc27a2cb875787b2ae50e2eb6c313b42c48c300f3ede95bdb41602e1ea4c5
+size 16513

CKPTS/checkpoint-37500/random_states_3.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a03efff8e8e3961b1a5cbfcead76c650edfbe5ff150d423c30a2989a86e8f411
+size 16513

CKPTS/checkpoint-37500/random_states_4.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:be238041d82b250a47a2c69560aebb669574ac5260bd584b278195a8505a60e2
+size 16513

CKPTS/checkpoint-37500/random_states_5.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:997e15b6e3708a90ededd2c11df328e110eaef5242fc1df4848c9ccf471047c0
+size 16513

CKPTS/checkpoint-37500/random_states_6.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b3b7294cf037f9976fee41d5d50ee030eb0f01984debb5ac341a4784067de0a
+size 16513

CKPTS/checkpoint-37500/random_states_7.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5fcfa8807619d9008ee0fbfc48ee106f59d94ec47990e5c13c8bcd0ca471986
+size 16513

CKPTS/checkpoint-37500/sampler_state.json

@@ -0,0 +1,8 @@
+{
+  "batches_consumed_in_epoch": 2346,
+  "dataset_pool_seed": 42,
+  "epoch": 9,
+  "global_step": 37500,
+  "gradient_accumulation_steps": 1,
+  "schema_version": 1
+}

CKPTS/checkpoint-37500/scheduler.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:041feb26c54cf05ae8f8b921f1868f76e620f57af2cc12e6d9c5488a31f514aa
+size 1465

CKPTS/checkpoint-37500/transformer/config.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "Flux2Transformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "_name_or_path": "black-forest-labs/FLUX.2-klein-base-9B",
+  "attention_head_dim": 128,
+  "axes_dims_rope": [
+    32,
+    32,
+    32,
+    32
+  ],
+  "eps": 1e-06,
+  "guidance_embeds": false,
+  "in_channels": 128,
+  "joint_attention_dim": 12288,
+  "mlp_ratio": 3.0,
+  "num_attention_heads": 32,
+  "num_layers": 8,
+  "num_single_layers": 24,
+  "out_channels": null,
+  "patch_size": 1,
+  "rope_theta": 2000,
+  "timestep_guidance_channels": 256
+}

CKPTS/checkpoint-37500/transformer/diffusion_pytorch_model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbb926f875e52a38fc40660a3ec0b37f4fd2d75e85624279bbcca926317dc116
+size 9804206048

CKPTS/checkpoint-37500/transformer/diffusion_pytorch_model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:372b1dc97c6f27fae073a0cfbc3913236c8bba56579aab1f90ded850443d4d7d
+size 9797913704

CKPTS/checkpoint-37500/transformer/diffusion_pytorch_model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd117aa0a5ba2dc606ece70c5c7ada7932f4a2bd65e5d2c15312bfe6d82fc6c7
+size 9596584288

CKPTS/checkpoint-37500/transformer/diffusion_pytorch_model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cf45bdc042c4f9712deaed5f00083dc86a03e6caf53981d1bf8ec174f43bb69
+size 7115649136

CKPTS/checkpoint-37500/transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,240 @@
+{
+  "metadata": {
+    "total_size": 36314324992
+  },
+  "weight_map": {
+    "context_embedder.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "double_stream_modulation_img.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "double_stream_modulation_txt.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "norm_out.linear.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "proj_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_stream_modulation.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "time_guidance_embed.timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "time_guidance_embed.timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "x_embedder.weight": "diffusion_pytorch_model-00001-of-00004.safetensors"
+  }
+}

CKPTS/dataset_specs.resolved.json

@@ -0,0 +1,69 @@
+{
+  "datasets": [
+    {
+      "filters": null,
+      "index_path": null,
+      "name": "semseg",
+      "num_rows": 172015,
+      "path": "/gpfs/haicore/projects/haicore470/hpc26/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/semseg",
+      "sampling": "row",
+      "type": "prepared_folder",
+      "weight": 0.247525
+    },
+    {
+      "filters": null,
+      "index_path": null,
+      "name": "pico_banana",
+      "num_rows": 156670,
+      "path": "/gpfs/haicore/projects/haicore470/hpc26/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/pico_banana_sft_150k",
+      "sampling": "row",
+      "type": "prepared_folder",
+      "weight": 0.118812
+    },
+    {
+      "filters": null,
+      "index_path": null,
+      "name": "identity",
+      "num_rows": 50000,
+      "path": "/gpfs/haicore/projects/haicore470/hpc26/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/identity_50k",
+      "sampling": "row",
+      "type": "prepared_folder",
+      "weight": 0.079207
+    },
+    {
+      "filters": null,
+      "index_path": null,
+      "name": "refseg",
+      "num_rows": 144947,
+      "path": "/gpfs/haicore/projects/haicore470/hpc26/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/refseg",
+      "sampling": "row",
+      "type": "prepared_folder",
+      "weight": 0.188119
+    },
+    {
+      "filters": null,
+      "index_path": null,
+      "name": "instanceseg",
+      "num_rows": 139580,
+      "path": "/gpfs/haicore/projects/haicore470/hpc26/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/instanceseg",
+      "sampling": "row",
+      "type": "prepared_folder",
+      "weight": 0.247525
+    },
+    {
+      "filters": null,
+      "index_path": null,
+      "name": "pd12m",
+      "num_rows": 150000,
+      "path": "/gpfs/haicore/projects/haicore470/hpc26/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/pd12m_150000",
+      "sampling": "row",
+      "type": "prepared_folder",
+      "weight": 0.118812
+    }
+  ],
+  "max_refseg_expressions_per_image": 2000000,
+  "samples_per_epoch": 250000,
+  "sampling_strategy": "without_replacement",
+  "seed": 42,
+  "shuffle": true
+}

CKPTS/logs/flux2-klein-partial-image2img/1786717770.5928092/events.out.tfevents.1786717770.gb005.2141029.1

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d88323dd62bc15015d858c4f85e2c51321f52e303ae8c073c5ef8066167cf92
+size 7665

CKPTS/logs/flux2-klein-partial-image2img/1786717770.5968676/hparams.yml

@@ -0,0 +1,129 @@
+adam_beta1: 0.9
+adam_beta2: 0.999
+adam_epsilon: 1.0e-08
+adam_weight_decay: 0.0001
+adam_weight_decay_text_encoder: 0.001
+allow_large_prompt_precompute_cache: false
+allow_legacy_hypersim_range_depth: false
+allow_tf32: true
+allow_unverified_normal_sources: false
+angular_normal_aux_weight: 0.0
+aspect_ratio_buckets: 512,512;576,448;448,576;608,416;416,608
+bnb_quantization_config_path: null
+cache_dir: null
+cache_latents: false
+caption_column: text
+center_crop: false
+checkpointing_steps: 625
+checkpoints_total_limit: 2
+class_data_dir: null
+cond_image_column: conditioning_image
+dataloader_num_workers: 8
+dataset_config_name: null
+dataset_index_num_workers: 16
+dataset_name: null
+dataset_pool_seed: 42
+dataset_specs: /home/hpc26/open_visionbanana/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/open_visionbanana_mix_v1_tmpdir.json
+depth_prompt_contract: campaign_detailed_inverse
+depth_target_encoding: relative_inverse_grayscale
+depth_target_resize_policy: metric_field_then_encode
+do_fp8_training: false
+enable_npu_flash_attention: false
+final_validation_prompt: null
+fsdp_text_encoder: false
+fsdp_transformer: true
+fsdp_transformer_cls_to_wrap: Flux2TransformerBlock,Flux2SingleTransformerBlock
+fsdp_transformer_dtype_policy: fp32_all
+gradient_accumulation_steps: 1
+gradient_checkpointing: false
+guidance_scale: 3.5
+hub_model_id: null
+hub_token: null
+image_column: image
+instance_data_dir: null
+instance_prompt: '[TASK: IMAGE_EDITING] Edit the image according to the instruction.'
+learning_rate: 2.0e-06
+local_rank: 0
+logging_dir: logs
+logit_mean: 0.0
+logit_std: 1.0
+lora_alpha: 64
+lora_dropout: 0.0
+lora_layers: null
+lora_rank: 64
+loss_mask_erosion_radius: 0
+loss_mask_min_sample_valid_fraction: 0.0
+loss_mask_min_valid_fraction: 0.5
+loss_mask_mode: hard
+loss_mask_policy: disabled
+loss_mask_reduction: per_sample
+lr_num_cycles: 1
+lr_power: 1.0
+lr_scheduler: constant_with_warmup
+lr_warmup_steps: 1000
+max_grad_norm: 1.0
+max_refseg_expressions_per_image: 2000000
+max_sequence_length: 512
+max_train_steps: 37500
+metric_depth_aux_max_depth: 80.0
+metric_depth_aux_weight: 0.0
+mixed_precision: bf16
+mode_scale: 1.29
+normal_prompt_contract: campaign_short
+normal_target_encoding: vision_banana_x_left_red
+normal_target_resize_policy: vector_field_renormalize_then_encode
+num_train_epochs: 10
+num_validation_images: 4
+offload: false
+optimizer: AdamW
+output_dir: /gpfs/haicore/projects/haicore470/hpc26/checkpoints/open_visionbanana/ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000
+packed_prompt_shard_bytes: 1073741824
+palette_manifold_loss_weight: 0.0
+permanent_checkpoint_steps: 6250,12500,18750,25000,31250,37500
+precompute_cache_dir: /gpfs/haicore/projects/haicore470/hpc26/precompute_cache/ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000
+pretrained_model_name_or_path: black-forest-labs/FLUX.2-klein-base-9B
+prodigy_beta3: null
+prodigy_decouple: true
+prodigy_safeguard_warmup: true
+prodigy_use_bias_correction: true
+prompt_embedding_batch_size: 128
+prompt_encoding_mode: packed_precompute
+prompt_precompute_max_unique: 2048
+push_to_hub: false
+random_flip: false
+repeats: 1
+report_to: tensorboard
+resample_dataset_each_epoch: false
+resolution: 512
+resume_from_checkpoint: latest
+revision: null
+sample_batch_size: 4
+samples_per_epoch: 250000
+sampling_strategy: without_replacement
+scale_lr: false
+seed: 42
+skip_final_inference: true
+task_homogeneous_batches: true
+train_all_transformer_blocks: true
+train_batch_size: 8
+train_full_transformer: false
+train_input_layers: false
+train_last_single_transformer_blocks: 1
+train_last_transformer_blocks: 0
+train_method: partial_full_rank
+train_non_block_layers: true
+train_output_layers: true
+train_single_transformer_block_indices: null
+train_transformer_block_indices: null
+unit_normal_loss_weight: 0.0
+upcast_before_saving: false
+use_8bit_adam: false
+validation_epochs: 50
+validation_height: null
+validation_image: null
+validation_num_inference_steps: 28
+validation_prompt: null
+validation_prompts_file: null
+validation_width: null
+variant: null
+weighting_scheme: none

CKPTS/logs/flux2-klein-partial-image2img/1786718917.173297/events.out.tfevents.1786718917.gb006.594302.1

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82b8475d921839981ff111d56615c51a49fa82414fc897de9270192ff16fcb11
+size 7665

CKPTS/logs/flux2-klein-partial-image2img/1786718917.1766176/hparams.yml

@@ -0,0 +1,129 @@
+adam_beta1: 0.9
+adam_beta2: 0.999
+adam_epsilon: 1.0e-08
+adam_weight_decay: 0.0001
+adam_weight_decay_text_encoder: 0.001
+allow_large_prompt_precompute_cache: false
+allow_legacy_hypersim_range_depth: false
+allow_tf32: true
+allow_unverified_normal_sources: false
+angular_normal_aux_weight: 0.0
+aspect_ratio_buckets: 512,512;576,448;448,576;608,416;416,608
+bnb_quantization_config_path: null
+cache_dir: null
+cache_latents: false
+caption_column: text
+center_crop: false
+checkpointing_steps: 625
+checkpoints_total_limit: 2
+class_data_dir: null
+cond_image_column: conditioning_image
+dataloader_num_workers: 8
+dataset_config_name: null
+dataset_index_num_workers: 16
+dataset_name: null
+dataset_pool_seed: 42
+dataset_specs: /home/hpc26/open_visionbanana/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/open_visionbanana_mix_v1_tmpdir.json
+depth_prompt_contract: campaign_detailed_inverse
+depth_target_encoding: relative_inverse_grayscale
+depth_target_resize_policy: metric_field_then_encode
+do_fp8_training: false
+enable_npu_flash_attention: false
+final_validation_prompt: null
+fsdp_text_encoder: false
+fsdp_transformer: true
+fsdp_transformer_cls_to_wrap: Flux2TransformerBlock,Flux2SingleTransformerBlock
+fsdp_transformer_dtype_policy: fp32_all
+gradient_accumulation_steps: 1
+gradient_checkpointing: false
+guidance_scale: 3.5
+hub_model_id: null
+hub_token: null
+image_column: image
+instance_data_dir: null
+instance_prompt: '[TASK: IMAGE_EDITING] Edit the image according to the instruction.'
+learning_rate: 2.0e-06
+local_rank: 0
+logging_dir: logs
+logit_mean: 0.0
+logit_std: 1.0
+lora_alpha: 64
+lora_dropout: 0.0
+lora_layers: null
+lora_rank: 64
+loss_mask_erosion_radius: 0
+loss_mask_min_sample_valid_fraction: 0.0
+loss_mask_min_valid_fraction: 0.5
+loss_mask_mode: hard
+loss_mask_policy: disabled
+loss_mask_reduction: per_sample
+lr_num_cycles: 1
+lr_power: 1.0
+lr_scheduler: constant_with_warmup
+lr_warmup_steps: 1000
+max_grad_norm: 1.0
+max_refseg_expressions_per_image: 2000000
+max_sequence_length: 512
+max_train_steps: 37500
+metric_depth_aux_max_depth: 80.0
+metric_depth_aux_weight: 0.0
+mixed_precision: bf16
+mode_scale: 1.29
+normal_prompt_contract: campaign_short
+normal_target_encoding: vision_banana_x_left_red
+normal_target_resize_policy: vector_field_renormalize_then_encode
+num_train_epochs: 10
+num_validation_images: 4
+offload: false
+optimizer: AdamW
+output_dir: /gpfs/haicore/projects/haicore470/hpc26/checkpoints/open_visionbanana/ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000
+packed_prompt_shard_bytes: 1073741824
+palette_manifold_loss_weight: 0.0
+permanent_checkpoint_steps: 6250,12500,18750,25000,31250,37500
+precompute_cache_dir: /gpfs/haicore/projects/haicore470/hpc26/precompute_cache/ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000
+pretrained_model_name_or_path: black-forest-labs/FLUX.2-klein-base-9B
+prodigy_beta3: null
+prodigy_decouple: true
+prodigy_safeguard_warmup: true
+prodigy_use_bias_correction: true
+prompt_embedding_batch_size: 128
+prompt_encoding_mode: packed_precompute
+prompt_precompute_max_unique: 2048
+push_to_hub: false
+random_flip: false
+repeats: 1
+report_to: tensorboard
+resample_dataset_each_epoch: false
+resolution: 512
+resume_from_checkpoint: latest
+revision: null
+sample_batch_size: 4
+samples_per_epoch: 250000
+sampling_strategy: without_replacement
+scale_lr: false
+seed: 42
+skip_final_inference: true
+task_homogeneous_batches: true
+train_all_transformer_blocks: true
+train_batch_size: 8
+train_full_transformer: false
+train_input_layers: false
+train_last_single_transformer_blocks: 1
+train_last_transformer_blocks: 0
+train_method: partial_full_rank
+train_non_block_layers: true
+train_output_layers: true
+train_single_transformer_block_indices: null
+train_transformer_block_indices: null
+unit_normal_loss_weight: 0.0
+upcast_before_saving: false
+use_8bit_adam: false
+validation_epochs: 50
+validation_height: null
+validation_image: null
+validation_num_inference_steps: 28
+validation_prompt: null
+validation_prompts_file: null
+validation_width: null
+variant: null
+weighting_scheme: none

CKPTS/logs/flux2-klein-partial-image2img/events.out.tfevents.1786717770.gb005.2141029.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c922af6bd70877ab539446237eb0e765594b977e614ad5045a7610e3ec4754e
+size 22862

CKPTS/logs/flux2-klein-partial-image2img/events.out.tfevents.1786718917.gb006.594302.0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ff6dd7a628c330658f18daf1d65a43a31820c472bbfb9cc05cc9604ca6b6130
+size 7471548

CKPTS/runtime_provenance.json

@@ -0,0 +1,16 @@
+{
+  "base_model": "black-forest-labs/FLUX.2-klein-base-9B",
+  "code_sha256": {
+    "common/geometry_auxiliary_losses.py": "29a96313b12e39317bd423bf06eb41ac4c8a68f9c7d4e5cea549898449844972",
+    "common/indexed_tar.py": "cae03af8d1c9b5a53ff2c09ca27eb0711079ab1d7a2599e0834db827efd73a82",
+    "scripts/training/run_open_visionbanana_pool_flux2_klein.sh": "5690d706d74bdf729ec744b0a0e5edfb5d4150a7b5f423bebbc85b923e2fc0bb",
+    "train_dreambooth_partial_flux2_klein_img2img.py": "794c84fa5e20b84dab0844190226916cb6ebb3f80d0383ec91e6498230c0844f"
+  },
+  "created_utc": "2026-08-14T14:04:37.020670+00:00",
+  "dataset_specs_sha256": "800e41b599031f71c27b3cb8b330465d02f5ebf18155bf75b769a584cc685e6d",
+  "diffusers_version": "0.39.0",
+  "git_commit": "84090870c046eff46822e8df5112e01d0dda0414",
+  "schema_version": 1,
+  "torch_version": "2.11.0+cu128",
+  "transformers_version": "4.57.6"
+}

CKPTS/semseg_diversify_config.json

@@ -0,0 +1,20 @@
+{
+  "background_color_prob": 0.0,
+  "canonical_anchor": false,
+  "color_mapping": "p100_randomnegdraw",
+  "color_min_delta_e": 25.0,
+  "full_palette_size": 200,
+  "full_set_ratio": 0.15,
+  "max_negative_classes": 30,
+  "min_class_area": 0.0,
+  "named_palette_size": 37,
+  "negative_set_ratio": 0.15,
+  "notation_styles": [
+    "name",
+    "rgb",
+    "hex"
+  ],
+  "seed": 42,
+  "subset_ratio": 0.7,
+  "variants_per_image": 1
+}

CKPTS/training_args.json

@@ -0,0 +1,131 @@
+{
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "adam_epsilon": 1e-08,
+  "adam_weight_decay": 0.0001,
+  "adam_weight_decay_text_encoder": 0.001,
+  "allow_large_prompt_precompute_cache": false,
+  "allow_legacy_hypersim_range_depth": false,
+  "allow_tf32": true,
+  "allow_unverified_normal_sources": false,
+  "angular_normal_aux_weight": 0.0,
+  "aspect_ratio_buckets": "512,512;576,448;448,576;608,416;416,608",
+  "bnb_quantization_config_path": null,
+  "cache_dir": null,
+  "cache_latents": false,
+  "caption_column": "text",
+  "center_crop": false,
+  "checkpointing_steps": 625,
+  "checkpoints_total_limit": 2,
+  "class_data_dir": null,
+  "cond_image_column": "conditioning_image",
+  "dataloader_num_workers": 8,
+  "dataset_config_name": null,
+  "dataset_index_num_workers": 16,
+  "dataset_name": null,
+  "dataset_pool_seed": 42,
+  "dataset_specs": "/home/hpc26/open_visionbanana/open_visionbanana_data/ovb_prepared_seg_geo_pd12m_v2_semseg_p100_iseg_fixeddraw_32_ordered_unlisted/open_visionbanana_mix_v1_tmpdir.json",
+  "depth_prompt_contract": "campaign_detailed_inverse",
+  "depth_target_encoding": "relative_inverse_grayscale",
+  "depth_target_resize_policy": "metric_field_then_encode",
+  "do_fp8_training": false,
+  "enable_npu_flash_attention": false,
+  "final_validation_prompt": null,
+  "fsdp_text_encoder": false,
+  "fsdp_transformer": true,
+  "fsdp_transformer_cls_to_wrap": "Flux2TransformerBlock,Flux2SingleTransformerBlock",
+  "fsdp_transformer_dtype_policy": "fp32_all",
+  "gradient_accumulation_steps": 1,
+  "gradient_checkpointing": false,
+  "guidance_scale": 3.5,
+  "hub_model_id": null,
+  "hub_token": null,
+  "image_column": "image",
+  "instance_data_dir": null,
+  "instance_prompt": "[TASK: IMAGE_EDITING] Edit the image according to the instruction.",
+  "learning_rate": 2e-06,
+  "local_rank": 0,
+  "logging_dir": "logs",
+  "logit_mean": 0.0,
+  "logit_std": 1.0,
+  "lora_alpha": 64,
+  "lora_dropout": 0.0,
+  "lora_layers": null,
+  "lora_rank": 64,
+  "loss_mask_erosion_radius": 0,
+  "loss_mask_min_sample_valid_fraction": 0.0,
+  "loss_mask_min_valid_fraction": 0.5,
+  "loss_mask_mode": "hard",
+  "loss_mask_policy": "disabled",
+  "loss_mask_reduction": "per_sample",
+  "lr_num_cycles": 1,
+  "lr_power": 1.0,
+  "lr_scheduler": "constant_with_warmup",
+  "lr_warmup_steps": 1000,
+  "max_grad_norm": 1.0,
+  "max_refseg_expressions_per_image": 2000000,
+  "max_sequence_length": 512,
+  "max_train_steps": 37500,
+  "metric_depth_aux_max_depth": 80.0,
+  "metric_depth_aux_weight": 0.0,
+  "mixed_precision": "bf16",
+  "mode_scale": 1.29,
+  "normal_prompt_contract": "campaign_short",
+  "normal_target_encoding": "vision_banana_x_left_red",
+  "normal_target_resize_policy": "vector_field_renormalize_then_encode",
+  "num_train_epochs": 1,
+  "num_validation_images": 4,
+  "offload": false,
+  "optimizer": "AdamW",
+  "output_dir": "/gpfs/haicore/projects/haicore470/hpc26/checkpoints/open_visionbanana/ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000",
+  "packed_prompt_shard_bytes": 1073741824,
+  "palette_manifold_loss_weight": 0.0,
+  "permanent_checkpoint_steps": "6250,12500,18750,25000,31250,37500",
+  "precompute_cache_dir": "/gpfs/haicore/projects/haicore470/hpc26/precompute_cache/ovb_f2k9b_v2-p100-iseg-fixeddraw32-ordered-unlisted_r512_allblocks-nonblock_lr2e-6_b8-ga1-gpu8_37500_wu1000",
+  "pretrained_model_name_or_path": "black-forest-labs/FLUX.2-klein-base-9B",
+  "prodigy_beta3": null,
+  "prodigy_decouple": true,
+  "prodigy_safeguard_warmup": true,
+  "prodigy_use_bias_correction": true,
+  "prompt_embedding_batch_size": 128,
+  "prompt_encoding_mode": "packed_precompute",
+  "prompt_precompute_max_unique": 2048,
+  "push_to_hub": false,
+  "random_flip": false,
+  "repeats": 1,
+  "report_to": "tensorboard",
+  "resample_dataset_each_epoch": false,
+  "resolution": 512,
+  "resume_from_checkpoint": "latest",
+  "revision": null,
+  "sample_batch_size": 4,
+  "samples_per_epoch": 250000,
+  "sampling_strategy": "without_replacement",
+  "scale_lr": false,
+  "seed": 42,
+  "skip_final_inference": true,
+  "task_homogeneous_batches": true,
+  "train_all_transformer_blocks": true,
+  "train_batch_size": 8,
+  "train_full_transformer": false,
+  "train_input_layers": false,
+  "train_last_single_transformer_blocks": 1,
+  "train_last_transformer_blocks": 0,
+  "train_method": "partial_full_rank",
+  "train_non_block_layers": true,
+  "train_output_layers": true,
+  "train_single_transformer_block_indices": null,
+  "train_transformer_block_indices": null,
+  "unit_normal_loss_weight": 0.0,
+  "upcast_before_saving": false,
+  "use_8bit_adam": false,
+  "validation_epochs": 50,
+  "validation_height": null,
+  "validation_image": null,
+  "validation_num_inference_steps": 28,
+  "validation_prompt": null,
+  "validation_prompts_file": null,
+  "validation_width": null,
+  "variant": null,
+  "weighting_scheme": "none"
+}

CKPTS/training_metrics/metrics.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:328de01dcd6de629e1495cdaec88534233bb2010b84694b336f0b4b11b948e5c
+size 48982375

CKPTS/transformer/config.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "Flux2Transformer2DModel",
+  "_diffusers_version": "0.39.0",
+  "_name_or_path": "black-forest-labs/FLUX.2-klein-base-9B",
+  "attention_head_dim": 128,
+  "axes_dims_rope": [
+    32,
+    32,
+    32,
+    32
+  ],
+  "eps": 1e-06,
+  "guidance_embeds": false,
+  "in_channels": 128,
+  "joint_attention_dim": 12288,
+  "mlp_ratio": 3.0,
+  "num_attention_heads": 32,
+  "num_layers": 8,
+  "num_single_layers": 24,
+  "out_channels": null,
+  "patch_size": 1,
+  "rope_theta": 2000,
+  "timestep_guidance_channels": 256
+}

CKPTS/transformer/diffusion_pytorch_model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbb926f875e52a38fc40660a3ec0b37f4fd2d75e85624279bbcca926317dc116
+size 9804206048

CKPTS/transformer/diffusion_pytorch_model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:372b1dc97c6f27fae073a0cfbc3913236c8bba56579aab1f90ded850443d4d7d
+size 9797913704

CKPTS/transformer/diffusion_pytorch_model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd117aa0a5ba2dc606ece70c5c7ada7932f4a2bd65e5d2c15312bfe6d82fc6c7
+size 9596584288

CKPTS/transformer/diffusion_pytorch_model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cf45bdc042c4f9712deaed5f00083dc86a03e6caf53981d1bf8ec174f43bb69
+size 7115649136

CKPTS/transformer/diffusion_pytorch_model.safetensors.index.json

@@ -0,0 +1,240 @@
+{
+  "metadata": {
+    "total_size": 36314324992
+  },
+  "weight_map": {
+    "context_embedder.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "double_stream_modulation_img.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "double_stream_modulation_txt.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "norm_out.linear.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "proj_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_stream_modulation.linear.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.1.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.10.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.11.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.12.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.13.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.14.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.15.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.16.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.17.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.18.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.19.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.2.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.20.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.21.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.22.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.norm_k.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.norm_q.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.to_out.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.23.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00004-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.3.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.to_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.4.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.5.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.6.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.7.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.8.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "single_transformer_blocks.9.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00004.safetensors",
+    "time_guidance_embed.timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "time_guidance_embed.timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.0.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.1.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.2.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.3.ff_context.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_k_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_q_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.add_v_proj.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_added_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_added_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_add_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff_context.linear_in.weight": "diffusion_pytorch_model-00001-of-00004.safetensors",
+    "transformer_blocks.4.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.5.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.6.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "transformer_blocks.7.ff_context.linear_out.weight": "diffusion_pytorch_model-00002-of-00004.safetensors",
+    "x_embedder.weight": "diffusion_pytorch_model-00001-of-00004.safetensors"
+  }
+}