Add files using upload-large-folder tool

.gitattributes

@@ -96,3 +96,4 @@ LTA_openwebtext_dualt/mini_owt_logdirichlet/cache/owt_llmclean_qwen36_35b_articl
 LTA_openwebtext_dualt/mini_owt_logdirichlet/cache/owt_llmclean_qwen36_35b_articlefull_300k_snapshot_rev8_docs600000_shards32/part-018.jsonl filter=lfs diff=lfs merge=lfs -text
 LTA_openwebtext_dualt/mini_owt_logdirichlet/cache/owt_llmclean_qwen36_35b_articlefull_300k_snapshot_rev8_docs600000_shards32/part-029.jsonl filter=lfs diff=lfs merge=lfs -text
 LTA_openwebtext_dualt/mini_owt_logdirichlet/cache/owt_llmclean_qwen36_35b_articlefull_300k_snapshot_rev8_docs600000_shards32/part-019.jsonl filter=lfs diff=lfs merge=lfs -text
+LTA_openwebtext_dualt/mini_owt_logdirichlet/cache/owt_llmclean_qwen36_35b_articlefull_full_rev8/audit.jsonl filter=lfs diff=lfs merge=lfs -text

LTA_openwebtext_dualt/logs/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_gumbel_sde_watch/infer_lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525_step_0003000.log

@@ -0,0 +1,136 @@
+[watch-gumbel] 2026-05-25_21:11:10 infer runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0003000.pt -> docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0003000
+[load] runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0003000.pt
+[ckpt] step=3000
+[sde] generated 2/128
+[sde] generated 4/128
+[sde] generated 6/128
+[sde] generated 8/128
+[sde] generated 10/128
+[sde] generated 12/128
+[sde] generated 14/128
+[sde] generated 16/128
+[sde] generated 18/128
+[sde] generated 20/128
+[sde] generated 22/128
+[sde] generated 24/128
+[sde] generated 26/128
+[sde] generated 28/128
+[sde] generated 30/128
+[sde] generated 32/128
+[sde] generated 34/128
+[sde] generated 36/128
+[sde] generated 38/128
+[sde] generated 40/128
+[sde] generated 42/128
+[sde] generated 44/128
+[sde] generated 46/128
+[sde] generated 48/128
+[sde] generated 50/128
+[sde] generated 52/128
+[sde] generated 54/128
+[sde] generated 56/128
+[sde] generated 58/128
+[sde] generated 60/128
+[sde] generated 62/128
+[sde] generated 64/128
+[sde] generated 66/128
+[sde] generated 68/128
+[sde] generated 70/128
+[sde] generated 72/128
+[sde] generated 74/128
+[sde] generated 76/128
+[sde] generated 78/128
+[sde] generated 80/128
+[sde] generated 82/128
+[sde] generated 84/128
+[sde] generated 86/128
+[sde] generated 88/128
+[sde] generated 90/128
+[sde] generated 92/128
+[sde] generated 94/128
+[sde] generated 96/128
+[sde] generated 98/128
+[sde] generated 100/128
+[sde] generated 102/128
+[sde] generated 104/128
+[sde] generated 106/128
+[sde] generated 108/128
+[sde] generated 110/128
+[sde] generated 112/128
+[sde] generated 114/128
+[sde] generated 116/128
+[sde] generated 118/128
+[sde] generated 120/128
+[sde] generated 122/128
+[sde] generated 124/128
+[sde] generated 126/128
+[sde] generated 128/128
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0003000.pt",
+  "step": 3000,
+  "decode": {
+    "decode_rule": "dirichlet_resample_sde",
+    "steps": 128,
+    "model_t_mode": "support_t",
+    "mean_mode": "endpoint_only",
+    "anchor_gamma": 1.0,
+    "endpoint_floor": 0.0,
+    "concentration_min": 30522.0,
+    "concentration_max": 61044.0,
+    "endpoint_temp": 1.45,
+    "endpoint_temp_start": null,
+    "endpoint_temp_end": null,
+    "endpoint_projection": "gumbel_softmax",
+    "endpoint_top_k": 0,
+    "endpoint_top_p": 0.95,
+    "gumbel_tau_start": 1.0,
+    "gumbel_tau_end": 0.2,
+    "gumbel_noise_scale_start": 1.0,
+    "gumbel_noise_scale_end": 1.0,
+    "ban_special_tokens": false,
+    "banned_endpoint_ids": [],
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "dirichlet",
+    "noise_sigma": -1.0,
+    "noise_dirichlet_concentration": 30522.0,
+    "sde_resample": "dirichlet",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 128,
+    "seed": 20260524
+  },
+  "raw_genppl": {
+    "ppl": 3.8147677280337264,
+    "nll_per_token": 1.3388797790123395,
+    "tokens": 125410,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 3.7623815682257042,
+    "nll_per_token": 1.3250521526960528,
+    "tokens": 124984,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 1.1669265594308051,
+    "unique_tokens": 557,
+    "token_count": 131072,
+    "distinct_1": 0.00424957275390625,
+    "distinct_2": 0.02731702101661779,
+    "top_token_mass": 0.5301742553710938
+  }
+}
+[done] docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0003000/sde_steps128_samples128_scored.jsonl
+[watch-gumbel] 2026-05-25_21:22:30 done step_0003000

LTA_openwebtext_dualt/logs/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_gumbel_sde_watch/infer_lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525_step_0006000.log

@@ -0,0 +1,136 @@
+[watch-gumbel] 2026-05-26_01:14:44 infer runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0006000.pt -> docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0006000
+[load] runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0006000.pt
+[ckpt] step=6000
+[sde] generated 2/128
+[sde] generated 4/128
+[sde] generated 6/128
+[sde] generated 8/128
+[sde] generated 10/128
+[sde] generated 12/128
+[sde] generated 14/128
+[sde] generated 16/128
+[sde] generated 18/128
+[sde] generated 20/128
+[sde] generated 22/128
+[sde] generated 24/128
+[sde] generated 26/128
+[sde] generated 28/128
+[sde] generated 30/128
+[sde] generated 32/128
+[sde] generated 34/128
+[sde] generated 36/128
+[sde] generated 38/128
+[sde] generated 40/128
+[sde] generated 42/128
+[sde] generated 44/128
+[sde] generated 46/128
+[sde] generated 48/128
+[sde] generated 50/128
+[sde] generated 52/128
+[sde] generated 54/128
+[sde] generated 56/128
+[sde] generated 58/128
+[sde] generated 60/128
+[sde] generated 62/128
+[sde] generated 64/128
+[sde] generated 66/128
+[sde] generated 68/128
+[sde] generated 70/128
+[sde] generated 72/128
+[sde] generated 74/128
+[sde] generated 76/128
+[sde] generated 78/128
+[sde] generated 80/128
+[sde] generated 82/128
+[sde] generated 84/128
+[sde] generated 86/128
+[sde] generated 88/128
+[sde] generated 90/128
+[sde] generated 92/128
+[sde] generated 94/128
+[sde] generated 96/128
+[sde] generated 98/128
+[sde] generated 100/128
+[sde] generated 102/128
+[sde] generated 104/128
+[sde] generated 106/128
+[sde] generated 108/128
+[sde] generated 110/128
+[sde] generated 112/128
+[sde] generated 114/128
+[sde] generated 116/128
+[sde] generated 118/128
+[sde] generated 120/128
+[sde] generated 122/128
+[sde] generated 124/128
+[sde] generated 126/128
+[sde] generated 128/128
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0006000.pt",
+  "step": 6000,
+  "decode": {
+    "decode_rule": "dirichlet_resample_sde",
+    "steps": 128,
+    "model_t_mode": "support_t",
+    "mean_mode": "endpoint_only",
+    "anchor_gamma": 1.0,
+    "endpoint_floor": 0.0,
+    "concentration_min": 30522.0,
+    "concentration_max": 61044.0,
+    "endpoint_temp": 1.45,
+    "endpoint_temp_start": null,
+    "endpoint_temp_end": null,
+    "endpoint_projection": "gumbel_softmax",
+    "endpoint_top_k": 0,
+    "endpoint_top_p": 0.95,
+    "gumbel_tau_start": 1.0,
+    "gumbel_tau_end": 0.2,
+    "gumbel_noise_scale_start": 1.0,
+    "gumbel_noise_scale_end": 1.0,
+    "ban_special_tokens": false,
+    "banned_endpoint_ids": [],
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "dirichlet",
+    "noise_sigma": -1.0,
+    "noise_dirichlet_concentration": 30522.0,
+    "sde_resample": "dirichlet",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 128,
+    "seed": 20260524
+  },
+  "raw_genppl": {
+    "ppl": 2.0555040043182493,
+    "nll_per_token": 0.7205210754532657,
+    "tokens": 128145,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 2.0238469957582215,
+    "nll_per_token": 0.7050001535866586,
+    "tokens": 127950,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 1.2342076949480105,
+    "unique_tokens": 923,
+    "token_count": 131072,
+    "distinct_1": 0.00704193115234375,
+    "distinct_2": 0.037894061583577714,
+    "top_token_mass": 0.26293182373046875
+  }
+}
+[done] docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0006000/sde_steps128_samples128_scored.jsonl
+[watch-gumbel] 2026-05-26_01:33:28 done step_0006000

LTA_openwebtext_dualt/logs/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_gumbel_sde_watch/infer_lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525_step_0009000.log

@@ -0,0 +1,136 @@
+[watch-gumbel] 2026-05-26_06:55:00 infer runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0009000.pt -> docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0009000
+[load] runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0009000.pt
+[ckpt] step=9000
+[sde] generated 2/128
+[sde] generated 4/128
+[sde] generated 6/128
+[sde] generated 8/128
+[sde] generated 10/128
+[sde] generated 12/128
+[sde] generated 14/128
+[sde] generated 16/128
+[sde] generated 18/128
+[sde] generated 20/128
+[sde] generated 22/128
+[sde] generated 24/128
+[sde] generated 26/128
+[sde] generated 28/128
+[sde] generated 30/128
+[sde] generated 32/128
+[sde] generated 34/128
+[sde] generated 36/128
+[sde] generated 38/128
+[sde] generated 40/128
+[sde] generated 42/128
+[sde] generated 44/128
+[sde] generated 46/128
+[sde] generated 48/128
+[sde] generated 50/128
+[sde] generated 52/128
+[sde] generated 54/128
+[sde] generated 56/128
+[sde] generated 58/128
+[sde] generated 60/128
+[sde] generated 62/128
+[sde] generated 64/128
+[sde] generated 66/128
+[sde] generated 68/128
+[sde] generated 70/128
+[sde] generated 72/128
+[sde] generated 74/128
+[sde] generated 76/128
+[sde] generated 78/128
+[sde] generated 80/128
+[sde] generated 82/128
+[sde] generated 84/128
+[sde] generated 86/128
+[sde] generated 88/128
+[sde] generated 90/128
+[sde] generated 92/128
+[sde] generated 94/128
+[sde] generated 96/128
+[sde] generated 98/128
+[sde] generated 100/128
+[sde] generated 102/128
+[sde] generated 104/128
+[sde] generated 106/128
+[sde] generated 108/128
+[sde] generated 110/128
+[sde] generated 112/128
+[sde] generated 114/128
+[sde] generated 116/128
+[sde] generated 118/128
+[sde] generated 120/128
+[sde] generated 122/128
+[sde] generated 124/128
+[sde] generated 126/128
+[sde] generated 128/128
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0009000.pt",
+  "step": 9000,
+  "decode": {
+    "decode_rule": "dirichlet_resample_sde",
+    "steps": 128,
+    "model_t_mode": "support_t",
+    "mean_mode": "endpoint_only",
+    "anchor_gamma": 1.0,
+    "endpoint_floor": 0.0,
+    "concentration_min": 30522.0,
+    "concentration_max": 61044.0,
+    "endpoint_temp": 1.45,
+    "endpoint_temp_start": null,
+    "endpoint_temp_end": null,
+    "endpoint_projection": "gumbel_softmax",
+    "endpoint_top_k": 0,
+    "endpoint_top_p": 0.95,
+    "gumbel_tau_start": 1.0,
+    "gumbel_tau_end": 0.2,
+    "gumbel_noise_scale_start": 1.0,
+    "gumbel_noise_scale_end": 1.0,
+    "ban_special_tokens": false,
+    "banned_endpoint_ids": [],
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "dirichlet",
+    "noise_sigma": -1.0,
+    "noise_dirichlet_concentration": 30522.0,
+    "sde_resample": "dirichlet",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 128,
+    "seed": 20260524
+  },
+  "raw_genppl": {
+    "ppl": 5.088105610903924,
+    "nll_per_token": 1.6269055826824483,
+    "tokens": 129820,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 5.014868702762704,
+    "nll_per_token": 1.612407240166428,
+    "tokens": 129623,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 1.3782953161976421,
+    "unique_tokens": 1229,
+    "token_count": 131072,
+    "distinct_1": 0.00937652587890625,
+    "distinct_2": 0.0396963587487781,
+    "top_token_mass": 0.366668701171875
+  }
+}
+[done] docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0009000/sde_steps128_samples128_scored.jsonl
+[watch-gumbel] 2026-05-26_07:13:03 done step_0009000

LTA_openwebtext_dualt/logs/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_gumbel_sde_watch/infer_lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525_step_0010000.log

@@ -0,0 +1,136 @@
+[watch-gumbel] 2026-05-26_08:47:07 infer runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0010000.pt -> docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0010000
+[load] runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0010000.pt
+[ckpt] step=10000
+[sde] generated 2/128
+[sde] generated 4/128
+[sde] generated 6/128
+[sde] generated 8/128
+[sde] generated 10/128
+[sde] generated 12/128
+[sde] generated 14/128
+[sde] generated 16/128
+[sde] generated 18/128
+[sde] generated 20/128
+[sde] generated 22/128
+[sde] generated 24/128
+[sde] generated 26/128
+[sde] generated 28/128
+[sde] generated 30/128
+[sde] generated 32/128
+[sde] generated 34/128
+[sde] generated 36/128
+[sde] generated 38/128
+[sde] generated 40/128
+[sde] generated 42/128
+[sde] generated 44/128
+[sde] generated 46/128
+[sde] generated 48/128
+[sde] generated 50/128
+[sde] generated 52/128
+[sde] generated 54/128
+[sde] generated 56/128
+[sde] generated 58/128
+[sde] generated 60/128
+[sde] generated 62/128
+[sde] generated 64/128
+[sde] generated 66/128
+[sde] generated 68/128
+[sde] generated 70/128
+[sde] generated 72/128
+[sde] generated 74/128
+[sde] generated 76/128
+[sde] generated 78/128
+[sde] generated 80/128
+[sde] generated 82/128
+[sde] generated 84/128
+[sde] generated 86/128
+[sde] generated 88/128
+[sde] generated 90/128
+[sde] generated 92/128
+[sde] generated 94/128
+[sde] generated 96/128
+[sde] generated 98/128
+[sde] generated 100/128
+[sde] generated 102/128
+[sde] generated 104/128
+[sde] generated 106/128
+[sde] generated 108/128
+[sde] generated 110/128
+[sde] generated 112/128
+[sde] generated 114/128
+[sde] generated 116/128
+[sde] generated 118/128
+[sde] generated 120/128
+[sde] generated 122/128
+[sde] generated 124/128
+[sde] generated 126/128
+[sde] generated 128/128
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0010000.pt",
+  "step": 10000,
+  "decode": {
+    "decode_rule": "dirichlet_resample_sde",
+    "steps": 128,
+    "model_t_mode": "support_t",
+    "mean_mode": "endpoint_only",
+    "anchor_gamma": 1.0,
+    "endpoint_floor": 0.0,
+    "concentration_min": 30522.0,
+    "concentration_max": 61044.0,
+    "endpoint_temp": 1.45,
+    "endpoint_temp_start": null,
+    "endpoint_temp_end": null,
+    "endpoint_projection": "gumbel_softmax",
+    "endpoint_top_k": 0,
+    "endpoint_top_p": 0.95,
+    "gumbel_tau_start": 1.0,
+    "gumbel_tau_end": 0.2,
+    "gumbel_noise_scale_start": 1.0,
+    "gumbel_noise_scale_end": 1.0,
+    "ban_special_tokens": false,
+    "banned_endpoint_ids": [],
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "dirichlet",
+    "noise_sigma": -1.0,
+    "noise_dirichlet_concentration": 30522.0,
+    "sde_resample": "dirichlet",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 128,
+    "seed": 20260524
+  },
+  "raw_genppl": {
+    "ppl": 7.382004975613158,
+    "nll_per_token": 1.9990452786272435,
+    "tokens": 108846,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 7.193271401079397,
+    "nll_per_token": 1.9731460614509597,
+    "tokens": 108120,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 1.9640364761127493,
+    "unique_tokens": 1041,
+    "token_count": 131072,
+    "distinct_1": 0.00794219970703125,
+    "distinct_2": 0.05942234848484849,
+    "top_token_mass": 0.31719970703125
+  }
+}
+[done] docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0010000/sde_steps128_samples128_scored.jsonl
+[watch-gumbel] 2026-05-26_09:06:14 done step_0010000

LTA_openwebtext_dualt/logs/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_gumbel_sde_watch/infer_lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525_step_0012000.log

@@ -0,0 +1,136 @@
+[watch-gumbel] 2026-05-26_12:52:47 infer runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0012000.pt -> docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0012000
+[load] runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0012000.pt
+[ckpt] step=12000
+[sde] generated 2/128
+[sde] generated 4/128
+[sde] generated 6/128
+[sde] generated 8/128
+[sde] generated 10/128
+[sde] generated 12/128
+[sde] generated 14/128
+[sde] generated 16/128
+[sde] generated 18/128
+[sde] generated 20/128
+[sde] generated 22/128
+[sde] generated 24/128
+[sde] generated 26/128
+[sde] generated 28/128
+[sde] generated 30/128
+[sde] generated 32/128
+[sde] generated 34/128
+[sde] generated 36/128
+[sde] generated 38/128
+[sde] generated 40/128
+[sde] generated 42/128
+[sde] generated 44/128
+[sde] generated 46/128
+[sde] generated 48/128
+[sde] generated 50/128
+[sde] generated 52/128
+[sde] generated 54/128
+[sde] generated 56/128
+[sde] generated 58/128
+[sde] generated 60/128
+[sde] generated 62/128
+[sde] generated 64/128
+[sde] generated 66/128
+[sde] generated 68/128
+[sde] generated 70/128
+[sde] generated 72/128
+[sde] generated 74/128
+[sde] generated 76/128
+[sde] generated 78/128
+[sde] generated 80/128
+[sde] generated 82/128
+[sde] generated 84/128
+[sde] generated 86/128
+[sde] generated 88/128
+[sde] generated 90/128
+[sde] generated 92/128
+[sde] generated 94/128
+[sde] generated 96/128
+[sde] generated 98/128
+[sde] generated 100/128
+[sde] generated 102/128
+[sde] generated 104/128
+[sde] generated 106/128
+[sde] generated 108/128
+[sde] generated 110/128
+[sde] generated 112/128
+[sde] generated 114/128
+[sde] generated 116/128
+[sde] generated 118/128
+[sde] generated 120/128
+[sde] generated 122/128
+[sde] generated 124/128
+[sde] generated 126/128
+[sde] generated 128/128
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0012000.pt",
+  "step": 12000,
+  "decode": {
+    "decode_rule": "dirichlet_resample_sde",
+    "steps": 128,
+    "model_t_mode": "support_t",
+    "mean_mode": "endpoint_only",
+    "anchor_gamma": 1.0,
+    "endpoint_floor": 0.0,
+    "concentration_min": 30522.0,
+    "concentration_max": 61044.0,
+    "endpoint_temp": 1.45,
+    "endpoint_temp_start": null,
+    "endpoint_temp_end": null,
+    "endpoint_projection": "gumbel_softmax",
+    "endpoint_top_k": 0,
+    "endpoint_top_p": 0.95,
+    "gumbel_tau_start": 1.0,
+    "gumbel_tau_end": 0.2,
+    "gumbel_noise_scale_start": 1.0,
+    "gumbel_noise_scale_end": 1.0,
+    "ban_special_tokens": false,
+    "banned_endpoint_ids": [],
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "dirichlet",
+    "noise_sigma": -1.0,
+    "noise_dirichlet_concentration": 30522.0,
+    "sde_resample": "dirichlet",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 128,
+    "seed": 20260524
+  },
+  "raw_genppl": {
+    "ppl": 9.787557925258849,
+    "nll_per_token": 2.281111979588946,
+    "tokens": 121050,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 9.575427309198085,
+    "nll_per_token": 2.2592001616634176,
+    "tokens": 120277,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 2.3105578319413413,
+    "unique_tokens": 1540,
+    "token_count": 131072,
+    "distinct_1": 0.011749267578125,
+    "distinct_2": 0.08389082355816227,
+    "top_token_mass": 0.17154693603515625
+  }
+}
+[done] docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0012000/sde_steps128_samples128_scored.jsonl
+[watch-gumbel] 2026-05-26_13:16:51 done step_0012000

LTA_openwebtext_dualt/logs/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_gumbel_sde_watch/infer_lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525_step_0013000.log

@@ -0,0 +1,136 @@
+[watch-gumbel] 2026-05-26_15:32:28 infer runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0013000.pt -> docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0013000
+[load] runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0013000.pt
+[ckpt] step=13000
+[sde] generated 2/128
+[sde] generated 4/128
+[sde] generated 6/128
+[sde] generated 8/128
+[sde] generated 10/128
+[sde] generated 12/128
+[sde] generated 14/128
+[sde] generated 16/128
+[sde] generated 18/128
+[sde] generated 20/128
+[sde] generated 22/128
+[sde] generated 24/128
+[sde] generated 26/128
+[sde] generated 28/128
+[sde] generated 30/128
+[sde] generated 32/128
+[sde] generated 34/128
+[sde] generated 36/128
+[sde] generated 38/128
+[sde] generated 40/128
+[sde] generated 42/128
+[sde] generated 44/128
+[sde] generated 46/128
+[sde] generated 48/128
+[sde] generated 50/128
+[sde] generated 52/128
+[sde] generated 54/128
+[sde] generated 56/128
+[sde] generated 58/128
+[sde] generated 60/128
+[sde] generated 62/128
+[sde] generated 64/128
+[sde] generated 66/128
+[sde] generated 68/128
+[sde] generated 70/128
+[sde] generated 72/128
+[sde] generated 74/128
+[sde] generated 76/128
+[sde] generated 78/128
+[sde] generated 80/128
+[sde] generated 82/128
+[sde] generated 84/128
+[sde] generated 86/128
+[sde] generated 88/128
+[sde] generated 90/128
+[sde] generated 92/128
+[sde] generated 94/128
+[sde] generated 96/128
+[sde] generated 98/128
+[sde] generated 100/128
+[sde] generated 102/128
+[sde] generated 104/128
+[sde] generated 106/128
+[sde] generated 108/128
+[sde] generated 110/128
+[sde] generated 112/128
+[sde] generated 114/128
+[sde] generated 116/128
+[sde] generated 118/128
+[sde] generated 120/128
+[sde] generated 122/128
+[sde] generated 124/128
+[sde] generated 126/128
+[sde] generated 128/128
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0013000.pt",
+  "step": 13000,
+  "decode": {
+    "decode_rule": "dirichlet_resample_sde",
+    "steps": 128,
+    "model_t_mode": "support_t",
+    "mean_mode": "endpoint_only",
+    "anchor_gamma": 1.0,
+    "endpoint_floor": 0.0,
+    "concentration_min": 30522.0,
+    "concentration_max": 61044.0,
+    "endpoint_temp": 1.45,
+    "endpoint_temp_start": null,
+    "endpoint_temp_end": null,
+    "endpoint_projection": "gumbel_softmax",
+    "endpoint_top_k": 0,
+    "endpoint_top_p": 0.95,
+    "gumbel_tau_start": 1.0,
+    "gumbel_tau_end": 0.2,
+    "gumbel_noise_scale_start": 1.0,
+    "gumbel_noise_scale_end": 1.0,
+    "ban_special_tokens": false,
+    "banned_endpoint_ids": [],
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "dirichlet",
+    "noise_sigma": -1.0,
+    "noise_dirichlet_concentration": 30522.0,
+    "sde_resample": "dirichlet",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 128,
+    "seed": 20260524
+  },
+  "raw_genppl": {
+    "ppl": 11.019391485874875,
+    "nll_per_token": 2.399656583132524,
+    "tokens": 130289,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 10.84729878838137,
+    "nll_per_token": 2.3839160893946967,
+    "tokens": 129965,
+    "kept_samples": 128,
+    "total_samples": 128,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 2.991297718897202,
+    "unique_tokens": 1773,
+    "token_count": 131072,
+    "distinct_1": 0.01352691650390625,
+    "distinct_2": 0.11194098240469208,
+    "top_token_mass": 0.14432525634765625
+  }
+}
+[done] docs/lta_samples/metrics_20260525/owt_bert_absrope_adaln_Cv_to_2v_mask1_sameT_every1k_sde_gumbel_topp0p95_tau1p0_to_0p2_blend_c30522_61044_n128/lta_owt_bert_absrope_adaln_dirichlet_len1024_Cv_to_2v_mask1_sameT_gbs512_b4x4_1m_save1k_watch_20260525/step_0013000/sde_steps128_samples128_scored.jsonl
+[watch-gumbel] 2026-05-26_15:50:49 done step_0013000

LTA_openwebtext_dualt/mini_owt_fit/runs/mini_owt_fit_bert_10/step_012500.pt

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

LTA_openwebtext_dualt/mini_owt_fit/runs/mini_owt_fit_bert_len1024_C1_to_1024_absrope_time4_d256_l3_h4_subset8192_8gpu_20260526_144131/step_012000.pt

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

LTA_openwebtext_dualt/mini_owt_fit/runs/mini_owt_fit_bert_len1024_C1_to_1024_absrope_time4_d256_l3_h4_subset8192_8gpu_20260526_144131/step_019000.pt

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

LTA_openwebtext_dualt/mini_owt_fit/runs/mini_owt_fit_t5_bernoulliwrong_len1024_bos_eos_C1_to_1024_absrope_time4_d768_l12_h12_native_nofloor_full_gbs512_8gpu_20260526_171143/step_004000.pt

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

LTA_openwebtext_dualt/mini_owt_fit/runs/mini_owt_fit_t5_bernoulliwrong_len1024_bos_eos_C1_to_1024_absrope_time4_d768_l12_h12_native_nofloor_full_gbs512_8gpu_20260526_171143/step_011000.pt

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

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/focalnet/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_focalnet import *
+    from .modeling_focalnet import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/focalnet/configuration_focalnet.py

@@ -0,0 +1,98 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""FocalNet model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="microsoft/focalnet-tiny")
+@strict
+class FocalNetConfig(BackboneConfigMixin, PreTrainedConfig):
+    r"""
+    use_conv_embed (`bool`, *optional*, defaults to `False`):
+        Whether to use convolutional embedding. The authors noted that using convolutional embedding usually
+        improve the performance, but it's not used by default.
+    focal_levels (`list(int)`, *optional*, defaults to `[2, 2, 2, 2]`):
+        Number of focal levels in each layer of the respective stages in the encoder.
+    focal_windows (`list(int)`, *optional*, defaults to `[3, 3, 3, 3]`):
+        Focal window size in each layer of the respective stages in the encoder.
+    hidden_dropout_prob (`float`, *optional*, defaults to 0.0):
+        The dropout probability for all fully connected layers in the embeddings and encoder.
+    use_layerscale (`bool`, *optional*, defaults to `False`):
+        Whether to use layer scale in the encoder.
+    layerscale_value (`float`, *optional*, defaults to 0.0001):
+        The initial value of the layer scale.
+    use_post_layernorm (`bool`, *optional*, defaults to `False`):
+        Whether to use post layer normalization in the encoder.
+    use_post_layernorm_in_modulation (`bool`, *optional*, defaults to `False`):
+        Whether to use post layer normalization in the modulation layer.
+    normalize_modulator (`bool`, *optional*, defaults to `False`):
+        Whether to normalize the modulator.
+    encoder_stride (`int`, *optional*, defaults to 32):
+        Factor to increase the spatial resolution by in the decoder head for masked image modeling.
+
+    Example:
+
+    ```python
+    >>> from transformers import FocalNetConfig, FocalNetModel
+
+    >>> # Initializing a FocalNet microsoft/focalnet-tiny style configuration
+    >>> configuration = FocalNetConfig()
+
+    >>> # Initializing a model (with random weights) from the microsoft/focalnet-tiny style configuration
+    >>> model = FocalNetModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "focalnet"
+
+    image_size: int | list[int] | tuple[int, int] = 224
+    patch_size: int | list[int] | tuple[int, int] = 4
+    num_channels: int = 3
+    embed_dim: int = 96
+    use_conv_embed: bool = False
+    hidden_sizes: list[int] | tuple[int, ...] = (192, 384, 768, 768)
+    depths: list[int] | tuple[int, ...] = (2, 2, 6, 2)
+    focal_levels: list[int] | tuple[int, ...] = (2, 2, 2, 2)
+    focal_windows: list[int] | tuple[int, ...] = (3, 3, 3, 3)
+    hidden_act: str = "gelu"
+    mlp_ratio: float = 4.0
+    hidden_dropout_prob: float | int = 0.0
+    drop_path_rate: float | int = 0.1
+    use_layerscale: bool = False
+    layerscale_value: float = 1e-4
+    use_post_layernorm: bool = False
+    use_post_layernorm_in_modulation: bool = False
+    normalize_modulator: bool = False
+    initializer_range: float = 0.02
+    layer_norm_eps: float = 1e-5
+    encoder_stride: int = 32
+    _out_features: list[str] | None = None
+    _out_indices: list[int] | None = None
+
+    def __post_init__(self, **kwargs):
+        self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
+        self.set_output_features_output_indices(
+            out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+        )
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["FocalNetConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/focalnet/modeling_focalnet.py

@@ -0,0 +1,928 @@
+# Copyright 2023 Microsoft Research and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch FocalNet model."""
+
+import collections.abc
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BackboneOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import ModelOutput, auto_docstring, logging
+from ...utils.generic import can_return_tuple
+from .configuration_focalnet import FocalNetConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet encoder's outputs, with potential hidden states.
+    """
+)
+@dataclass
+class FocalNetEncoderOutput(ModelOutput):
+    r"""
+    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+        shape `(batch_size, hidden_size, height, width)`.
+
+        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+        include the spatial dimensions.
+    """
+
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    reshaped_hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet model's outputs that also contains a pooling of the last hidden states.
+    """
+)
+@dataclass
+class FocalNetModelOutput(ModelOutput):
+    r"""
+    pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
+        Average pooling of the last layer hidden-state.
+    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+        shape `(batch_size, hidden_size, height, width)`.
+
+        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+        include the spatial dimensions.
+    """
+
+    last_hidden_state: torch.FloatTensor | None = None
+    pooler_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    reshaped_hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet masked image model outputs.
+    """
+)
+@dataclass
+class FocalNetMaskedImageModelingOutput(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `bool_masked_pos` is provided):
+        Masked image modeling (MLM) loss.
+    reconstruction (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+        Reconstructed pixel values.
+    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+        shape `(batch_size, hidden_size, height, width)`.
+
+        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+        include the spatial dimensions.
+    """
+
+    loss: torch.FloatTensor | None = None
+    reconstruction: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    reshaped_hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet outputs for image classification.
+    """
+)
+@dataclass
+class FocalNetImageClassifierOutput(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Classification (or regression if config.num_labels==1) loss.
+    logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+        Classification (or regression if config.num_labels==1) scores (before SoftMax).
+    reshaped_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage) of
+        shape `(batch_size, hidden_size, height, width)`.
+
+        Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+        include the spatial dimensions.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    reshaped_hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+class FocalNetEmbeddings(nn.Module):
+    """
+    Construct the patch embeddings and layernorm. Optionally, also the mask token.
+    """
+
+    def __init__(self, config, use_mask_token=False):
+        super().__init__()
+
+        self.patch_embeddings = FocalNetPatchEmbeddings(
+            config=config,
+            image_size=config.image_size,
+            patch_size=config.patch_size,
+            num_channels=config.num_channels,
+            embed_dim=config.embed_dim,
+            use_conv_embed=config.use_conv_embed,
+            is_stem=True,
+        )
+        self.patch_grid = self.patch_embeddings.grid_size
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None
+
+        self.norm = nn.LayerNorm(config.embed_dim, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(
+        self, pixel_values: torch.FloatTensor | None, bool_masked_pos: torch.BoolTensor | None = None
+    ) -> tuple[torch.Tensor]:
+        embeddings, output_dimensions = self.patch_embeddings(pixel_values)
+        embeddings = self.norm(embeddings)
+        batch_size, seq_len, _ = embeddings.size()
+
+        if bool_masked_pos is not None:
+            mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
+            # replace the masked visual tokens by mask_tokens
+            mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
+            embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
+
+        embeddings = self.dropout(embeddings)
+        return embeddings, output_dimensions
+
+
+class FocalNetPatchEmbeddings(nn.Module):
+    def __init__(
+        self,
+        config,
+        image_size,
+        patch_size,
+        num_channels,
+        embed_dim,
+        add_norm=False,
+        use_conv_embed=False,
+        is_stem=False,
+    ):
+        super().__init__()
+        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+        self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1])
+
+        if use_conv_embed:
+            # if we choose to use conv embedding, then we treat the stem and non-stem differently
+            if is_stem:
+                kernel_size = 7
+                padding = 2
+                stride = 4
+            else:
+                kernel_size = 3
+                padding = 1
+                stride = 2
+            self.projection = nn.Conv2d(
+                num_channels, embed_dim, kernel_size=kernel_size, stride=stride, padding=padding
+            )
+        else:
+            self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+        if add_norm:
+            self.norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+        else:
+            self.norm = None
+
+    def maybe_pad(self, pixel_values, height, width):
+        if width % self.patch_size[1] != 0:
+            pad_values = (0, self.patch_size[1] - width % self.patch_size[1])
+            pixel_values = nn.functional.pad(pixel_values, pad_values)
+        if height % self.patch_size[0] != 0:
+            pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0])
+            pixel_values = nn.functional.pad(pixel_values, pad_values)
+        return pixel_values
+
+    def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor, tuple[int]]:
+        _, num_channels, height, width = pixel_values.shape
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+            )
+        # pad the input to be divisible by self.patch_size, if needed
+        pixel_values = self.maybe_pad(pixel_values, height, width)
+        embeddings = self.projection(pixel_values)
+        _, _, height, width = embeddings.shape
+        output_dimensions = (height, width)
+        embeddings = embeddings.flatten(2).transpose(1, 2)
+
+        if self.norm is not None:
+            embeddings = self.norm(embeddings)
+
+        return embeddings, output_dimensions
+
+
+class FocalNetModulation(nn.Module):
+    def __init__(self, config, index, dim, focal_factor=2, bias=True, projection_dropout=0.0):
+        super().__init__()
+
+        self.dim = dim
+        self.focal_window = config.focal_windows[index]
+        self.focal_level = config.focal_levels[index]
+        self.focal_factor = focal_factor
+        self.use_post_layernorm_in_modulation = config.use_post_layernorm_in_modulation
+        self.normalize_modulator = config.normalize_modulator
+
+        self.projection_in = nn.Linear(dim, 2 * dim + (self.focal_level + 1), bias=bias)
+        self.projection_context = nn.Conv2d(dim, dim, kernel_size=1, stride=1, bias=bias)
+
+        self.activation = nn.GELU()
+        self.projection_out = nn.Linear(dim, dim)
+        self.projection_dropout = nn.Dropout(projection_dropout)
+        self.focal_layers = nn.ModuleList()
+
+        self.kernel_sizes = []
+        for k in range(self.focal_level):
+            kernel_size = self.focal_factor * k + self.focal_window
+            self.focal_layers.append(
+                nn.Sequential(
+                    nn.Conv2d(
+                        dim, dim, kernel_size=kernel_size, stride=1, groups=dim, padding=kernel_size // 2, bias=False
+                    ),
+                    nn.GELU(),
+                )
+            )
+            self.kernel_sizes.append(kernel_size)
+        if self.use_post_layernorm_in_modulation:
+            self.layernorm = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_state):
+        """
+        Args:
+            hidden_state:
+                Input features with shape of (batch_size, height, width, num_channels)
+        """
+        num_channels = hidden_state.shape[-1]
+
+        # pre linear projection
+        x = self.projection_in(hidden_state).permute(0, 3, 1, 2).contiguous()
+        q, ctx, gates = torch.split(x, (num_channels, num_channels, self.focal_level + 1), 1)
+
+        # context aggregation
+        ctx_all = 0
+        for level in range(self.focal_level):
+            ctx = self.focal_layers[level](ctx)
+            ctx_all = ctx_all + ctx * gates[:, level : level + 1]
+        ctx_global = self.activation(ctx.mean(2, keepdim=True).mean(3, keepdim=True))
+        ctx_all = ctx_all + ctx_global * gates[:, self.focal_level :]
+
+        # normalize context
+        if self.normalize_modulator:
+            ctx_all = ctx_all / (self.focal_level + 1)
+
+        # focal modulation
+        modulator = self.projection_context(ctx_all)
+        x_out = q * modulator
+        x_out = x_out.permute(0, 2, 3, 1).contiguous()
+        if self.use_post_layernorm_in_modulation:
+            x_out = self.layernorm(x_out)
+
+        # post linear projection
+        x_out = self.projection_out(x_out)
+        x_out = self.projection_dropout(x_out)
+        return x_out
+
+
+class FocalNetMlp(nn.Module):
+    def __init__(self, config, in_features, hidden_features=None, out_features=None, drop=0.0):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.activation = ACT2FN[config.hidden_act]
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, hidden_state):
+        hidden_state = self.fc1(hidden_state)
+        hidden_state = self.activation(hidden_state)
+        hidden_state = self.drop(hidden_state)
+        hidden_state = self.fc2(hidden_state)
+        hidden_state = self.drop(hidden_state)
+        return hidden_state
+
+
+# Copied from transformers.models.swin.modular_swin.SwinDropPath with SwinDropPath->FocalNetDropPath
+class FocalNetDropPath(nn.Module):
+    """Stochastic depth (DropPath) per sample, for residual blocks.
+
+    Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth
+    <https://arxiv.org/abs/1603.09382>`_.
+    """
+
+    def __init__(self, drop_prob: float = 0.0) -> None:
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return hidden_states
+        keep_prob = 1 - self.drop_prob
+        shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
+        random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
+        random_tensor = torch.floor(random_tensor + keep_prob)
+        return hidden_states.div(keep_prob) * random_tensor
+
+    def extra_repr(self) -> str:
+        return f"p={self.drop_prob}"
+
+
+class FocalNetLayer(nn.Module):
+    r"""Focal Modulation Network layer (block).
+
+    Args:
+        config (`FocalNetConfig`):
+            Model config.
+        index (`int`):
+            Layer index.
+        dim (`int`):
+            Number of input channels.
+        input_resolution (`tuple[int]`):
+            Input resolution.
+        drop_path (`float`, *optional*, defaults to 0.0):
+            Stochastic depth rate.
+    """
+
+    def __init__(self, config, index, dim, input_resolution, drop_path=0.0):
+        super().__init__()
+
+        self.config = config
+
+        # layer-specific attributes
+        self.dim = dim
+        self.input_resolution = input_resolution
+
+        # general attributes
+        self.drop = config.hidden_dropout_prob
+        self.use_post_layernorm = config.use_post_layernorm
+
+        self.norm1 = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+        self.modulation = FocalNetModulation(
+            config=config,
+            index=index,
+            dim=dim,
+            projection_dropout=self.drop,
+        )
+
+        self.drop_path = FocalNetDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+        mlp_hidden_dim = int(dim * config.mlp_ratio)
+        self.mlp = FocalNetMlp(config=config, in_features=dim, hidden_features=mlp_hidden_dim, drop=self.drop)
+
+        self.gamma_1 = 1.0
+        self.gamma_2 = 1.0
+        if config.use_layerscale:
+            self.gamma_1 = nn.Parameter(config.layerscale_value * torch.ones(dim), requires_grad=True)
+            self.gamma_2 = nn.Parameter(config.layerscale_value * torch.ones(dim), requires_grad=True)
+
+    def forward(self, hidden_state, input_dimensions):
+        height, width = input_dimensions
+        batch_size, _, num_channels = hidden_state.shape
+        shortcut = hidden_state
+
+        # Focal Modulation
+        hidden_state = hidden_state if self.use_post_layernorm else self.norm1(hidden_state)
+        hidden_state = hidden_state.view(batch_size, height, width, num_channels)
+        hidden_state = self.modulation(hidden_state).view(batch_size, height * width, num_channels)
+        hidden_state = hidden_state if not self.use_post_layernorm else self.norm1(hidden_state)
+
+        # FFN
+        hidden_state = shortcut + self.drop_path(self.gamma_1 * hidden_state)
+        hidden_state = hidden_state + self.drop_path(
+            self.gamma_2
+            * (self.norm2(self.mlp(hidden_state)) if self.use_post_layernorm else self.mlp(self.norm2(hidden_state)))
+        )
+
+        return hidden_state
+
+
+class FocalNetStage(GradientCheckpointingLayer):
+    def __init__(self, config, index, input_resolution):
+        super().__init__()
+
+        self.config = config
+        self.num_stages = len(config.depths)
+
+        embed_dim = [config.embed_dim * (2**i) for i in range(self.num_stages)]
+        dim = embed_dim[index]
+        out_dim = embed_dim[index + 1] if (index < self.num_stages - 1) else None
+        downsample = FocalNetPatchEmbeddings if (index < self.num_stages - 1) else None
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")]
+        drop_path = dpr[sum(config.depths[:index]) : sum(config.depths[: index + 1])]
+
+        self.layers = nn.ModuleList(
+            [
+                FocalNetLayer(
+                    config=config,
+                    index=index,
+                    dim=dim,
+                    input_resolution=input_resolution,
+                    drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                )
+                for i in range(config.depths[index])
+            ]
+        )
+
+        if downsample is not None:
+            self.downsample = downsample(
+                config=config,
+                image_size=input_resolution,
+                patch_size=2,
+                num_channels=dim,
+                embed_dim=out_dim,
+                add_norm=True,
+                use_conv_embed=config.use_conv_embed,
+                is_stem=False,
+            )
+        else:
+            self.downsample = None
+
+        self.pointing = False
+
+    def forward(self, hidden_states: torch.Tensor, input_dimensions: tuple[int, int]) -> tuple[torch.Tensor]:
+        height, width = input_dimensions
+        for layer_module in self.layers:
+            hidden_states = layer_module(hidden_states, input_dimensions)
+
+        hidden_states_before_downsampling = hidden_states
+        if self.downsample is not None:
+            height, width = input_dimensions
+            hidden_states = hidden_states.transpose(1, 2).reshape(
+                hidden_states_before_downsampling.shape[0], -1, height, width
+            )
+            hidden_states, output_dimensions = self.downsample(hidden_states)
+
+        else:
+            output_dimensions = (height, width, height, width)
+
+        stage_outputs = (hidden_states, hidden_states_before_downsampling, output_dimensions)
+
+        return stage_outputs
+
+
+class FocalNetEncoder(nn.Module):
+    def __init__(self, config, grid_size):
+        super().__init__()
+        self.num_stages = len(config.depths)
+        self.config = config
+
+        self.stages = nn.ModuleList(
+            [
+                FocalNetStage(
+                    config=config,
+                    index=i_layer,
+                    input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)),
+                )
+                for i_layer in range(self.num_stages)
+            ]
+        )
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        input_dimensions: tuple[int, int],
+        output_hidden_states: bool | None = False,
+        output_hidden_states_before_downsampling: bool | None = False,
+        return_dict: bool | None = True,
+    ) -> tuple | FocalNetEncoderOutput:
+        all_hidden_states = () if output_hidden_states else None
+        all_reshaped_hidden_states = () if output_hidden_states else None
+
+        if output_hidden_states:
+            batch_size, _, hidden_size = hidden_states.shape
+            # rearrange b (h w) c -> b c h w
+            reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
+            reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
+            all_hidden_states += (hidden_states,)
+            all_reshaped_hidden_states += (reshaped_hidden_state,)
+
+        for i, stage_module in enumerate(self.stages):
+            stage_outputs = stage_module(hidden_states, input_dimensions)
+
+            hidden_states = stage_outputs[0]
+            hidden_states_before_downsampling = stage_outputs[1]
+            output_dimensions = stage_outputs[2]
+
+            input_dimensions = (output_dimensions[-2], output_dimensions[-1])
+
+            if output_hidden_states and output_hidden_states_before_downsampling:
+                batch_size, _, hidden_size = hidden_states_before_downsampling.shape
+                # rearrange b (h w) c -> b c h w
+                # here we use the original (not downsampled) height and width
+                reshaped_hidden_state = hidden_states_before_downsampling.view(
+                    batch_size, *(output_dimensions[0], output_dimensions[1]), hidden_size
+                )
+                reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
+                all_hidden_states += (hidden_states_before_downsampling,)
+                all_reshaped_hidden_states += (reshaped_hidden_state,)
+            elif output_hidden_states and not output_hidden_states_before_downsampling:
+                batch_size, _, hidden_size = hidden_states.shape
+                # rearrange b (h w) c -> b c h w
+                reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
+                reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
+                all_hidden_states += (hidden_states,)
+                all_reshaped_hidden_states += (reshaped_hidden_state,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
+
+        return FocalNetEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            reshaped_hidden_states=all_reshaped_hidden_states,
+        )
+
+
+@auto_docstring
+class FocalNetPreTrainedModel(PreTrainedModel):
+    config: FocalNetConfig
+    base_model_prefix = "focalnet"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["FocalNetStage"]
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        super()._init_weights(module)
+        if isinstance(module, FocalNetEmbeddings):
+            if module.mask_token is not None:
+                init.zeros_(module.mask_token)
+        elif isinstance(module, FocalNetLayer):
+            if self.config.use_layerscale:
+                init.constant_(module.gamma_1, self.config.layerscale_value)
+                init.constant_(module.gamma_2, self.config.layerscale_value)
+
+
+@auto_docstring
+class FocalNetModel(FocalNetPreTrainedModel):
+    def __init__(self, config, add_pooling_layer=True, use_mask_token=False):
+        r"""
+        add_pooling_layer (bool, *optional*, defaults to `True`):
+            Whether to add a pooling layer
+        use_mask_token (`bool`, *optional*, defaults to `False`):
+            Whether to use a mask token for masked image modeling.
+        """
+        super().__init__(config)
+        self.config = config
+        self.num_stages = len(config.depths)
+        self.num_features = int(config.embed_dim * 2 ** (self.num_stages - 1))
+
+        self.embeddings = FocalNetEmbeddings(config, use_mask_token=use_mask_token)
+        self.encoder = FocalNetEncoder(config, self.embeddings.patch_grid)
+
+        self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
+        self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.patch_embeddings
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        bool_masked_pos: torch.BoolTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | FocalNetModelOutput:
+        r"""
+        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
+            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        if pixel_values is None:
+            raise ValueError("You have to specify pixel_values")
+
+        embedding_output, input_dimensions = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            input_dimensions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = encoder_outputs[0]
+        sequence_output = self.layernorm(sequence_output)
+
+        pooled_output = None
+        if self.pooler is not None:
+            pooled_output = self.pooler(sequence_output.transpose(1, 2))
+            pooled_output = torch.flatten(pooled_output, 1)
+
+        if not return_dict:
+            output = (sequence_output, pooled_output) + encoder_outputs[1:]
+
+            return output
+
+        return FocalNetModelOutput(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet Model with a decoder on top for masked image modeling.
+
+    This follows the same implementation as in [SimMIM](https://huggingface.co/papers/2111.09886).
+
+    <Tip>
+
+    Note that we provide a script to pre-train this model on custom data in our [examples
+    directory](https://github.com/huggingface/transformers/tree/main/examples/pytorch/image-pretraining).
+
+    </Tip>
+    """
+)
+class FocalNetForMaskedImageModeling(FocalNetPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.focalnet = FocalNetModel(config, add_pooling_layer=False, use_mask_token=True)
+
+        self.num_stages = len(config.depths)
+        num_features = int(config.embed_dim * 2 ** (self.num_stages - 1))
+        self.decoder = nn.Sequential(
+            nn.Conv2d(
+                in_channels=num_features, out_channels=config.encoder_stride**2 * config.num_channels, kernel_size=1
+            ),
+            nn.PixelShuffle(config.encoder_stride),
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        bool_masked_pos: torch.BoolTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | FocalNetMaskedImageModelingOutput:
+        r"""
+        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`):
+            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+
+        Examples:
+        ```python
+        >>> from transformers import AutoImageProcessor, FocalNetConfig, FocalNetForMaskedImageModeling
+        >>> import torch
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/focalnet-base-simmim-window6-192")
+        >>> config = FocalNetConfig()
+        >>> model = FocalNetForMaskedImageModeling(config)
+
+        >>> num_patches = (model.config.image_size // model.config.patch_size) ** 2
+        >>> pixel_values = image_processor(images=image, return_tensors="pt").pixel_values
+        >>> # create random boolean mask of shape (batch_size, num_patches)
+        >>> bool_masked_pos = torch.randint(low=0, high=2, size=(1, num_patches)).bool()
+
+        >>> outputs = model(pixel_values, bool_masked_pos=bool_masked_pos)
+        >>> loss, reconstructed_pixel_values = outputs.loss, outputs.logits
+        >>> list(reconstructed_pixel_values.shape)
+        [1, 3, 192, 192]
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        outputs = self.focalnet(
+            pixel_values,
+            bool_masked_pos=bool_masked_pos,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        # Reshape to (batch_size, num_channels, height, width)
+        sequence_output = sequence_output.transpose(1, 2)
+        batch_size, num_channels, sequence_length = sequence_output.shape
+        height = width = math.floor(sequence_length**0.5)
+        sequence_output = sequence_output.reshape(batch_size, num_channels, height, width)
+
+        # Reconstruct pixel values
+        reconstructed_pixel_values = self.decoder(sequence_output)
+
+        masked_im_loss = None
+        if bool_masked_pos is not None:
+            size = self.config.image_size // self.config.patch_size
+            bool_masked_pos = bool_masked_pos.reshape(-1, size, size)
+            mask = (
+                bool_masked_pos.repeat_interleave(self.config.patch_size, 1)
+                .repeat_interleave(self.config.patch_size, 2)
+                .unsqueeze(1)
+                .contiguous()
+            )
+            reconstruction_loss = nn.functional.l1_loss(pixel_values, reconstructed_pixel_values, reduction="none")
+            masked_im_loss = (reconstruction_loss * mask).sum() / (mask.sum() + 1e-5) / self.config.num_channels
+
+        if not return_dict:
+            output = (reconstructed_pixel_values,) + outputs[2:]
+            return ((masked_im_loss,) + output) if masked_im_loss is not None else output
+
+        return FocalNetMaskedImageModelingOutput(
+            loss=masked_im_loss,
+            reconstruction=reconstructed_pixel_values,
+            hidden_states=outputs.hidden_states,
+            reshaped_hidden_states=outputs.reshaped_hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet Model with an image classification head on top (a linear layer on top of the pooled output) e.g. for
+    ImageNet.
+    """
+)
+class FocalNetForImageClassification(FocalNetPreTrainedModel):
+    # Copied from transformers.models.swin.modeling_swin.SwinForImageClassification.__init__ with Swin->FocalNet, swin->focalnet
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.num_labels = config.num_labels
+        self.focalnet = FocalNetModel(config)
+
+        # Classifier head
+        self.classifier = (
+            nn.Linear(self.focalnet.num_features, config.num_labels) if config.num_labels > 0 else nn.Identity()
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | FocalNetImageClassifierOutput:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        outputs = self.focalnet(
+            pixel_values,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        pooled_output = outputs[1]
+
+        logits = self.classifier(pooled_output)
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(labels, logits, self.config)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return FocalNetImageClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            reshaped_hidden_states=outputs.reshaped_hidden_states,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    FocalNet backbone, to be used with frameworks like X-Decoder.
+    """
+)
+class FocalNetBackbone(BackboneMixin, FocalNetPreTrainedModel):
+    has_attentions = False
+
+    def __init__(self, config: FocalNetConfig):
+        super().__init__(config)
+
+        self.num_features = [config.embed_dim] + config.hidden_sizes
+        self.focalnet = FocalNetModel(config)
+
+        # initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @filter_output_hidden_states
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> BackboneOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, AutoBackbone
+        >>> import torch
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> processor = AutoImageProcessor.from_pretrained("microsoft/focalnet-tiny-lrf")
+        >>> model = AutoBackbone.from_pretrained("microsoft/focalnet-tiny-lrf")
+
+        >>> inputs = processor(image, return_tensors="pt")
+        >>> outputs = model(**inputs)
+        ```"""
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        outputs = self.focalnet(pixel_values, output_hidden_states=True, return_dict=True)
+
+        hidden_states = outputs.reshaped_hidden_states
+
+        feature_maps = ()
+        for idx, stage in enumerate(self.stage_names):
+            if stage in self.out_features:
+                feature_maps += (hidden_states[idx],)
+
+        if not return_dict:
+            output = (feature_maps,)
+            if output_hidden_states:
+                output += (outputs.hidden_states,)
+            return output
+
+        return BackboneOutput(
+            feature_maps=feature_maps,
+            hidden_states=outputs.hidden_states if output_hidden_states else None,
+            attentions=None,
+        )
+
+
+__all__ = [
+    "FocalNetForImageClassification",
+    "FocalNetForMaskedImageModeling",
+    "FocalNetBackbone",
+    "FocalNetModel",
+    "FocalNetPreTrainedModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/superglue/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_superglue import *
+    from .image_processing_pil_superglue import *
+    from .image_processing_superglue import *
+    from .modeling_superglue import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/superglue/configuration_superglue.py

@@ -0,0 +1,92 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+@auto_docstring(checkpoint="magic-leap-community/superglue_indoor")
+@strict
+class SuperGlueConfig(PreTrainedConfig):
+    r"""
+    keypoint_detector_config (`Union[AutoConfig, dict]`,  *optional*, defaults to `SuperPointConfig`):
+        The config object or dictionary of the keypoint detector.
+    keypoint_encoder_sizes (`list[int]`, *optional*, defaults to `[32, 64, 128, 256]`):
+        The sizes of the keypoint encoder layers.
+    gnn_layers_types (`list[str]`, *optional*, defaults to `['self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross', 'self', 'cross']`):
+        The types of the GNN layers. Must be either 'self' or 'cross'.
+    sinkhorn_iterations (`int`, *optional*, defaults to 100):
+        The number of Sinkhorn iterations.
+    matching_threshold (`float`, *optional*, defaults to 0.0):
+        The matching threshold.
+
+    Examples:
+        ```python
+        >>> from transformers import SuperGlueConfig, SuperGlueModel
+
+        >>> # Initializing a SuperGlue superglue style configuration
+        >>> configuration = SuperGlueConfig()
+
+        >>> # Initializing a model from the superglue style configuration
+        >>> model = SuperGlueModel(configuration)
+
+        >>> # Accessing the model configuration
+        >>> configuration = model.config
+        ```
+    """
+
+    model_type = "superglue"
+    sub_configs = {"keypoint_detector_config": AutoConfig}
+
+    keypoint_detector_config: dict | PreTrainedConfig | None = None
+    hidden_size: int = 256
+    keypoint_encoder_sizes: list[int] | None = None
+    gnn_layers_types: list[str] | None = None
+    num_attention_heads: int = 4
+    sinkhorn_iterations: int = 100
+    matching_threshold: float = 0.0
+    initializer_range: float = 0.02
+    is_decoder: bool = False
+    attention_probs_dropout_prob: int | float = 0.0
+
+    def __post_init__(self, **kwargs):
+        self.gnn_layers_types = self.gnn_layers_types if self.gnn_layers_types is not None else ["self", "cross"] * 9
+        self.keypoint_encoder_sizes = (
+            self.keypoint_encoder_sizes if self.keypoint_encoder_sizes is not None else [32, 64, 128, 256]
+        )
+
+        if isinstance(self.keypoint_detector_config, dict):
+            self.keypoint_detector_config["model_type"] = self.keypoint_detector_config.get("model_type", "superpoint")
+            self.keypoint_detector_config = CONFIG_MAPPING[self.keypoint_detector_config["model_type"]](
+                **self.keypoint_detector_config
+            )
+        elif self.keypoint_detector_config is None:
+            self.keypoint_detector_config = CONFIG_MAPPING["superpoint"]()
+
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        # Check whether all gnn_layers_types are either 'self' or 'cross'
+        if not all(layer_type in ["self", "cross"] for layer_type in self.gnn_layers_types):
+            raise ValueError("All gnn_layers_types must be either 'self' or 'cross'")
+
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError("hidden_size % num_attention_heads is different from zero")
+
+
+__all__ = ["SuperGlueConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/superglue/image_processing_pil_superglue.py

@@ -0,0 +1,299 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor class for SuperGlue."""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature
+from ...image_utils import (
+    ImageInput,
+    ImageType,
+    PILImageResampling,
+    SizeDict,
+    get_image_type,
+    is_pil_image,
+    is_valid_image,
+    to_numpy_array,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+from ...utils.import_utils import requires
+
+
+if TYPE_CHECKING:
+    import torch
+
+    from .modeling_superglue import SuperGlueKeypointMatchingOutput
+
+
+def is_grayscale(image: np.ndarray):
+    if image.shape[0] == 1:
+        return True
+    return np.all(image[0, ...] == image[1, ...]) and np.all(image[1, ...] == image[2, ...])
+
+
+def convert_to_grayscale(image: ImageInput) -> ImageInput:
+    """
+    Converts an image to grayscale format using the NTSC formula. Only support numpy and PIL Image.
+
+    This function is supposed to return a 1-channel image, but it returns a 3-channel image with the same value in each
+    channel, because of an issue that is discussed in :
+    https://github.com/huggingface/transformers/pull/25786#issuecomment-1730176446
+
+    Args:
+        image (Image):
+            The image to convert.
+    """
+
+    if isinstance(image, np.ndarray):
+        if is_grayscale(image):
+            return image
+
+        gray_image = image[0, ...] * 0.2989 + image[1, ...] * 0.5870 + image[2, ...] * 0.1140
+        gray_image = np.stack([gray_image] * 3, axis=0)
+        return gray_image
+
+    if not isinstance(image, Image.Image):
+        return image
+
+    image = image.convert("L")
+    return image
+
+
+# Adapted from transformers.models.superglue.image_processing_superglue.validate_and_format_image_pairs
+def validate_and_format_image_pairs(images: ImageInput):
+    error_message = (
+        "Input images must be a one of the following :",
+        " - A pair of PIL images.",
+        " - A pair of 3D arrays.",
+        " - A list of pairs of PIL images.",
+        " - A list of pairs of 3D arrays.",
+    )
+
+    def _is_valid_image(image):
+        """images is a PIL Image or a 3D array."""
+        return is_pil_image(image) or (
+            is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3
+        )
+
+    if isinstance(images, list):
+        if len(images) == 2 and all((_is_valid_image(image)) for image in images):
+            return images
+        if all(
+            isinstance(image_pair, list)
+            and len(image_pair) == 2
+            and all(_is_valid_image(image) for image in image_pair)
+            for image_pair in images
+        ):
+            return [image for image_pair in images for image in image_pair]
+    raise ValueError(error_message)
+
+
+class SuperGlueImageProcessorKwargs(ImagesKwargs, total=False):
+    r"""
+    do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`):
+        Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method.
+    """
+
+    do_grayscale: bool
+
+
+@auto_docstring
+class SuperGlueImageProcessorPil(PilBackend):
+    valid_kwargs = SuperGlueImageProcessorKwargs
+    resample = PILImageResampling.BILINEAR
+    size = {"height": 480, "width": 640}
+    default_to_square = False
+    do_resize = True
+    do_rescale = True
+    rescale_factor = 1 / 255
+    do_normalize = None
+    do_grayscale = True
+
+    def __init__(self, **kwargs: Unpack[SuperGlueImageProcessorKwargs]):
+        super().__init__(**kwargs)
+
+    @auto_docstring
+    def preprocess(self, images: ImageInput, **kwargs: Unpack[SuperGlueImageProcessorKwargs]) -> BatchFeature:
+        return super().preprocess(images, **kwargs)
+
+    def _prepare_images_structure(self, images: ImageInput, **kwargs) -> ImageInput:
+        # we need to handle image pairs validation and flattening
+        images = self.fetch_images(images)
+        return validate_and_format_image_pairs(images)
+
+    def _preprocess(
+        self,
+        images: list[np.ndarray],
+        do_resize: bool,
+        size: SizeDict,
+        resample: PILImageResampling | None,
+        do_rescale: bool,
+        rescale_factor: float,
+        return_tensors: str | TensorType | None,
+        do_grayscale: bool = True,
+        **kwargs,
+    ) -> BatchFeature:
+        all_images = []
+        for image in images:
+            if do_resize:
+                image = self.resize(image=image, size=size, resample=resample)
+
+            if do_rescale:
+                image = self.rescale(image=image, scale=rescale_factor)
+
+            if do_grayscale:
+                image = convert_to_grayscale(image)
+
+            all_images.append(image)
+
+        # Convert back the flattened list of images into a list of pairs of images.
+        image_pairs = [all_images[i : i + 2] for i in range(0, len(all_images), 2)]
+
+        data = {"pixel_values": image_pairs}
+
+        return BatchFeature(data=data, tensor_type=return_tensors)
+
+    @requires(backends=("torch",))
+    def post_process_keypoint_matching(
+        self,
+        outputs: "SuperGlueKeypointMatchingOutput",
+        target_sizes: TensorType | list[tuple],
+        threshold: float = 0.0,
+    ) -> list[dict[str, "torch.Tensor"]]:
+        """
+        Converts the raw output of [`SuperGlueKeypointMatchingOutput`] into lists of keypoints, scores and descriptors
+        with coordinates absolute to the original image sizes.
+        Args:
+            outputs ([`SuperGlueKeypointMatchingOutput`]):
+                Raw outputs of the model.
+            target_sizes (`torch.Tensor` or `list[tuple[tuple[int, int]]]`, *optional*):
+                Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`tuple[int, int]`) containing the
+                target size `(height, width)` of each image in the batch. This must be the original image size (before
+                any processing).
+            threshold (`float`, *optional*, defaults to `0.0`):
+                Threshold to filter out the matches with low scores.
+        Returns:
+            `list[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image
+            of the pair, the matching scores and the matching indices.
+        """
+        import torch
+
+        if outputs.mask.shape[0] != len(target_sizes):
+            raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask")
+        if not all(len(target_size) == 2 for target_size in target_sizes):
+            raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
+
+        if isinstance(target_sizes, list):
+            image_pair_sizes = torch.tensor(target_sizes, device=outputs.mask.device)
+        else:
+            if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2:
+                raise ValueError(
+                    "Each element of target_sizes must contain the size (h, w) of each image of the batch"
+                )
+            image_pair_sizes = target_sizes
+
+        keypoints = outputs.keypoints.clone()
+        keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2)
+        keypoints = keypoints.to(torch.int32)
+
+        results = []
+        for mask_pair, keypoints_pair, matches, scores in zip(
+            outputs.mask, keypoints, outputs.matches[:, 0], outputs.matching_scores[:, 0]
+        ):
+            mask0 = mask_pair[0] > 0
+            mask1 = mask_pair[1] > 0
+            keypoints0 = keypoints_pair[0][mask0]
+            keypoints1 = keypoints_pair[1][mask1]
+            matches0 = matches[mask0]
+            scores0 = scores[mask0]
+
+            # Filter out matches with low scores, invalid matches, and out-of-bounds indices
+            valid_matches = (scores0 > threshold) & (matches0 > -1) & (matches0 < keypoints1.shape[0])
+
+            matched_keypoints0 = keypoints0[valid_matches]
+            matched_keypoints1 = keypoints1[matches0[valid_matches]]
+            matching_scores = scores0[valid_matches]
+
+            results.append(
+                {
+                    "keypoints0": matched_keypoints0,
+                    "keypoints1": matched_keypoints1,
+                    "matching_scores": matching_scores,
+                }
+            )
+
+        return results
+
+    def visualize_keypoint_matching(
+        self, images: ImageInput, keypoint_matching_output: list[dict[str, "torch.Tensor"]]
+    ) -> list["Image.Image"]:
+        """
+        Plots the image pairs side by side with the detected keypoints as well as the matching between them.
+
+        Args:
+            images (`ImageInput`):
+                Image pairs to plot. Same as `SuperGlueImageProcessor.preprocess`. Expects either a list of 2
+                images or a list of list of 2 images list with pixel values ranging from 0 to 255.
+            keypoint_matching_output (List[Dict[str, torch.Tensor]]]):
+                A post processed keypoint matching output
+
+        Returns:
+            `List[PIL.Image.Image]`: A list of PIL images, each containing the image pairs side by side with the detected
+            keypoints as well as the matching between them.
+        """
+        images = validate_and_format_image_pairs(images)
+        images = [to_numpy_array(image) for image in images]
+        image_pairs = [images[i : i + 2] for i in range(0, len(images), 2)]
+
+        results = []
+        for image_pair, pair_output in zip(image_pairs, keypoint_matching_output):
+            height0, width0 = image_pair[0].shape[:2]
+            height1, width1 = image_pair[1].shape[:2]
+            plot_image = np.zeros((max(height0, height1), width0 + width1, 3), dtype=np.uint8)
+            plot_image[:height0, :width0] = image_pair[0]
+            plot_image[:height1, width0:] = image_pair[1]
+
+            plot_image_pil = Image.fromarray(plot_image)
+            draw = ImageDraw.Draw(plot_image_pil)
+
+            keypoints0_x, keypoints0_y = pair_output["keypoints0"].unbind(1)
+            keypoints1_x, keypoints1_y = pair_output["keypoints1"].unbind(1)
+            for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip(
+                keypoints0_x, keypoints0_y, keypoints1_x, keypoints1_y, pair_output["matching_scores"]
+            ):
+                color = self._get_color(matching_score)
+                draw.line((keypoint0_x, keypoint0_y, keypoint1_x + width0, keypoint1_y), fill=color, width=3)
+                draw.ellipse((keypoint0_x - 2, keypoint0_y - 2, keypoint0_x + 2, keypoint0_y + 2), fill="black")
+                draw.ellipse(
+                    (keypoint1_x + width0 - 2, keypoint1_y - 2, keypoint1_x + width0 + 2, keypoint1_y + 2),
+                    fill="black",
+                )
+
+            results.append(plot_image_pil)
+        return results
+
+    def _get_color(self, score):
+        """Maps a score to a color."""
+        r = int(255 * (1 - score))
+        g = int(255 * score)
+        b = 0
+        return (r, g, b)
+
+
+__all__ = ["SuperGlueImageProcessorPil"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/superglue/image_processing_superglue.py

@@ -0,0 +1,325 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Image processor class for SuperGlue."""
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import group_images_by_shape, reorder_images
+from ...image_utils import (
+    ImageInput,
+    ImageType,
+    PILImageResampling,
+    SizeDict,
+    get_image_type,
+    is_pil_image,
+    is_valid_image,
+    to_numpy_array,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import (
+    TensorType,
+    auto_docstring,
+    is_vision_available,
+)
+
+
+if TYPE_CHECKING:
+    from .modeling_superglue import SuperGlueKeypointMatchingOutput
+
+if is_vision_available():
+    from PIL import Image, ImageDraw
+
+from torchvision.transforms.v2 import functional as tvF
+
+
+def _is_valid_image(image):
+    return is_pil_image(image) or (
+        is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3
+    )
+
+
+def validate_and_format_image_pairs(images: ImageInput):
+    error_message = (
+        "Input images must be a one of the following :",
+        " - A pair of PIL images.",
+        " - A pair of 3D arrays.",
+        " - A list of pairs of PIL images.",
+        " - A list of pairs of 3D arrays.",
+    )
+
+    def _is_valid_image(image):
+        """images is a PIL Image or a 3D array."""
+        return is_pil_image(image) or (
+            is_valid_image(image) and get_image_type(image) != ImageType.PIL and len(image.shape) == 3
+        )
+
+    if isinstance(images, list):
+        if len(images) == 2 and all((_is_valid_image(image)) for image in images):
+            return images
+        if all(
+            isinstance(image_pair, list)
+            and len(image_pair) == 2
+            and all(_is_valid_image(image) for image in image_pair)
+            for image_pair in images
+        ):
+            return [image for image_pair in images for image in image_pair]
+    raise ValueError(error_message)
+
+
+def is_grayscale(
+    image: "torch.Tensor",
+):
+    """Checks if an image is grayscale (all RGB channels are identical)."""
+    if image.ndim < 3 or image.shape[0 if image.ndim == 3 else 1] == 1:
+        return True
+    return torch.all(image[..., 0, :, :] == image[..., 1, :, :]) and torch.all(
+        image[..., 1, :, :] == image[..., 2, :, :]
+    )
+
+
+def convert_to_grayscale(
+    image: "torch.Tensor",
+) -> "torch.Tensor":
+    """
+    Converts an image to grayscale format using the NTSC formula. Only support torch.Tensor.
+
+    This function is supposed to return a 1-channel image, but it returns a 3-channel image with the same value in each
+    channel, because of an issue that is discussed in :
+    https://github.com/huggingface/transformers/pull/25786#issuecomment-1730176446
+
+    Args:
+        image (torch.Tensor):
+            The image to convert.
+    """
+    if is_grayscale(image):
+        return image
+    return tvF.rgb_to_grayscale(image, num_output_channels=3)
+
+
+class SuperGlueImageProcessorKwargs(ImagesKwargs, total=False):
+    r"""
+    do_grayscale (`bool`, *optional*, defaults to `self.do_grayscale`):
+        Whether to convert the image to grayscale. Can be overridden by `do_grayscale` in the `preprocess` method.
+    """
+
+    do_grayscale: bool
+
+
+@auto_docstring
+class SuperGlueImageProcessor(TorchvisionBackend):
+    valid_kwargs = SuperGlueImageProcessorKwargs
+    resample = PILImageResampling.BILINEAR
+    size = {"height": 480, "width": 640}
+    default_to_square = False
+    do_resize = True
+    do_rescale = True
+    rescale_factor = 1 / 255
+    do_normalize = None
+    do_grayscale = True
+
+    def __init__(self, **kwargs: Unpack[SuperGlueImageProcessorKwargs]):
+        super().__init__(**kwargs)
+
+    @auto_docstring
+    def preprocess(self, images: ImageInput, **kwargs: Unpack[SuperGlueImageProcessorKwargs]) -> BatchFeature:
+        return super().preprocess(images, **kwargs)
+
+    def _prepare_images_structure(
+        self,
+        images: ImageInput,
+        **kwargs,
+    ) -> ImageInput:
+        # we need to handle image pairs validation and flattening
+        images = self.fetch_images(images)
+        return validate_and_format_image_pairs(images)
+
+    def _preprocess(
+        self,
+        images: list["torch.Tensor"],
+        do_resize: bool,
+        size: SizeDict,
+        resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+        do_rescale: bool,
+        rescale_factor: float,
+        disable_grouping: bool | None,
+        return_tensors: str | TensorType | None,
+        do_grayscale: bool = True,
+        **kwargs,
+    ) -> BatchFeature:
+        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+        processed_images_grouped = {}
+
+        for shape, stacked_images in grouped_images.items():
+            if do_resize:
+                stacked_images = self.resize(stacked_images, size=size, resample=resample)
+            processed_images_grouped[shape] = stacked_images
+        resized_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+        processed_images_grouped = {}
+        for shape, stacked_images in grouped_images.items():
+            if do_rescale:
+                stacked_images = self.rescale(stacked_images, rescale_factor)
+            if do_grayscale:
+                stacked_images = convert_to_grayscale(stacked_images)
+            processed_images_grouped[shape] = stacked_images
+
+        processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+        # Convert back to pairs format
+        image_pairs = [processed_images[i : i + 2] for i in range(0, len(processed_images), 2)]
+
+        # Stack each pair into a single tensor to match slow processor format
+        stacked_pairs = [torch.stack(pair, dim=0) for pair in image_pairs]
+
+        # Return in same format as slow processor
+
+        return BatchFeature(data={"pixel_values": stacked_pairs}, tensor_type=return_tensors)
+
+    def post_process_keypoint_matching(
+        self,
+        outputs: "SuperGlueKeypointMatchingOutput",
+        target_sizes: TensorType | list[tuple],
+        threshold: float = 0.0,
+    ) -> list[dict[str, torch.Tensor]]:
+        """
+        Converts the raw output of [`SuperGlueKeypointMatchingOutput`] into lists of keypoints, scores and descriptors
+        with coordinates absolute to the original image sizes.
+        Args:
+            outputs ([`SuperGlueKeypointMatchingOutput`]):
+                Raw outputs of the model.
+            target_sizes (`torch.Tensor` or `list[tuple[tuple[int, int]]]`, *optional*):
+                Tensor of shape `(batch_size, 2, 2)` or list of tuples of tuples (`tuple[int, int]`) containing the
+                target size `(height, width)` of each image in the batch. This must be the original image size (before
+                any processing).
+            threshold (`float`, *optional*, defaults to `0.0`):
+                Threshold to filter out the matches with low scores.
+        Returns:
+            `list[Dict]`: A list of dictionaries, each dictionary containing the keypoints in the first and second image
+            of the pair, the matching scores and the matching indices.
+        """
+        if outputs.mask.shape[0] != len(target_sizes):
+            raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the mask")
+        if not all(len(target_size) == 2 for target_size in target_sizes):
+            raise ValueError("Each element of target_sizes must contain the size (h, w) of each image of the batch")
+
+        if isinstance(target_sizes, list):
+            image_pair_sizes = torch.tensor(target_sizes, device=outputs.mask.device)
+        else:
+            if target_sizes.shape[1] != 2 or target_sizes.shape[2] != 2:
+                raise ValueError(
+                    "Each element of target_sizes must contain the size (h, w) of each image of the batch"
+                )
+            image_pair_sizes = target_sizes
+
+        keypoints = outputs.keypoints.clone()
+        keypoints = keypoints * image_pair_sizes.flip(-1).reshape(-1, 2, 1, 2)
+        keypoints = keypoints.to(torch.int32)
+
+        results = []
+        for mask_pair, keypoints_pair, matches, scores in zip(
+            outputs.mask, keypoints, outputs.matches[:, 0], outputs.matching_scores[:, 0]
+        ):
+            mask0 = mask_pair[0] > 0
+            mask1 = mask_pair[1] > 0
+            keypoints0 = keypoints_pair[0][mask0]
+            keypoints1 = keypoints_pair[1][mask1]
+            matches0 = matches[mask0]
+            scores0 = scores[mask0]
+
+            # Filter out matches with low scores, invalid matches, and out-of-bounds indices
+            valid_matches = (scores0 > threshold) & (matches0 > -1) & (matches0 < keypoints1.shape[0])
+
+            matched_keypoints0 = keypoints0[valid_matches]
+            matched_keypoints1 = keypoints1[matches0[valid_matches]]
+            matching_scores = scores0[valid_matches]
+
+            results.append(
+                {
+                    "keypoints0": matched_keypoints0,
+                    "keypoints1": matched_keypoints1,
+                    "matching_scores": matching_scores,
+                }
+            )
+
+        return results
+
+    def visualize_keypoint_matching(
+        self,
+        images,
+        keypoint_matching_output: list[dict[str, torch.Tensor]],
+    ) -> list["Image.Image"]:
+        """
+        Plots the image pairs side by side with the detected keypoints as well as the matching between them.
+
+        Args:
+            images:
+                Image pairs to plot. Same as `EfficientLoFTRImageProcessor.preprocess`. Expects either a list of 2
+                images or a list of list of 2 images list with pixel values ranging from 0 to 255.
+            keypoint_matching_output (List[Dict[str, torch.Tensor]]]):
+                A post processed keypoint matching output
+
+        Returns:
+            `List[PIL.Image.Image]`: A list of PIL images, each containing the image pairs side by side with the detected
+            keypoints as well as the matching between them.
+        """
+
+        images = validate_and_format_image_pairs(images)
+        images = [to_numpy_array(image) for image in images]
+        image_pairs = [images[i : i + 2] for i in range(0, len(images), 2)]
+
+        results = []
+        for image_pair, pair_output in zip(image_pairs, keypoint_matching_output):
+            height0, width0 = image_pair[0].shape[:2]
+            height1, width1 = image_pair[1].shape[:2]
+            plot_image = torch.zeros((max(height0, height1), width0 + width1, 3), dtype=torch.uint8)
+            plot_image[:height0, :width0] = torch.from_numpy(image_pair[0])
+            plot_image[:height1, width0:] = torch.from_numpy(image_pair[1])
+
+            plot_image_pil = Image.fromarray(plot_image.numpy())
+            draw = ImageDraw.Draw(plot_image_pil)
+
+            keypoints0_x, keypoints0_y = pair_output["keypoints0"].unbind(1)
+            keypoints1_x, keypoints1_y = pair_output["keypoints1"].unbind(1)
+            for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip(
+                keypoints0_x, keypoints0_y, keypoints1_x, keypoints1_y, pair_output["matching_scores"]
+            ):
+                color = self._get_color(matching_score)
+                draw.line(
+                    (keypoint0_x, keypoint0_y, keypoint1_x + width0, keypoint1_y),
+                    fill=color,
+                    width=3,
+                )
+                draw.ellipse((keypoint0_x - 2, keypoint0_y - 2, keypoint0_x + 2, keypoint0_y + 2), fill="black")
+                draw.ellipse(
+                    (keypoint1_x + width0 - 2, keypoint1_y - 2, keypoint1_x + width0 + 2, keypoint1_y + 2),
+                    fill="black",
+                )
+
+            results.append(plot_image_pil)
+        return results
+
+    def _get_color(self, score):
+        """Maps a score to a color."""
+        r = int(255 * (1 - score))
+        g = int(255 * score)
+        b = 0
+        return r, g, b
+
+
+__all__ = ["SuperGlueImageProcessor"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/superglue/modeling_superglue.py

@@ -0,0 +1,758 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch SuperGlue model."""
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from transformers import PreTrainedModel
+from transformers.models.superglue.configuration_superglue import SuperGlueConfig
+
+from ... import initialization as init
+from ...masking_utils import create_bidirectional_mask
+from ...utils import ModelOutput, auto_docstring, logging
+from ..auto import AutoModelForKeypointDetection
+
+
+logger = logging.get_logger(__name__)
+
+
+def concat_pairs(tensor_tuple0: tuple[torch.Tensor], tensor_tuple1: tuple[torch.Tensor]) -> tuple[torch.Tensor]:
+    """
+    Concatenate two tuples of tensors pairwise
+
+    Args:
+        tensor_tuple0 (`tuple[torch.Tensor]`):
+            Tuple of tensors.
+        tensor_tuple1 (`tuple[torch.Tensor]`):
+            Tuple of tensors.
+
+    Returns:
+        (`tuple[torch.Tensor]`): Tuple of concatenated tensors.
+    """
+    return tuple(torch.cat([tensor0, tensor1]) for tensor0, tensor1 in zip(tensor_tuple0, tensor_tuple1))
+
+
+def normalize_keypoints(keypoints: torch.Tensor, height: int, width: int) -> torch.Tensor:
+    """
+    Normalize keypoints locations based on image image_shape
+
+    Args:
+        keypoints (`torch.Tensor` of shape `(batch_size, num_keypoints, 2)`):
+            Keypoints locations in (x, y) format.
+        height (`int`):
+            Image height.
+        width (`int`):
+            Image width.
+
+    Returns:
+        Normalized keypoints locations of shape (`torch.Tensor` of shape `(batch_size, num_keypoints, 2)`).
+    """
+    size = torch.tensor([width, height], device=keypoints.device, dtype=keypoints.dtype)[None]
+    center = size / 2
+    scaling = size.max(1, keepdim=True).values * 0.7
+    return (keypoints - center[:, None, :]) / scaling[:, None, :]
+
+
+def log_sinkhorn_iterations(
+    log_cost_matrix: torch.Tensor,
+    log_source_distribution: torch.Tensor,
+    log_target_distribution: torch.Tensor,
+    num_iterations: int,
+) -> torch.Tensor:
+    """
+    Perform Sinkhorn Normalization in Log-space for stability
+
+    Args:
+        log_cost_matrix (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`):
+            Logarithm of the cost matrix.
+        log_source_distribution (`torch.Tensor` of shape `(batch_size, num_rows)`):
+            Logarithm of the source distribution.
+        log_target_distribution (`torch.Tensor` of shape `(batch_size, num_columns)`):
+            Logarithm of the target distribution.
+
+    Returns:
+        log_cost_matrix (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`): Logarithm of the optimal
+        transport matrix.
+    """
+    log_u_scaling = torch.zeros_like(log_source_distribution)
+    log_v_scaling = torch.zeros_like(log_target_distribution)
+    for _ in range(num_iterations):
+        log_u_scaling = log_source_distribution - torch.logsumexp(log_cost_matrix + log_v_scaling.unsqueeze(1), dim=2)
+        log_v_scaling = log_target_distribution - torch.logsumexp(log_cost_matrix + log_u_scaling.unsqueeze(2), dim=1)
+    return log_cost_matrix + log_u_scaling.unsqueeze(2) + log_v_scaling.unsqueeze(1)
+
+
+def log_optimal_transport(scores: torch.Tensor, reg_param: torch.Tensor, iterations: int) -> torch.Tensor:
+    """
+    Perform Differentiable Optimal Transport in Log-space for stability
+
+    Args:
+        scores: (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`):
+            Cost matrix.
+        reg_param: (`torch.Tensor` of shape `(batch_size, 1, 1)`):
+            Regularization parameter.
+        iterations: (`int`):
+            Number of Sinkhorn iterations.
+
+    Returns:
+        log_optimal_transport_matrix: (`torch.Tensor` of shape `(batch_size, num_rows, num_columns)`): Logarithm of the
+        optimal transport matrix.
+    """
+    batch_size, num_rows, num_columns = scores.shape
+    one_tensor = scores.new_tensor(1)
+    num_rows_tensor, num_columns_tensor = (num_rows * one_tensor).to(scores), (num_columns * one_tensor).to(scores)
+
+    source_reg_param = reg_param.expand(batch_size, num_rows, 1)
+    target_reg_param = reg_param.expand(batch_size, 1, num_columns)
+    reg_param = reg_param.expand(batch_size, 1, 1)
+
+    couplings = torch.cat([torch.cat([scores, source_reg_param], -1), torch.cat([target_reg_param, reg_param], -1)], 1)
+
+    log_normalization = -(num_rows_tensor + num_columns_tensor).log()
+    log_source_distribution = torch.cat(
+        [log_normalization.expand(num_rows), num_columns_tensor.log()[None] + log_normalization]
+    )
+    log_target_distribution = torch.cat(
+        [log_normalization.expand(num_columns), num_rows_tensor.log()[None] + log_normalization]
+    )
+    log_source_distribution, log_target_distribution = (
+        log_source_distribution[None].expand(batch_size, -1),
+        log_target_distribution[None].expand(batch_size, -1),
+    )
+
+    log_optimal_transport_matrix = log_sinkhorn_iterations(
+        couplings, log_source_distribution, log_target_distribution, num_iterations=iterations
+    )
+    log_optimal_transport_matrix = log_optimal_transport_matrix - log_normalization  # multiply probabilities by M+N
+    return log_optimal_transport_matrix
+
+
+def arange_like(x, dim: int) -> torch.Tensor:
+    return x.new_ones(x.shape[dim]).cumsum(0) - 1
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for outputs of SuperGlue keypoint matching models. Due to the nature of keypoint detection and matching, the number
+    of keypoints is not fixed and can vary from image to image, which makes batching non-trivial. In the batch of
+    images, the maximum number of matches is set as the dimension of the matches and matching scores. The mask tensor is
+    used to indicate which values in the keypoints, matches and matching_scores tensors are keypoint matching
+    information.
+    """
+)
+@dataclass
+class SuperGlueKeypointMatchingOutput(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*):
+        Loss computed during training.
+    matches (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`):
+        Index of keypoint matched in the other image.
+    matching_scores (`torch.FloatTensor` of shape `(batch_size, 2, num_matches)`):
+        Scores of predicted matches.
+    keypoints (`torch.FloatTensor` of shape `(batch_size, num_keypoints, 2)`):
+        Absolute (x, y) coordinates of predicted keypoints in a given image.
+    mask (`torch.IntTensor` of shape `(batch_size, num_keypoints)`):
+        Mask indicating which values in matches and matching_scores are keypoint matching information.
+    hidden_states (`tuple[torch.FloatTensor, ...]`, *optional*):
+        Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(batch_size, 2, num_channels,
+        num_keypoints)`, returned when `output_hidden_states=True` is passed or when
+        `config.output_hidden_states=True`)
+    attentions (`tuple[torch.FloatTensor, ...]`, *optional*):
+        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, 2, num_heads, num_keypoints,
+        num_keypoints)`, returned when `output_attentions=True` is passed or when `config.output_attentions=True`)
+    """
+
+    loss: torch.FloatTensor | None = None
+    matches: torch.FloatTensor | None = None
+    matching_scores: torch.FloatTensor | None = None
+    keypoints: torch.FloatTensor | None = None
+    mask: torch.IntTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+
+
+class SuperGlueMultiLayerPerceptron(nn.Module):
+    def __init__(self, config: SuperGlueConfig, in_channels: int, out_channels: int) -> None:
+        super().__init__()
+        self.linear = nn.Linear(in_channels, out_channels)
+        self.batch_norm = nn.BatchNorm1d(out_channels)
+        self.activation = nn.ReLU()
+
+    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+        hidden_state = self.linear(hidden_state)
+        hidden_state = hidden_state.transpose(-1, -2)
+        hidden_state = self.batch_norm(hidden_state)
+        hidden_state = hidden_state.transpose(-1, -2)
+        hidden_state = self.activation(hidden_state)
+        return hidden_state
+
+
+class SuperGlueKeypointEncoder(nn.Module):
+    def __init__(self, config: SuperGlueConfig) -> None:
+        super().__init__()
+        layer_sizes = config.keypoint_encoder_sizes
+        hidden_size = config.hidden_size
+        # 3 here consists of 2 for the (x, y) coordinates and 1 for the score of the keypoint
+        encoder_channels = [3] + layer_sizes + [hidden_size]
+
+        layers = [
+            SuperGlueMultiLayerPerceptron(config, encoder_channels[i - 1], encoder_channels[i])
+            for i in range(1, len(encoder_channels) - 1)
+        ]
+        layers.append(nn.Linear(encoder_channels[-2], encoder_channels[-1]))
+        self.encoder = nn.ModuleList(layers)
+
+    def forward(
+        self,
+        keypoints: torch.Tensor,
+        scores: torch.Tensor,
+        output_hidden_states: bool | None = False,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor] | None]:
+        scores = scores.unsqueeze(2)
+        hidden_state = torch.cat([keypoints, scores], dim=2)
+        all_hidden_states = () if output_hidden_states else None
+        for layer in self.encoder:
+            hidden_state = layer(hidden_state)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_state,)
+        return hidden_state, all_hidden_states
+
+
+class SuperGlueSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({config.num_attention_heads})"
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+        self.is_decoder = config.is_decoder
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor | None = None,
+        encoder_attention_mask: torch.FloatTensor | None = None,
+        output_attentions: bool | None = False,
+    ) -> tuple[torch.Tensor]:
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+        current_states = encoder_hidden_states if is_cross_attention else hidden_states
+        attention_mask = encoder_attention_mask if is_cross_attention else attention_mask
+
+        batch_size = hidden_states.shape[0]
+        key_layer = (
+            self.key(current_states)
+            .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+            .transpose(1, 2)
+        )
+        value_layer = (
+            self.value(current_states)
+            .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+            .transpose(1, 2)
+        )
+        query_layer = (
+            self.query(hidden_states)
+            .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+            .transpose(1, 2)
+        )
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in SuperGlueModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs = self.dropout(attention_probs)
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+        if self.is_decoder:
+            outputs = outputs + (None,)
+        return outputs
+
+
+class SuperGlueSelfOutput(nn.Module):
+    def __init__(self, config: SuperGlueConfig):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor, *args) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        return hidden_states
+
+
+SUPERGLUE_SELF_ATTENTION_CLASSES = {
+    "eager": SuperGlueSelfAttention,
+}
+
+
+class SuperGlueAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = SUPERGLUE_SELF_ATTENTION_CLASSES[config._attn_implementation](config)
+        self.output = SuperGlueSelfOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        output_attentions: bool | None = False,
+    ) -> tuple[torch.Tensor]:
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+class SuperGlueAttentionalPropagation(nn.Module):
+    def __init__(self, config: SuperGlueConfig) -> None:
+        super().__init__()
+        hidden_size = config.hidden_size
+        self.attention = SuperGlueAttention(config)
+        mlp_channels = [hidden_size * 2, hidden_size * 2, hidden_size]
+        layers = [
+            SuperGlueMultiLayerPerceptron(config, mlp_channels[i - 1], mlp_channels[i])
+            for i in range(1, len(mlp_channels) - 1)
+        ]
+        layers.append(nn.Linear(mlp_channels[-2], mlp_channels[-1]))
+        self.mlp = nn.ModuleList(layers)
+
+    def forward(
+        self,
+        descriptors: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor] | None, tuple[torch.Tensor] | None]:
+        attention_outputs = self.attention(
+            descriptors,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+        )
+        output = attention_outputs[0]
+        attention = attention_outputs[1:]
+
+        hidden_state = torch.cat([descriptors, output], dim=2)
+
+        all_hidden_states = () if output_hidden_states else None
+        for layer in self.mlp:
+            hidden_state = layer(hidden_state)
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_state,)
+
+        return hidden_state, all_hidden_states, attention
+
+
+class SuperGlueAttentionalGNN(nn.Module):
+    def __init__(self, config: SuperGlueConfig) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.layers_types = config.gnn_layers_types
+        self.layers = nn.ModuleList([SuperGlueAttentionalPropagation(config) for _ in range(len(self.layers_types))])
+
+    def forward(
+        self,
+        descriptors: torch.Tensor,
+        mask: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool | None = False,
+    ) -> tuple[torch.Tensor, tuple | None, tuple | None]:
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        batch_size, num_keypoints, _ = descriptors.shape
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (descriptors,)
+
+        for gnn_layer, layer_type in zip(self.layers, self.layers_types):
+            encoder_hidden_states = None
+            encoder_attention_mask = None
+            if layer_type == "cross":
+                encoder_hidden_states = (
+                    descriptors.reshape(-1, 2, num_keypoints, self.hidden_size)
+                    .flip(1)
+                    .reshape(batch_size, num_keypoints, self.hidden_size)
+                )
+                encoder_attention_mask = (
+                    mask.reshape(-1, 2, 1, 1, num_keypoints).flip(1).reshape(batch_size, 1, 1, num_keypoints)
+                    if mask is not None
+                    else None
+                )
+
+            gnn_outputs = gnn_layer(
+                descriptors,
+                attention_mask=mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                output_hidden_states=output_hidden_states,
+                output_attentions=output_attentions,
+            )
+            delta = gnn_outputs[0]
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + gnn_outputs[1]
+            if output_attentions:
+                all_attentions = all_attentions + gnn_outputs[2]
+
+            descriptors = descriptors + delta
+        return descriptors, all_hidden_states, all_attentions
+
+
+class SuperGlueFinalProjection(nn.Module):
+    def __init__(self, config: SuperGlueConfig) -> None:
+        super().__init__()
+        hidden_size = config.hidden_size
+        self.final_proj = nn.Linear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, descriptors: torch.Tensor) -> torch.Tensor:
+        return self.final_proj(descriptors)
+
+
+@auto_docstring
+class SuperGluePreTrainedModel(PreTrainedModel):
+    config: SuperGlueConfig
+    base_model_prefix = "superglue"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+
+    @torch.no_grad()
+    def _init_weights(self, module: nn.Module) -> None:
+        """Initialize the weights"""
+        super()._init_weights(module)
+        if hasattr(module, "bin_score"):
+            init.ones_(module.bin_score)
+
+
+@auto_docstring(
+    custom_intro="""
+    SuperGlue model taking images as inputs and outputting the matching of them.
+    """
+)
+class SuperGlueForKeypointMatching(SuperGluePreTrainedModel):
+    """SuperGlue feature matching middle-end
+
+    Given two sets of keypoints and locations, we determine the
+    correspondences by:
+      1. Keypoint Encoding (normalization + visual feature and location fusion)
+      2. Graph Neural Network with multiple self and cross-attention layers
+      3. Final projection layer
+      4. Optimal Transport Layer (a differentiable Hungarian matching algorithm)
+      5. Thresholding matrix based on mutual exclusivity and a match_threshold
+
+    The correspondence ids use -1 to indicate non-matching points.
+
+    Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew
+    Rabinovich. SuperGlue: Learning Feature Matching with Graph Neural
+    Networks. In CVPR, 2020. https://huggingface.co/papers/1911.11763
+    """
+
+    def __init__(self, config: SuperGlueConfig) -> None:
+        super().__init__(config)
+
+        self.keypoint_detector = AutoModelForKeypointDetection.from_config(config.keypoint_detector_config)
+
+        self.keypoint_encoder = SuperGlueKeypointEncoder(config)
+        self.gnn = SuperGlueAttentionalGNN(config)
+        self.final_projection = SuperGlueFinalProjection(config)
+
+        bin_score = torch.nn.Parameter(torch.tensor(1.0))
+        self.register_parameter("bin_score", bin_score)
+
+        self.post_init()
+
+    def _match_image_pair(
+        self,
+        keypoints: torch.Tensor,
+        descriptors: torch.Tensor,
+        scores: torch.Tensor,
+        height: int,
+        width: int,
+        mask: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, tuple, tuple]:
+        """
+        Perform keypoint matching between two images.
+
+        Args:
+            keypoints (`torch.Tensor` of shape `(batch_size, 2, num_keypoints, 2)`):
+                Keypoints detected in the pair of image.
+            descriptors (`torch.Tensor` of shape `(batch_size, 2, descriptor_dim, num_keypoints)`):
+                Descriptors of the keypoints detected in the image pair.
+            scores (`torch.Tensor` of shape `(batch_size, 2, num_keypoints)`):
+                Confidence scores of the keypoints detected in the image pair.
+            height (`int`): Image height.
+            width (`int`): Image width.
+            mask (`torch.Tensor` of shape `(batch_size, 2, num_keypoints)`, *optional*):
+                Mask indicating which values in the keypoints, matches and matching_scores tensors are keypoint matching
+                information.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors. Default to `config.output_attentions`.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. Default to `config.output_hidden_states`.
+
+        Returns:
+            matches (`torch.Tensor` of shape `(batch_size, 2, num_keypoints)`):
+                For each image pair, for each keypoint in image0, the index of the keypoint in image1 that was matched
+                with. And for each keypoint in image1, the index of the keypoint in image0 that was matched with.
+            matching_scores (`torch.Tensor` of shape `(batch_size, 2, num_keypoints)`):
+                Scores of predicted matches for each image pair
+            all_hidden_states (`tuple(torch.FloatTensor)`, *optional*):
+                Tuple of `torch.FloatTensor` (one for the output of each stage) of shape `(1, 2, num_keypoints,
+                num_channels)`.
+            all_attentions (`tuple(torch.FloatTensor)`, *optional*):
+                Tuple of `torch.FloatTensor` (one for each layer) of shape `(1, 2, num_heads, num_keypoints,
+                num_keypoints)`.
+        """
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        if keypoints.shape[2] == 0:  # no keypoints
+            shape = keypoints.shape[:-1]
+            return (
+                keypoints.new_full(shape, -1, dtype=torch.int),
+                keypoints.new_zeros(shape),
+                all_hidden_states,
+                all_attentions,
+            )
+
+        batch_size, _, num_keypoints, _ = keypoints.shape
+        # (batch_size, 2, num_keypoints, 2) -> (batch_size * 2, num_keypoints, 2)
+        keypoints = keypoints.reshape(batch_size * 2, num_keypoints, 2)
+        descriptors = descriptors.reshape(batch_size * 2, num_keypoints, self.config.hidden_size)
+        scores = scores.reshape(batch_size * 2, num_keypoints)
+        mask = mask.reshape(batch_size * 2, num_keypoints) if mask is not None else None
+
+        # Keypoint normalization
+        keypoints = normalize_keypoints(keypoints, height, width)
+
+        encoded_keypoints = self.keypoint_encoder(keypoints, scores, output_hidden_states=output_hidden_states)
+
+        last_hidden_state = encoded_keypoints[0]
+
+        # Keypoint MLP encoder.
+        descriptors = descriptors + last_hidden_state
+
+        extended_attention_mask = create_bidirectional_mask(
+            config=self.config,
+            inputs_embeds=descriptors[:, 0:1, :],  # force q_len == 1
+            attention_mask=mask,
+        )
+
+        # Multi-layer Transformer network.
+        gnn_outputs = self.gnn(
+            descriptors,
+            mask=extended_attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+        descriptors = gnn_outputs[0]
+
+        # Final MLP projection.
+        projected_descriptors = self.final_projection(descriptors)
+
+        # (batch_size * 2, num_keypoints, descriptor_dim) -> (batch_size, 2, num_keypoints, descriptor_dim)
+        final_descriptors = projected_descriptors.reshape(batch_size, 2, num_keypoints, self.config.hidden_size)
+        final_descriptors0 = final_descriptors[:, 0]
+        final_descriptors1 = final_descriptors[:, 1]
+
+        # Compute matching descriptor distance.
+        scores = final_descriptors0 @ final_descriptors1.transpose(1, 2)
+        scores = scores / self.config.hidden_size**0.5
+
+        if mask is not None:
+            mask = mask.reshape(batch_size, 2, num_keypoints)
+            mask0 = mask[:, 0].unsqueeze(2)
+            mask1 = mask[:, 1].unsqueeze(1)
+            mask = torch.logical_and(mask0, mask1)
+            scores = scores.masked_fill(mask == 0, torch.finfo(scores.dtype).min)
+
+        # Run the optimal transport.
+        scores = log_optimal_transport(scores, self.bin_score, iterations=self.config.sinkhorn_iterations)
+
+        # Get the matches with score above "match_threshold".
+        max0 = scores[:, :-1, :-1].max(2)
+        max1 = scores[:, :-1, :-1].max(1)
+        indices0 = max0.indices
+        indices1 = max1.indices
+        mutual0 = arange_like(indices0, 1)[None] == indices1.gather(1, indices0)
+        mutual1 = arange_like(indices1, 1)[None] == indices0.gather(1, indices1)
+        zero = scores.new_tensor(0)
+        matching_scores0 = torch.where(mutual0, max0.values.exp(), zero)
+        matching_scores0 = torch.where(matching_scores0 > self.config.matching_threshold, matching_scores0, zero)
+        matching_scores1 = torch.where(mutual1, matching_scores0.gather(1, indices1), zero)
+        valid0 = mutual0 & (matching_scores0 > zero)
+        valid1 = mutual1 & valid0.gather(1, indices1)
+        matches0 = torch.where(valid0, indices0, indices0.new_tensor(-1))
+        matches1 = torch.where(valid1, indices1, indices1.new_tensor(-1))
+
+        matches = torch.cat([matches0, matches1], dim=1).reshape(batch_size, 2, -1)
+        matching_scores = torch.cat([matching_scores0, matching_scores1], dim=1).reshape(batch_size, 2, -1)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + encoded_keypoints[1]
+            all_hidden_states = all_hidden_states + gnn_outputs[1]
+            all_hidden_states = all_hidden_states + (projected_descriptors,)
+            all_hidden_states = tuple(
+                x.reshape(batch_size, 2, num_keypoints, -1).transpose(-1, -2) for x in all_hidden_states
+            )
+        if output_attentions:
+            all_attentions = all_attentions + gnn_outputs[2]
+            all_attentions = tuple(x.reshape(batch_size, 2, -1, num_keypoints, num_keypoints) for x in all_attentions)
+
+        return (
+            matches,
+            matching_scores,
+            all_hidden_states,
+            all_attentions,
+        )
+
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
+        labels: torch.LongTensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | SuperGlueKeypointMatchingOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, AutoModel
+        >>> import torch
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "https://github.com/magicleap/SuperGluePretrainedNetwork/blob/master/assets/phototourism_sample_images/london_bridge_78916675_4568141288.jpg?raw=true"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image_1 = Image.open(BytesIO(response.read()))
+
+        >>> url = "https://github.com/magicleap/SuperGluePretrainedNetwork/blob/master/assets/phototourism_sample_images/london_bridge_19481797_2295892421.jpg?raw=true"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image_2 = Image.open(BytesIO(response.read()))
+
+        >>> images = [image_1, image_2]
+
+        >>> processor = AutoImageProcessor.from_pretrained("magic-leap-community/superglue_outdoor")
+        >>> model = AutoModel.from_pretrained("magic-leap-community/superglue_outdoor")
+
+        >>> with torch.no_grad():
+        >>>     inputs = processor(images, return_tensors="pt")
+        >>>     outputs = model(**inputs)
+        ```"""
+        loss = None
+        if labels is not None:
+            raise ValueError("SuperGlue is not trainable, no labels should be provided.")
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        if pixel_values.ndim != 5 or pixel_values.size(1) != 2:
+            raise ValueError("Input must be a 5D tensor of shape (batch_size, 2, num_channels, height, width)")
+
+        batch_size, _, channels, height, width = pixel_values.shape
+        pixel_values = pixel_values.reshape(batch_size * 2, channels, height, width)
+        keypoint_detections = self.keypoint_detector(pixel_values)
+
+        keypoints, scores, descriptors, mask = keypoint_detections[:4]
+        keypoints = keypoints.reshape(batch_size, 2, -1, 2).to(pixel_values)
+        scores = scores.reshape(batch_size, 2, -1).to(pixel_values)
+        descriptors = descriptors.reshape(batch_size, 2, -1, self.config.hidden_size).to(pixel_values)
+        mask = mask.reshape(batch_size, 2, -1)
+
+        absolute_keypoints = keypoints.clone()
+        absolute_keypoints[:, :, :, 0] = absolute_keypoints[:, :, :, 0] * width
+        absolute_keypoints[:, :, :, 1] = absolute_keypoints[:, :, :, 1] * height
+
+        matches, matching_scores, hidden_states, attentions = self._match_image_pair(
+            absolute_keypoints,
+            descriptors,
+            scores,
+            height,
+            width,
+            mask=mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+        )
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [loss, matches, matching_scores, keypoints, mask, hidden_states, attentions]
+                if v is not None
+            )
+
+        return SuperGlueKeypointMatchingOutput(
+            loss=loss,
+            matches=matches,
+            matching_scores=matching_scores,
+            keypoints=keypoints,
+            mask=mask,
+            hidden_states=hidden_states,
+            attentions=attentions,
+        )
+
+
+__all__ = ["SuperGluePreTrainedModel", "SuperGlueForKeypointMatching"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/cache/owt_llmclean_qwen36_35b_articlefull_full_rev8/audit.jsonl

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