Add files using upload-large-folder tool

LICENSES.md

@@ -0,0 +1,31 @@
+# Licenses for the TRELLIS.2 MLX 8-bit Bundle
+
+This repository combines derivatives of three separately licensed models. No
+single license replaces the component licenses below.
+
+| Component | Files | Governing terms |
+| --- | --- | --- |
+| Microsoft TRELLIS.2-4B | `ckpts/`, `pipeline.json`, `texturing_pipeline.json` | [MIT License](LICENSE_TRELLIS2) |
+| Meta DINOv3 ViT-L/16 | `dinov3/` | [DINOv3 License](LICENSE_DINOV3.md) |
+| BRIA RMBG-2.0 | `rmbg/` | [CC BY-NC 4.0 attribution and change notice](LICENSE_RMBG2.md) |
+
+The quantization metadata, packaging, and model card describe changes made for
+`mlx-spatial`; they do not relax or replace any upstream term. You must comply
+with every license that applies to the files you use or redistribute.
+
+Because the bundle includes RMBG-2.0, use of the complete bundle is limited to
+non-commercial purposes unless you obtain separate commercial rights from
+BRIA. The DINOv3 agreement also includes redistribution, trade-control, and
+other conditions that must be read in full.
+
+This repository is an unofficial derivative and is not affiliated with or
+endorsed by Microsoft, Meta, BRIA, or the original authors.
+
+## Upstream References
+
+- TRELLIS.2-4B: <https://huggingface.co/microsoft/TRELLIS.2-4B>
+- TRELLIS.2 source: <https://github.com/microsoft/TRELLIS.2>
+- DINOv3 ViT-L/16: <https://huggingface.co/facebook/dinov3-vitl16-pretrain-lvd1689m>
+- DINOv3 source: <https://github.com/facebookresearch/dinov3>
+- RMBG-2.0: <https://huggingface.co/briaai/RMBG-2.0>
+- RMBG-2.0 source: <https://github.com/Bria-AI/RMBG-2.0>

LICENSE_DINOV3.md

@@ -0,0 +1,66 @@
+# DINOv3 License
+
+*Last Updated: August 19, 2025*
+
+**“Agreement”** means the terms and conditions for use, reproduction, distribution and modification of the DINO Materials set forth herein.
+
+**“DINO Materials”** means, collectively, Documentation and the models, software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, and other elements of the foregoing distributed by Meta and made available under this Agreement.
+
+**“Documentation”** means the specifications, manuals and documentation accompanying
+DINO Materials distributed by Meta.
+
+**“Licensee”** or **“you”** means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
+
+**“Meta”** or **“we”** means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) or Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+**“Sanctions”** means any economic or trade sanctions or restrictions administered or enforced by the United States (including the Office of Foreign Assets Control of the U.S. Department of the Treasury (“OFAC”), the U.S. Department of State and the U.S. Department of Commerce), the United Nations, the European Union, or the United Kingdom.
+
+**“Trade Controls”** means any of the following: Sanctions and applicable export and import controls.
+
+By clicking “I Accept” below or by using or distributing any portion or element of the DINO Materials, you agree to be bound by this Agreement.
+
+## 1. License Rights and Redistribution.
+
+a. <ins>Grant of Rights</ins>. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Meta’s intellectual property or other rights owned by Meta embodied in the DINO Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the DINO Materials.
+
+b. <ins>Redistribution and Use</ins>.
+
+i. Distribution of DINO Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the DINO Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such DINO Materials.
+
+ii.  If you submit for publication the results of research you perform on, using, or otherwise in connection with DINO Materials, you must acknowledge the use of DINO Materials in your publication.
+
+iii. Your use of the DINO Materials must comply with applicable laws and regulations, including Trade Control Laws and applicable privacy and data protection laws.
+
+iv. Your use of the DINO Materials will not involve or encourage others to reverse engineer, decompile or discover the underlying components of the DINO Materials.
+
+v. You are not the target of Trade Controls and your use of DINO Materials must comply with Trade Controls. You agree not to use, or permit others to use, DINO Materials for any activities subject to the International Traffic in Arms Regulations (ITAR) or end uses prohibited by Trade Controls, including those related to military or warfare purposes, nuclear industries or applications, espionage, or the development or use of guns or illegal weapons.
+
+## 2. User Support.
+
+Your use of the DINO Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use.  Meta is under no obligation to provide any support services for the DINO Materials. Any support provided is “as is”, “with all faults”, and without warranty of any kind.
+
+## 3. Disclaimer of Warranty.
+
+UNLESS REQUIRED BY APPLICABLE LAW, THE DINO MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE DINO MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE DINO MATERIALS AND ANY OUTPUT AND RESULTS.
+
+## 4. Limitation of Liability.
+
+IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+
+## 5. Intellectual Property.
+
+a. Subject to Meta’s ownership of DINO Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the DINO Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the DINO Materials, outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the DINO Materials.
+
+## 6. Term and Termination.
+
+The term of this Agreement will commence upon your acceptance of this Agreement or access to the DINO Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the DINO Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
+
+## 7. Governing Law and Jurisdiction.
+
+This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+## 8. Modifications and Amendments.
+
+Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the DINO Materials after any modification to this Agreement constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.

LICENSE_RMBG2.md

@@ -0,0 +1,26 @@
+# BRIA RMBG-2.0 License and Attribution
+
+This distribution includes a selectively quantized derivative of BRIA
+RMBG-2.0.
+
+- Original work: **BRIA RMBG-2.0**
+- Creator: **BRIA AI**
+- Model: <https://huggingface.co/briaai/RMBG-2.0>
+- Source: <https://github.com/Bria-AI/RMBG-2.0>
+- License: [Creative Commons Attribution-NonCommercial 4.0
+  International](https://creativecommons.org/licenses/by-nc/4.0/)
+
+The original model is available for non-commercial use under CC BY-NC 4.0.
+Commercial use requires a separate agreement with BRIA. Refer to the original
+model page for BRIA's current commercial licensing information.
+
+## Changes in This Distribution
+
+Eligible attention and MLP matrix weights were repacked for direct MLX affine
+8-bit execution with group size 64. Other weights remain in their source
+precision. The derivative preserves the original logical tensor names through
+checkpoint metadata and is distributed as part of the TRELLIS.2 MLX 8-bit
+runtime bundle.
+
+No endorsement by BRIA AI is implied. This notice summarizes the source,
+attribution, and changes; the linked CC BY-NC 4.0 terms remain controlling.

LICENSE_TRELLIS2

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Microsoft Corporation.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,275 @@
+---
+license: other
+license_name: trellis2-dinov3-rmbg-multiple-licenses
+license_link: https://huggingface.co/appautomaton/trellis2-mlx-8bit/blob/main/LICENSES.md
+library_name: mlx
+pipeline_tag: image-to-3d
+base_model:
+  - microsoft/TRELLIS.2-4B
+  - facebook/dinov3-vitl16-pretrain-lvd1689m
+  - briaai/RMBG-2.0
+base_model_relation: quantized
+language:
+  - en
+tags:
+  - mlx
+  - apple-silicon
+  - safetensors
+  - trellis2
+  - o-voxel
+  - image-to-3d
+  - mesh-generation
+  - pbr
+  - glb
+  - 8-bit
+  - affine-quantization
+  - mixed-precision
+  - non-commercial
+---
+
+# TRELLIS.2 MLX 8-bit Affine for `mlx-spatial`
+
+<p align="center">
+  <a href="https://appautomaton.github.io"><img alt="App Automaton project" src="https://img.shields.io/badge/App_Automaton-Project-5B5BD6?style=for-the-badge"></a>
+  <a href="https://github.com/appautomaton/mlx-spatial"><img alt="GitHub — appautomaton/mlx-spatial" src="https://img.shields.io/badge/GitHub-mlx--spatial-181717?style=for-the-badge&amp;logo=github&amp;logoColor=white"></a>
+  <a href="https://appautomaton.renocrypt.com/mlx-spatial/"><img alt="mlx-spatial documentation" src="https://img.shields.io/badge/Documentation-mlx--spatial-0A7BBB?style=for-the-badge&amp;logo=readthedocs&amp;logoColor=white"></a>
+  <a href="https://pypi.org/project/mlx-spatial/"><img alt="mlx-spatial on PyPI" src="https://img.shields.io/pypi/v/mlx-spatial?style=for-the-badge&amp;logo=pypi&amp;logoColor=white&amp;label=PyPI"></a>
+</p>
+
+<p align="center">
+  <a href="https://huggingface.co/appautomaton/trellis2-mlx-8bit"><img alt="TRELLIS.2 MLX 8-bit" src="https://img.shields.io/badge/Hugging_Face-MLX_8--bit-FF9D00?style=for-the-badge&amp;logo=huggingface&amp;logoColor=000"></a>
+</p>
+
+A self-contained TRELLIS.2 image-to-3D bundle with selective 8-bit affine
+weights for direct execution by
+[`mlx-spatial`](https://github.com/appautomaton/mlx-spatial) on Apple Silicon.
+The bundle includes the TRELLIS.2 checkpoints, DINOv3 image conditioner, and
+RMBG-2.0 foreground extractor required by the local inference pipeline.
+
+Large transformer and block-internal MLP matrices execute through MLX's packed
+quantized matrix multiplication. Convolutions and accuracy-sensitive boundary
+tensors remain in their source precision.
+
+> This is an unofficial quantized derivative. It is not a Microsoft, Meta, or
+> BRIA release. Because the bundle contains RMBG-2.0, it is restricted to
+> non-commercial use unless the user obtains separate commercial rights from
+> BRIA.
+
+## Variants
+
+| Variant | Precision | Complete bundle size | Model |
+| --- | --- | ---: | --- |
+| Source dependency set | BF16, FP16, and FP32 | 18.483 GB | TRELLIS.2-4B + DINOv3 ViT-L/16 + RMBG-2.0 |
+| MLX 8-bit bundle | Selective affine INT8 with retained BF16, FP16, and FP32 | 11.248 GB | **This model** |
+
+The complete 8-bit bundle is 39.1% smaller and does not require duplicate
+full-precision checkpoints. Both forms expose the same 5,587 logical tensors
+with identical names, shapes, and declared source dtypes.
+
+## Compatibility
+
+This format requires an `mlx-spatial` build that includes TRELLIS.2 affine
+checkpoint support. Until that support is available in a tagged PyPI release,
+install the current project revision:
+
+```bash
+pip install \
+  "mlx-spatial @ git+https://github.com/appautomaton/mlx-spatial.git@main"
+```
+
+The runtime targets Apple Silicon, Python 3.13, and MLX `0.32.x`. It does not
+use Torch, CUDA, or a dequantized full-precision checkpoint.
+
+## Use
+
+Download the complete bundle:
+
+```bash
+hf download appautomaton/trellis2-mlx-8bit \
+  --local-dir weights/trellis2-mlx-8bit
+```
+
+Validate its three runtime roots:
+
+```bash
+mlx-spatial-trellis2 validate \
+  --root weights/trellis2-mlx-8bit
+
+mlx-spatial-trellis2 dinov3-validate \
+  --root weights/trellis2-mlx-8bit/dinov3
+
+mlx-spatial-trellis2 rmbg-validate \
+  --root weights/trellis2-mlx-8bit/rmbg
+```
+
+Generate a textured GLB:
+
+```bash
+mlx-spatial-trellis2 generate-textured \
+  weights/trellis2-mlx-8bit \
+  inputs/trellis2/object.png \
+  --dino-root weights/trellis2-mlx-8bit/dinov3 \
+  --rmbg-root weights/trellis2-mlx-8bit/rmbg \
+  --output outputs/trellis2/object-8bit/model.glb \
+  --pipeline-type 1024_cascade \
+  --seed 42
+```
+
+`1024_cascade` is the recommended quality tier; use `512` when lower memory
+use or faster iteration matters more.
+
+Do not pass `--slat-steps` for a quality run; the model configuration uses 12
+steps. `--slat-steps 1` is intended only for a quick runtime smoke test.
+
+RGBA inputs use their alpha channel directly. RGB inputs are passed through
+the bundled RMBG-2.0 model. A clean single-object foreground and an uncropped
+silhouette generally produce the most useful reconstruction.
+
+## Bundle
+
+| Component | Logical tensors | INT8 matrices | Safetensors bytes |
+| --- | ---: | ---: | ---: |
+| TRELLIS.2 flow, VAE, and decoder checkpoints | 4,418 | 1,306 | 10,560,255,296 |
+| `dinov3/model.safetensors` | 415 | 144 | 344,386,453 |
+| `rmbg/model.safetensors` | 754 | 96 | 343,122,500 |
+| **Total** | **5,587** | **1,546** | **11,247,764,249** |
+
+Configuration, model card, license, and RMBG support files account for the
+small difference between checkpoint bytes and the complete directory size.
+
+The repository preserves the TRELLIS.2 checkpoint names and layout used by the
+source model. DINOv3 and RMBG-2.0 are bundled under `dinov3/` and `rmbg/` so a
+single download contains every runtime weight.
+
+## Quantization
+
+The quantization scheme is affine 8-bit with group size 64. Packed weights are
+stored as `uint32` with FP32 scales and biases, then executed directly through
+`mx.quantized_matmul`.
+
+The following block-internal two-dimensional weights are quantized:
+
+- self-attention QKV and output projections in all five TRELLIS.2 flow DiTs;
+- cross-attention query, key/value, and output projections in those DiTs;
+- DiT MLP input and output matrices;
+- internal ConvNeXt MLP matrices in the shape and texture encoders and
+  decoders;
+- DINOv3 attention Q/K/V/output and MLP up/down projections;
+- RMBG-2.0 Swin attention QKV/output and MLP FC matrices.
+
+The following tensors remain in their source precision:
+
+- all dense and sparse convolutions, including the sparse-structure decoder;
+- input, output, patch, latent, subdivision, and final prediction boundaries;
+- timestep and condition projections, embeddings, learned tokens, and
+  positional tensors;
+- normalization parameters, RMSNorm parameters, modulation tensors, and
+  biases;
+- RMBG stage-transition reductions;
+- matrices whose input dimension is incompatible with group size 64.
+
+Physical packed arrays use internal qweight, scale, and bias suffixes. The
+runtime reconstructs the original logical names from safetensors metadata, so
+the existing TRELLIS.2 pipeline and configuration continue to use the same
+checkpoint contract. Format details are embedded under
+`mlx_spatial.trellis2.quantization`.
+
+## Reproducing the Bundle
+
+Starting from the three source weight roots:
+
+```bash
+mlx-spatial-trellis2-quantize \
+  weights/trellis2 \
+  weights/trellis2-mlx-8bit \
+  --dinov3-root weights/dinov3-vitl16-pretrain-lvd1689m \
+  --rmbg-root weights/rmbg2 \
+  --bits 8 \
+  --group-size 64
+```
+
+The command creates a complete runtime root. Do not add full-precision copies
+of the quantized checkpoints to the 8-bit repository.
+
+## Verification
+
+- Logical checkpoint inspection matched all 5,587 source tensor names,
+  shapes, and dtypes exactly.
+- The quantization inventory contains 1,546 packed matrices: 1,050 DiT
+  matrices, 256 encoder/decoder MLP matrices, 144 DINOv3 matrices, and 96
+  RMBG-2.0 matrices.
+- The complete `mlx-spatial` repository suite passed 1,175 tests; 10 tests
+  were skipped, 50 were deselected, and 3 expected failures remained.
+- A 512-resolution, 12-step run completed background handling, DINOv3
+  conditioning, sparse-structure sampling, shape and texture SLat sampling,
+  both decoders, and textured GLB export with no blocker.
+- That run produced 4,022 shape tokens, 1,041,541 decoded shape/texture points,
+  and a 12,083,344-byte GLB. The rendered asset retained coherent geometry and
+  recognizable appearance.
+- The same run completed in 152.23 seconds, observed 3.516 GB peak MLX
+  allocator use, and recorded zero swap growth.
+- A separate `1024_cascade`, 12-step run completed the same end-to-end path
+  and produced a Blender-readable 12,670,448-byte GLB with 199,884 faces and
+  embedded 1024 x 1024 PBR textures.
+
+The 512 runtime and memory figures are one local Apple Silicon observation,
+not a general benchmark. The `1024_cascade` run overlapped another MLX workload,
+so it establishes compatibility and artifact health rather than performance.
+Neither run is a formal claim of visual equivalence to the source weights.
+
+## Limitations
+
+- Quantization changes the sampling trajectory. Geometry, pose, topology,
+  texture placement, material values, and unseen surfaces can differ from the
+  source-precision model even with the same seed.
+- End-to-end validation covers `512` and `1024_cascade`. The standalone `1024`
+  and `1536_cascade` routes have not received an equivalent quality evaluation.
+- Single-view reconstruction cannot determine unseen geometry with certainty.
+- Fine detail depends on foreground extraction, cropping, occlusion,
+  reflections, transparency, and thin structures.
+- Generated meshes can contain small holes, non-manifold regions, or other
+  topology requiring downstream cleanup.
+- The packed checkpoint format requires `mlx-spatial`; generic safetensors
+  readers expose the physical packed arrays rather than the logical matrices.
+- This bundle supports inference, not training or fine-tuning.
+- Commercial use of the bundled RMBG-2.0 weights requires separate rights
+  from BRIA.
+
+## License and Attribution
+
+This is a combined derivative; no single permissive license covers every
+included component. Read [`LICENSES.md`](LICENSES.md) and all referenced terms
+before downloading, using, or redistributing the bundle.
+
+- Microsoft TRELLIS.2 files are provided under the MIT License. A copy is
+  included as [`LICENSE_TRELLIS2`](LICENSE_TRELLIS2).
+- Meta DINOv3 files and derivatives are governed by the DINOv3 License. The
+  complete agreement is included as
+  [`LICENSE_DINOV3.md`](LICENSE_DINOV3.md), as required for redistribution.
+- BRIA RMBG-2.0 is provided for non-commercial use under CC BY-NC 4.0. Its
+  attribution and change notice are included in
+  [`LICENSE_RMBG2.md`](LICENSE_RMBG2.md). Commercial use requires separate
+  authorization from BRIA.
+
+The combined bundle must therefore be treated as non-commercial unless the
+user has obtained all additional rights that their use requires. This
+repository is not affiliated with or endorsed by Microsoft, Meta, BRIA, or the
+original authors.
+
+## Links
+
+- [App Automaton](https://appautomaton.github.io)
+- [AppAutomaton models on Hugging Face](https://huggingface.co/appautomaton)
+- [`appautomaton/mlx-spatial`](https://github.com/appautomaton/mlx-spatial) — MLX-native 3D and spatial inference for Apple Silicon.
+- [`mlx-spatial` documentation](https://appautomaton.renocrypt.com/mlx-spatial/)
+- [`mlx-spatial` on PyPI](https://pypi.org/project/mlx-spatial/)
+- [TRELLIS.2 guide](https://github.com/appautomaton/mlx-spatial/blob/main/docs/trellis2.md)
+- [TRELLIS.2 MLX 8-bit model](https://huggingface.co/appautomaton/trellis2-mlx-8bit)
+- [Microsoft TRELLIS.2-4B model](https://huggingface.co/microsoft/TRELLIS.2-4B)
+- [Microsoft TRELLIS.2 source](https://github.com/microsoft/TRELLIS.2)
+- [TRELLIS.2 project page](https://microsoft.github.io/TRELLIS.2/)
+- [Meta DINOv3 ViT-L/16 model](https://huggingface.co/facebook/dinov3-vitl16-pretrain-lvd1689m)
+- [Meta DINOv3 source](https://github.com/facebookresearch/dinov3)
+- [BRIA RMBG-2.0 model](https://huggingface.co/briaai/RMBG-2.0)
+- [BRIA RMBG-2.0 source](https://github.com/Bria-AI/RMBG-2.0)
+- [MLX](https://github.com/ml-explore/mlx)

ckpts/shape_dec_next_dc_f16c32_fp16.json

@@ -0,0 +1,24 @@
+{
+    "name": "FlexiDualGridVaeDecoder",
+    "args": {
+        "resolution": 256,
+        "model_channels": [1024, 512, 256, 128, 64],
+        "latent_channels": 32,
+        "num_blocks": [4, 16, 8, 4, 0],
+        "block_type": [
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d"
+        ],
+        "up_block_type": [
+            "SparseResBlockC2S3d",
+            "SparseResBlockC2S3d",
+            "SparseResBlockC2S3d",
+            "SparseResBlockC2S3d"
+        ],
+        "block_args": [{}, {}, {}, {}, {}],
+        "use_fp16": true
+    }
+}

ckpts/shape_dec_next_dc_f16c32_fp16.safetensors

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

ckpts/shape_enc_next_dc_f16c32_fp16.json

@@ -0,0 +1,23 @@
+{
+    "name": "FlexiDualGridVaeEncoder",
+    "args": {
+        "model_channels": [64, 128, 256, 512, 1024],
+        "latent_channels": 32,
+        "num_blocks": [0, 4, 8, 16, 4],
+        "block_type": [
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d"
+        ],
+        "down_block_type": [
+            "SparseResBlockS2C3d",
+            "SparseResBlockS2C3d",
+            "SparseResBlockS2C3d",
+            "SparseResBlockS2C3d"
+        ],
+        "block_args": [{}, {}, {}, {}, {}],
+        "use_fp16": true
+    }
+}

ckpts/shape_enc_next_dc_f16c32_fp16.safetensors

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

ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.json

@@ -0,0 +1,19 @@
+{
+    "name": "SLatFlowModel",
+    "args": {
+        "resolution": 64,
+        "in_channels": 32,
+        "out_channels": 32,
+        "model_channels": 1536,
+        "cond_channels": 1024,
+        "num_blocks": 30,
+        "num_heads": 12,
+        "mlp_ratio": 5.3334,
+        "pe_mode": "rope",
+        "share_mod": true,
+        "initialization": "scaled",
+        "qk_rms_norm": true,
+        "qk_rms_norm_cross": true,
+        "dtype": "bfloat16"
+    }
+}

ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16.safetensors

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

ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.json

@@ -0,0 +1,19 @@
+{
+    "name": "SLatFlowModel",
+    "args": {
+        "resolution": 32,
+        "in_channels": 32,
+        "out_channels": 32,
+        "model_channels": 1536,
+        "cond_channels": 1024,
+        "num_blocks": 30,
+        "num_heads": 12,
+        "mlp_ratio": 5.3334,
+        "pe_mode": "rope",
+        "share_mod": true,
+        "initialization": "scaled",
+        "qk_rms_norm": true,
+        "qk_rms_norm_cross": true,
+        "dtype": "bfloat16"
+    }
+}

ckpts/slat_flow_img2shape_dit_1_3B_512_bf16.safetensors

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

ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16.json

@@ -0,0 +1,19 @@
+{
+    "name": "SLatFlowModel",
+    "args": {
+        "resolution": 64,
+        "in_channels": 64,
+        "out_channels": 32,
+        "model_channels": 1536,
+        "cond_channels": 1024,
+        "num_blocks": 30,
+        "num_heads": 12,
+        "mlp_ratio": 5.3334,
+        "pe_mode": "rope",
+        "share_mod": true,
+        "initialization": "scaled",
+        "qk_rms_norm": true,
+        "qk_rms_norm_cross": true,
+        "dtype": "bfloat16"
+    }
+}

ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16.safetensors

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

ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16.json

@@ -0,0 +1,19 @@
+{
+    "name": "SLatFlowModel",
+    "args": {
+        "resolution": 32,
+        "in_channels": 64,
+        "out_channels": 32,
+        "model_channels": 1536,
+        "cond_channels": 1024,
+        "num_blocks": 30,
+        "num_heads": 12,
+        "mlp_ratio": 5.3334,
+        "pe_mode": "rope",
+        "share_mod": true,
+        "initialization": "scaled",
+        "qk_rms_norm": true,
+        "qk_rms_norm_cross": true,
+        "dtype": "bfloat16"
+    }
+}

ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16.safetensors

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

ckpts/ss_flow_img_dit_1_3B_64_bf16.json

@@ -0,0 +1,19 @@
+{
+    "name": "SparseStructureFlowModel",
+    "args": {
+        "resolution": 16,
+        "in_channels": 8,
+        "out_channels": 8,
+        "model_channels": 1536,
+        "cond_channels": 1024,
+        "num_blocks": 30,
+        "num_heads": 12,
+        "mlp_ratio": 5.3334,
+        "pe_mode": "rope",
+        "share_mod": true,
+        "initialization": "scaled",
+        "qk_rms_norm": true,
+        "qk_rms_norm_cross": true,
+        "dtype": "bfloat16"
+    }
+}

ckpts/ss_flow_img_dit_1_3B_64_bf16.safetensors

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

ckpts/tex_dec_next_dc_f16c32_fp16.json

@@ -0,0 +1,25 @@
+{
+    "name": "SparseUnetVaeDecoder",
+    "args": {
+        "out_channels": 6,
+        "model_channels": [1024, 512, 256, 128, 64],
+        "latent_channels": 32,
+        "num_blocks": [4, 16, 8, 4, 0],
+        "block_type": [
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d"
+        ],
+        "up_block_type": [
+            "SparseResBlockC2S3d",
+            "SparseResBlockC2S3d",
+            "SparseResBlockC2S3d",
+            "SparseResBlockC2S3d"
+        ],
+        "block_args": [{}, {}, {}, {}, {}],
+        "pred_subdiv": false,
+        "use_fp16": true
+    }
+}

ckpts/tex_dec_next_dc_f16c32_fp16.safetensors

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

ckpts/tex_enc_next_dc_f16c32_fp16.json

@@ -0,0 +1,24 @@
+{
+    "name": "SparseUnetVaeEncoder",
+    "args": {
+        "in_channels": 6,
+        "model_channels": [64, 128, 256, 512, 1024],
+        "latent_channels": 32,
+        "num_blocks": [0, 4, 8, 16, 4],
+        "block_type": [
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d",
+            "SparseConvNeXtBlock3d"
+        ],
+        "down_block_type": [
+            "SparseResBlockS2C3d",
+            "SparseResBlockS2C3d",
+            "SparseResBlockS2C3d",
+            "SparseResBlockS2C3d"
+        ],
+        "block_args": [{}, {}, {}, {}, {}],
+        "use_fp16": true
+    }
+}

ckpts/tex_enc_next_dc_f16c32_fp16.safetensors

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

dinov3/config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "DINOv3ViTModel"
+  ],
+  "attention_dropout": 0.0,
+  "drop_path_rate": 0.0,
+  "hidden_act": "gelu",
+  "hidden_size": 1024,
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "key_bias": false,
+  "layer_norm_eps": 1e-05,
+  "layerscale_value": 1.0,
+  "mlp_bias": true,
+  "model_type": "dinov3_vit",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "num_register_tokens": 4,
+  "patch_size": 16,
+  "pos_embed_jitter": null,
+  "pos_embed_rescale": 2.0,
+  "pos_embed_shift": null,
+  "proj_bias": true,
+  "query_bias": true,
+  "rope_theta": 100.0,
+  "torch_dtype": "float32",
+  "transformers_version": "4.56.0.dev0",
+  "use_gated_mlp": false,
+  "value_bias": true
+}

dinov3/model.safetensors

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

microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16.json

@@ -0,0 +1,12 @@
+
+{
+    "name": "SparseStructureDecoder",
+    "args": {
+        "out_channels": 1,
+        "latent_channels": 8,
+        "num_res_blocks": 2,
+        "num_res_blocks_middle": 2,
+        "channels": [512, 128, 32],
+        "use_fp16": true
+    }
+}

microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16.safetensors

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

pipeline.json

@@ -0,0 +1,95 @@
+{
+    "name": "Trellis2ImageTo3DPipeline",
+    "args": {
+        "models": {
+            "sparse_structure_decoder": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16",
+            "sparse_structure_flow_model": "ckpts/ss_flow_img_dit_1_3B_64_bf16",
+            "shape_slat_decoder": "ckpts/shape_dec_next_dc_f16c32_fp16",
+            "shape_slat_flow_model_512": "ckpts/slat_flow_img2shape_dit_1_3B_512_bf16",
+            "shape_slat_flow_model_1024": "ckpts/slat_flow_img2shape_dit_1_3B_1024_bf16",
+            "tex_slat_decoder": "ckpts/tex_dec_next_dc_f16c32_fp16",
+            "tex_slat_flow_model_512": "ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16",
+            "tex_slat_flow_model_1024": "ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16"
+        },
+        "sparse_structure_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-5
+            },
+            "params": {
+                "steps": 12,
+                "guidance_strength": 7.5,
+                "guidance_rescale": 0.7,
+                "guidance_interval": [0.6, 1.0],
+                "rescale_t": 5.0
+            }
+        },
+        "shape_slat_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-5
+            },
+            "params": {
+                "steps": 12,
+                "guidance_strength": 7.5,
+                "guidance_rescale": 0.5,
+                "guidance_interval": [0.6, 1.0],
+                "rescale_t": 3.0
+            }
+        },
+        "shape_slat_normalization": {
+            "mean": [
+                0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+            ],
+            "std": [
+                5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+            ]
+        },
+        "tex_slat_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-5
+            },
+            "params": {
+                "steps": 12,
+                "guidance_strength": 1.0,
+                "guidance_rescale": 0.0,
+                "guidance_interval": [0.6, 0.9],
+                "rescale_t": 3.0
+            }
+        },
+        "tex_slat_normalization": {
+            "mean": [
+                3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+                0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+                -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+                1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+            ],
+            "std": [
+                2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+                2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+                2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+                2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+            ]
+        },
+        "image_cond_model": {
+            "name": "DinoV3FeatureExtractor",
+            "args": {
+                "model_name": "facebook/dinov3-vitl16-pretrain-lvd1689m"
+            }
+        },
+        "rembg_model": {
+            "name": "BiRefNet",
+            "args": {
+                "model_name": "briaai/RMBG-2.0"
+            }
+        },
+        "default_pipeline_type": "1024_cascade"
+    }
+}

rmbg/BiRefNet_config.py

@@ -0,0 +1,11 @@
+from transformers import PretrainedConfig
+
+class BiRefNetConfig(PretrainedConfig):
+    model_type = "SegformerForSemanticSegmentation"
+    def __init__(
+        self,
+        bb_pretrained=False,
+        **kwargs
+    ):
+        self.bb_pretrained = bb_pretrained
+        super().__init__(**kwargs)

rmbg/birefnet.py

@@ -0,0 +1,2248 @@
+### config.py
+
+import os
+import math
+from transformers import PretrainedConfig
+
+class Config(PretrainedConfig):
+    def __init__(self) -> None:
+        super().__init__()
+        # PATH settings
+        self.sys_home_dir = os.path.expanduser('~')    # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx
+
+        # TASK settings
+        self.task = ['DIS5K', 'COD', 'HRSOD', 'DIS5K+HRSOD+HRS10K', 'P3M-10k'][0]
+        self.training_set = {
+            'DIS5K': ['DIS-TR', 'DIS-TR+DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4'][0],
+            'COD': 'TR-COD10K+TR-CAMO',
+            'HRSOD': ['TR-DUTS', 'TR-HRSOD', 'TR-UHRSD', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][5],
+            'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD',     # leave DIS-VD for evaluation.
+            'P3M-10k': 'TR-P3M-10k',
+        }[self.task]
+        self.prompt4loc = ['dense', 'sparse'][0]
+
+        # Faster-Training settings
+        self.load_all = True
+        self.compile = True     # 1. Trigger CPU memory leak in some extend, which is an inherent problem of PyTorch.
+                                #   Machines with > 70GB CPU memory can run the whole training on DIS5K with default setting.
+                                # 2. Higher PyTorch version may fix it: https://github.com/pytorch/pytorch/issues/119607.
+                                # 3. But compile in Pytorch > 2.0.1 seems to bring no acceleration for training.
+        self.precisionHigh = True
+
+        # MODEL settings
+        self.ms_supervision = True
+        self.out_ref = self.ms_supervision and True
+        self.dec_ipt = True
+        self.dec_ipt_split = True
+        self.cxt_num = [0, 3][1]    # multi-scale skip connections from encoder
+        self.mul_scl_ipt = ['', 'add', 'cat'][2]
+        self.dec_att = ['', 'ASPP', 'ASPPDeformable'][2]
+        self.squeeze_block = ['', 'BasicDecBlk_x1', 'ResBlk_x4', 'ASPP_x3', 'ASPPDeformable_x3'][1]
+        self.dec_blk = ['BasicDecBlk', 'ResBlk', 'HierarAttDecBlk'][0]
+
+        # TRAINING settings
+        self.batch_size = 4
+        self.IoU_finetune_last_epochs = [
+            0,
+            {
+                'DIS5K': -50,
+                'COD': -20,
+                'HRSOD': -20,
+                'DIS5K+HRSOD+HRS10K': -20,
+                'P3M-10k': -20,
+            }[self.task]
+        ][1]    # choose 0 to skip
+        self.lr = (1e-4 if 'DIS5K' in self.task else 1e-5) * math.sqrt(self.batch_size / 4)     # DIS needs high lr to converge faster. Adapt the lr linearly
+        self.size = 1024
+        self.num_workers = max(4, self.batch_size)          # will be decrease to min(it, batch_size) at the initialization of the data_loader
+
+        # Backbone settings
+        self.bb = [
+            'vgg16', 'vgg16bn', 'resnet50',         # 0, 1, 2
+            'swin_v1_t', 'swin_v1_s',               # 3, 4
+            'swin_v1_b', 'swin_v1_l',               # 5-bs9, 6-bs4
+            'pvt_v2_b0', 'pvt_v2_b1',               # 7, 8
+            'pvt_v2_b2', 'pvt_v2_b5',               # 9-bs10, 10-bs5
+        ][6]
+        self.lateral_channels_in_collection = {
+            'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
+            'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
+            'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
+            'swin_v1_t': [768, 384, 192, 96], 'swin_v1_s': [768, 384, 192, 96],
+            'pvt_v2_b0': [256, 160, 64, 32], 'pvt_v2_b1': [512, 320, 128, 64],
+        }[self.bb]
+        if self.mul_scl_ipt == 'cat':
+            self.lateral_channels_in_collection = [channel * 2 for channel in self.lateral_channels_in_collection]
+        self.cxt = self.lateral_channels_in_collection[1:][::-1][-self.cxt_num:] if self.cxt_num else []
+
+        # MODEL settings - inactive
+        self.lat_blk = ['BasicLatBlk'][0]
+        self.dec_channels_inter = ['fixed', 'adap'][0]
+        self.refine = ['', 'itself', 'RefUNet', 'Refiner', 'RefinerPVTInChannels4'][0]
+        self.progressive_ref = self.refine and True
+        self.ender = self.progressive_ref and False
+        self.scale = self.progressive_ref and 2
+        self.auxiliary_classification = False       # Only for DIS5K, where class labels are saved in `dataset.py`.
+        self.refine_iteration = 1
+        self.freeze_bb = False
+        self.model = [
+            'BiRefNet',
+        ][0]
+        if self.dec_blk == 'HierarAttDecBlk':
+            self.batch_size = 2 ** [0, 1, 2, 3, 4][2]
+
+        # TRAINING settings - inactive
+        self.preproc_methods = ['flip', 'enhance', 'rotate', 'pepper', 'crop'][:4]
+        self.optimizer = ['Adam', 'AdamW'][1]
+        self.lr_decay_epochs = [1e5]    # Set to negative N to decay the lr in the last N-th epoch.
+        self.lr_decay_rate = 0.5
+        # Loss
+        self.lambdas_pix_last = {
+            # not 0 means opening this loss
+            # original rate -- 1 : 30 : 1.5 : 0.2, bce x 30
+            'bce': 30 * 1,          # high performance
+            'iou': 0.5 * 1,         # 0 / 255
+            'iou_patch': 0.5 * 0,   # 0 / 255, win_size = (64, 64)
+            'mse': 150 * 0,         # can smooth the saliency map
+            'triplet': 3 * 0,
+            'reg': 100 * 0,
+            'ssim': 10 * 1,          # help contours,
+            'cnt': 5 * 0,          # help contours
+            'structure': 5 * 0,    # structure loss from codes of MVANet. A little improvement on DIS-TE[1,2,3], a bit more decrease on DIS-TE4.
+        }
+        self.lambdas_cls = {
+            'ce': 5.0
+        }
+        # Adv
+        self.lambda_adv_g = 10. * 0        # turn to 0 to avoid adv training
+        self.lambda_adv_d = 3. * (self.lambda_adv_g > 0)
+
+        # PATH settings - inactive
+        self.data_root_dir = os.path.join(self.sys_home_dir, 'datasets/dis')
+        self.weights_root_dir = os.path.join(self.sys_home_dir, 'weights')
+        self.weights = {
+            'pvt_v2_b2': os.path.join(self.weights_root_dir, 'pvt_v2_b2.pth'),
+            'pvt_v2_b5': os.path.join(self.weights_root_dir, ['pvt_v2_b5.pth', 'pvt_v2_b5_22k.pth'][0]),
+            'swin_v1_b': os.path.join(self.weights_root_dir, ['swin_base_patch4_window12_384_22kto1k.pth', 'swin_base_patch4_window12_384_22k.pth'][0]),
+            'swin_v1_l': os.path.join(self.weights_root_dir, ['swin_large_patch4_window12_384_22kto1k.pth', 'swin_large_patch4_window12_384_22k.pth'][0]),
+            'swin_v1_t': os.path.join(self.weights_root_dir, ['swin_tiny_patch4_window7_224_22kto1k_finetune.pth'][0]),
+            'swin_v1_s': os.path.join(self.weights_root_dir, ['swin_small_patch4_window7_224_22kto1k_finetune.pth'][0]),
+            'pvt_v2_b0': os.path.join(self.weights_root_dir, ['pvt_v2_b0.pth'][0]),
+            'pvt_v2_b1': os.path.join(self.weights_root_dir, ['pvt_v2_b1.pth'][0]),
+        }
+
+        # Callbacks - inactive
+        self.verbose_eval = True
+        self.only_S_MAE = False
+        self.use_fp16 = False   # Bugs. It may cause nan in training.
+        self.SDPA_enabled = False    # Bugs. Slower and errors occur in multi-GPUs
+
+        # others
+        self.device = [0, 'cpu'][0]     # .to(0) == .to('cuda:0')
+
+        self.batch_size_valid = 1
+        self.rand_seed = 7
+        # run_sh_file = [f for f in os.listdir('.') if 'train.sh' == f] + [os.path.join('..', f) for f in os.listdir('..') if 'train.sh' == f]
+        # with open(run_sh_file[0], 'r') as f:
+        #     lines = f.readlines()
+        #     self.save_last = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'val_last=' in l][0].split('val_last=')[-1].split()[0])
+        #     self.save_step = int([l.strip() for l in lines if '"{}")'.format(self.task) in l and 'step=' in l][0].split('step=')[-1].split()[0])
+        # self.val_step = [0, self.save_step][0]
+
+    def print_task(self) -> None:
+        # Return task for choosing settings in shell scripts.
+        print(self.task)
+
+
+
+### models/backbones/pvt_v2.py
+
+import torch
+import torch.nn as nn
+from functools import partial
+
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+
+import math
+
+# from config import Config
+
+# config = Config()
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=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.dwconv = DWConv(hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x, H, W):
+        x = self.fc1(x)
+        x = self.dwconv(x, H, W)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0., sr_ratio=1):
+        super().__init__()
+        assert dim % num_heads == 0, f"dim {dim} should be divided by num_heads {num_heads}."
+
+        self.dim = dim
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
+        self.attn_drop_prob = attn_drop
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.sr_ratio = sr_ratio
+        if sr_ratio > 1:
+            self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
+            self.norm = nn.LayerNorm(dim)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x, H, W):
+        B, N, C = x.shape
+        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+
+        if self.sr_ratio > 1:
+            x_ = x.permute(0, 2, 1).reshape(B, C, H, W)
+            x_ = self.sr(x_).reshape(B, C, -1).permute(0, 2, 1)
+            x_ = self.norm(x_)
+            kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        else:
+            kv = self.kv(x).reshape(B, -1, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+
+        if config.SDPA_enabled:
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False
+            ).transpose(1, 2).reshape(B, N, C)
+        else:
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, sr_ratio=1):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, sr_ratio=sr_ratio)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x, H, W):
+        x = x + self.drop_path(self.attn(self.norm1(x), H, W))
+        x = x + self.drop_path(self.mlp(self.norm2(x), H, W))
+
+        return x
+
+
+class OverlapPatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=7, stride=4, in_channels=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
+        self.num_patches = self.H * self.W
+        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=stride,
+                              padding=(patch_size[0] // 2, patch_size[1] // 2))
+        self.norm = nn.LayerNorm(embed_dim)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def forward(self, x):
+        x = self.proj(x)
+        _, _, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+
+        return x, H, W
+
+
+class PyramidVisionTransformerImpr(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=1000, embed_dims=[64, 128, 256, 512],
+                 num_heads=[1, 2, 4, 8], mlp_ratios=[4, 4, 4, 4], qkv_bias=False, qk_scale=None, drop_rate=0.,
+                 attn_drop_rate=0., drop_path_rate=0., norm_layer=nn.LayerNorm,
+                 depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1]):
+        super().__init__()
+        self.num_classes = num_classes
+        self.depths = depths
+
+        # patch_embed
+        self.patch_embed1 = OverlapPatchEmbed(img_size=img_size, patch_size=7, stride=4, in_channels=in_channels,
+                                              embed_dim=embed_dims[0])
+        self.patch_embed2 = OverlapPatchEmbed(img_size=img_size // 4, patch_size=3, stride=2, in_channels=embed_dims[0],
+                                              embed_dim=embed_dims[1])
+        self.patch_embed3 = OverlapPatchEmbed(img_size=img_size // 8, patch_size=3, stride=2, in_channels=embed_dims[1],
+                                              embed_dim=embed_dims[2])
+        self.patch_embed4 = OverlapPatchEmbed(img_size=img_size // 16, patch_size=3, stride=2, in_channels=embed_dims[2],
+                                              embed_dim=embed_dims[3])
+
+        # transformer encoder
+        dpr = np.linspace(0, drop_path_rate, sum(depths)).tolist() # stochastic depth decay rule
+        cur = 0
+        self.block1 = nn.ModuleList([Block(
+            dim=embed_dims[0], num_heads=num_heads[0], mlp_ratio=mlp_ratios[0], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[0])
+            for i in range(depths[0])])
+        self.norm1 = norm_layer(embed_dims[0])
+
+        cur += depths[0]
+        self.block2 = nn.ModuleList([Block(
+            dim=embed_dims[1], num_heads=num_heads[1], mlp_ratio=mlp_ratios[1], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[1])
+            for i in range(depths[1])])
+        self.norm2 = norm_layer(embed_dims[1])
+
+        cur += depths[1]
+        self.block3 = nn.ModuleList([Block(
+            dim=embed_dims[2], num_heads=num_heads[2], mlp_ratio=mlp_ratios[2], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[2])
+            for i in range(depths[2])])
+        self.norm3 = norm_layer(embed_dims[2])
+
+        cur += depths[2]
+        self.block4 = nn.ModuleList([Block(
+            dim=embed_dims[3], num_heads=num_heads[3], mlp_ratio=mlp_ratios[3], qkv_bias=qkv_bias, qk_scale=qk_scale,
+            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[cur + i], norm_layer=norm_layer,
+            sr_ratio=sr_ratios[3])
+            for i in range(depths[3])])
+        self.norm4 = norm_layer(embed_dims[3])
+
+        # classification head
+        # self.head = nn.Linear(embed_dims[3], num_classes) if num_classes > 0 else nn.Identity()
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+        elif isinstance(m, nn.Conv2d):
+            fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            fan_out //= m.groups
+            m.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def init_weights(self, pretrained=None):
+        if isinstance(pretrained, str):
+            logger = 1
+            #load_checkpoint(self, pretrained, map_location='cpu', strict=False, logger=logger)
+
+    def reset_drop_path(self, drop_path_rate):
+        dpr = np.linspace(0, drop_path_rate, sum(self.depths)).tolist()
+        cur = 0
+        for i in range(self.depths[0]):
+            self.block1[i].drop_path.drop_prob = dpr[cur + i]
+
+        cur += self.depths[0]
+        for i in range(self.depths[1]):
+            self.block2[i].drop_path.drop_prob = dpr[cur + i]
+
+        cur += self.depths[1]
+        for i in range(self.depths[2]):
+            self.block3[i].drop_path.drop_prob = dpr[cur + i]
+
+        cur += self.depths[2]
+        for i in range(self.depths[3]):
+            self.block4[i].drop_path.drop_prob = dpr[cur + i]
+
+    def freeze_patch_emb(self):
+        self.patch_embed1.requires_grad = False
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed1', 'pos_embed2', 'pos_embed3', 'pos_embed4', 'cls_token'}  # has pos_embed may be better
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        outs = []
+
+        # stage 1
+        x, H, W = self.patch_embed1(x)
+        for i, blk in enumerate(self.block1):
+            x = blk(x, H, W)
+        x = self.norm1(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+
+        # stage 2
+        x, H, W = self.patch_embed2(x)
+        for i, blk in enumerate(self.block2):
+            x = blk(x, H, W)
+        x = self.norm2(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+
+        # stage 3
+        x, H, W = self.patch_embed3(x)
+        for i, blk in enumerate(self.block3):
+            x = blk(x, H, W)
+        x = self.norm3(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+
+        # stage 4
+        x, H, W = self.patch_embed4(x)
+        for i, blk in enumerate(self.block4):
+            x = blk(x, H, W)
+        x = self.norm4(x)
+        x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+
+        return outs
+
+        # return x.mean(dim=1)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        # x = self.head(x)
+
+        return x
+
+
+class DWConv(nn.Module):
+    def __init__(self, dim=768):
+        super(DWConv, self).__init__()
+        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
+
+    def forward(self, x, H, W):
+        B, N, C = x.shape
+        x = x.transpose(1, 2).view(B, C, H, W).contiguous()
+        x = self.dwconv(x)
+        x = x.flatten(2).transpose(1, 2)
+
+        return x
+
+
+def _conv_filter(state_dict, patch_size=16):
+    """ convert patch embedding weight from manual patchify + linear proj to conv"""
+    out_dict = {}
+    for k, v in state_dict.items():
+        if 'patch_embed.proj.weight' in k:
+            v = v.reshape((v.shape[0], 3, patch_size, patch_size))
+        out_dict[k] = v
+
+    return out_dict
+
+
+## @register_model
+class pvt_v2_b0(PyramidVisionTransformerImpr):
+    def __init__(self, **kwargs):
+        super(pvt_v2_b0, self).__init__(
+            patch_size=4, embed_dims=[32, 64, 160, 256], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+
+
+
+## @register_model
+class pvt_v2_b1(PyramidVisionTransformerImpr):
+    def __init__(self, **kwargs):
+        super(pvt_v2_b1, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[2, 2, 2, 2], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+
+## @register_model
+class pvt_v2_b2(PyramidVisionTransformerImpr):
+    def __init__(self, in_channels=3, **kwargs):
+        super(pvt_v2_b2, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 6, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1, in_channels=in_channels)
+
+## @register_model
+class pvt_v2_b3(PyramidVisionTransformerImpr):
+    def __init__(self, **kwargs):
+        super(pvt_v2_b3, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 4, 18, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+
+## @register_model
+class pvt_v2_b4(PyramidVisionTransformerImpr):
+    def __init__(self, **kwargs):
+        super(pvt_v2_b4, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[8, 8, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 8, 27, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+
+
+## @register_model
+class pvt_v2_b5(PyramidVisionTransformerImpr):
+    def __init__(self, **kwargs):
+        super(pvt_v2_b5, self).__init__(
+            patch_size=4, embed_dims=[64, 128, 320, 512], num_heads=[1, 2, 5, 8], mlp_ratios=[4, 4, 4, 4],
+            qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), depths=[3, 6, 40, 3], sr_ratios=[8, 4, 2, 1],
+            drop_rate=0.0, drop_path_rate=0.1)
+
+
+
+### models/backbones/swin_v1.py
+
+# --------------------------------------------------------
+# Swin Transformer
+# Copyright (c) 2021 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Ze Liu, Yutong Lin, Yixuan Wei
+# --------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+import numpy as np
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+
+# from config import Config
+
+
+# config = Config()
+
+class Mlp(nn.Module):
+    """ Multilayer perceptron."""
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=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.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    """ Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
+
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop_prob = attn_drop
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """ Forward function.
+
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+
+        if config.SDPA_enabled:
+            x = torch.nn.functional.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=None, dropout_p=self.attn_drop_prob, is_causal=False
+            ).transpose(1, 2).reshape(B_, N, C)
+        else:
+            attn = (q @ k.transpose(-2, -1))
+
+            relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)  # Wh*Ww,Wh*Ww,nH
+            relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+            if mask is not None:
+                nW = mask.shape[0]
+                attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+                attn = attn.view(-1, self.num_heads, N, N)
+                attn = self.softmax(attn)
+            else:
+                attn = self.softmax(attn)
+
+            attn = self.attn_drop(attn)
+
+            x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwinTransformerBlock(nn.Module):
+    """ Swin Transformer Block.
+
+    Args:
+        dim (int): Number of input channels.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, dim, num_heads, window_size=7, shift_size=0,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,
+                 act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        self.H = None
+        self.W = None
+
+    def forward(self, x, mask_matrix):
+        """ Forward function.
+
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+            mask_matrix: Attention mask for cyclic shift.
+        """
+        B, L, C = x.shape
+        H, W = self.H, self.W
+        assert L == H * W, "input feature has wrong size"
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # pad feature maps to multiples of window size
+        pad_l = pad_t = 0
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+            attn_mask = mask_matrix
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(-1, self.window_size * self.window_size, C)  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows = self.attn(x_windows, mask=attn_mask)  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+        else:
+            x = shifted_x
+
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x
+
+
+class PatchMerging(nn.Module):
+    """ Patch Merging Layer
+
+    Args:
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+    def __init__(self, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x, H, W):
+        """ Forward function.
+
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+
+        # padding
+        pad_input = (H % 2 == 1) or (W % 2 == 1)
+        if pad_input:
+            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+
+class BasicLayer(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+
+    Args:
+        dim (int): Number of feature channels
+        depth (int): Depths of this stage.
+        num_heads (int): Number of attention head.
+        window_size (int): Local window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self,
+                 dim,
+                 depth,
+                 num_heads,
+                 window_size=7,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop=0.,
+                 attn_drop=0.,
+                 drop_path=0.,
+                 norm_layer=nn.LayerNorm,
+                 downsample=None,
+                 use_checkpoint=False):
+        super().__init__()
+        self.window_size = window_size
+        self.shift_size = window_size // 2
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(
+                dim=dim,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift_size=0 if (i % 2 == 0) else window_size // 2,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                norm_layer=norm_layer)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x, H, W):
+        """ Forward function.
+
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+
+        # calculate attention mask for SW-MSA
+        Hp = int(np.ceil(H / self.window_size)) * self.window_size
+        Wp = int(np.ceil(W / self.window_size)) * self.window_size
+        img_mask = torch.zeros((1, Hp, Wp, 1), device=x.device)  # 1 Hp Wp 1
+        h_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        w_slices = (slice(0, -self.window_size),
+                    slice(-self.window_size, -self.shift_size),
+                    slice(-self.shift_size, None))
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)).to(x.dtype)
+
+        for blk in self.blocks:
+            blk.H, blk.W = H, W
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x, attn_mask)
+            else:
+                x = blk(x, attn_mask)
+        if self.downsample is not None:
+            x_down = self.downsample(x, H, W)
+            Wh, Ww = (H + 1) // 2, (W + 1) // 2
+            return x, H, W, x_down, Wh, Ww
+        else:
+            return x, H, W, x, H, W
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+
+    Args:
+        patch_size (int): Patch token size. Default: 4.
+        in_channels (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        norm_layer (nn.Module, optional): Normalization layer. Default: None
+    """
+
+    def __init__(self, patch_size=4, in_channels=3, embed_dim=96, norm_layer=None):
+        super().__init__()
+        patch_size = to_2tuple(patch_size)
+        self.patch_size = patch_size
+
+        self.in_channels = in_channels
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        """Forward function."""
+        # padding
+        _, _, H, W = x.size()
+        if W % self.patch_size[1] != 0:
+            x = F.pad(x, (0, self.patch_size[1] - W % self.patch_size[1]))
+        if H % self.patch_size[0] != 0:
+            x = F.pad(x, (0, 0, 0, self.patch_size[0] - H % self.patch_size[0]))
+
+        x = self.proj(x)  # B C Wh Ww
+        if self.norm is not None:
+            Wh, Ww = x.size(2), x.size(3)
+            x = x.flatten(2).transpose(1, 2)
+            x = self.norm(x)
+            x = x.transpose(1, 2).view(-1, self.embed_dim, Wh, Ww)
+
+        return x
+
+
+class SwinTransformer(nn.Module):
+    """ Swin Transformer backbone.
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+
+    Args:
+        pretrain_img_size (int): Input image size for training the pretrained model,
+            used in absolute postion embedding. Default 224.
+        patch_size (int | tuple(int)): Patch size. Default: 4.
+        in_channels (int): Number of input image channels. Default: 3.
+        embed_dim (int): Number of linear projection output channels. Default: 96.
+        depths (tuple[int]): Depths of each Swin Transformer stage.
+        num_heads (tuple[int]): Number of attention head of each stage.
+        window_size (int): Window size. Default: 7.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
+        drop_rate (float): Dropout rate.
+        attn_drop_rate (float): Attention dropout rate. Default: 0.
+        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
+        out_indices (Sequence[int]): Output from which stages.
+        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
+            -1 means not freezing any parameters.
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self,
+                 pretrain_img_size=224,
+                 patch_size=4,
+                 in_channels=3,
+                 embed_dim=96,
+                 depths=[2, 2, 6, 2],
+                 num_heads=[3, 6, 12, 24],
+                 window_size=7,
+                 mlp_ratio=4.,
+                 qkv_bias=True,
+                 qk_scale=None,
+                 drop_rate=0.,
+                 attn_drop_rate=0.,
+                 drop_path_rate=0.2,
+                 norm_layer=nn.LayerNorm,
+                 ape=False,
+                 patch_norm=True,
+                 out_indices=(0, 1, 2, 3),
+                 frozen_stages=-1,
+                 use_checkpoint=False):
+        super().__init__()
+
+        self.pretrain_img_size = pretrain_img_size
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.out_indices = out_indices
+        self.frozen_stages = frozen_stages
+
+        # split image into non-overlapping patches
+        self.patch_embed = PatchEmbed(
+            patch_size=patch_size, in_channels=in_channels, embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None)
+
+        # absolute position embedding
+        if self.ape:
+            pretrain_img_size = to_2tuple(pretrain_img_size)
+            patch_size = to_2tuple(patch_size)
+            patches_resolution = [pretrain_img_size[0] // patch_size[0], pretrain_img_size[1] // patch_size[1]]
+
+            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, patches_resolution[0], patches_resolution[1]))
+            trunc_normal_(self.absolute_pos_embed, std=.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth
+        dpr = np.linspace(0, drop_path_rate, sum(depths)).tolist() # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(
+                dim=int(embed_dim * 2 ** i_layer),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
+                norm_layer=norm_layer,
+                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+                use_checkpoint=use_checkpoint)
+            self.layers.append(layer)
+
+        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
+        self.num_features = num_features
+
+        # add a norm layer for each output
+        for i_layer in out_indices:
+            layer = norm_layer(num_features[i_layer])
+            layer_name = f'norm{i_layer}'
+            self.add_module(layer_name, layer)
+
+        self._freeze_stages()
+
+    def _freeze_stages(self):
+        if self.frozen_stages >= 0:
+            self.patch_embed.eval()
+            for param in self.patch_embed.parameters():
+                param.requires_grad = False
+
+        if self.frozen_stages >= 1 and self.ape:
+            self.absolute_pos_embed.requires_grad = False
+
+        if self.frozen_stages >= 2:
+            self.pos_drop.eval()
+            for i in range(0, self.frozen_stages - 1):
+                m = self.layers[i]
+                m.eval()
+                for param in m.parameters():
+                    param.requires_grad = False
+
+
+    def forward(self, x):
+        """Forward function."""
+        x = self.patch_embed(x)
+
+        Wh, Ww = x.size(2), x.size(3)
+        if self.ape:
+            # interpolate the position embedding to the corresponding size
+            absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wh, Ww), mode='bicubic')
+            x = (x + absolute_pos_embed) # B Wh*Ww C
+            
+        outs = []#x.contiguous()]
+        x = x.flatten(2).transpose(1, 2)
+        x = self.pos_drop(x)
+        for i in range(self.num_layers):
+            layer = self.layers[i]
+            x_out, H, W, x, Wh, Ww = layer(x, Wh, Ww)
+
+            if i in self.out_indices:
+                norm_layer = getattr(self, f'norm{i}')
+                x_out = norm_layer(x_out)
+
+                out = x_out.view(-1, H, W, self.num_features[i]).permute(0, 3, 1, 2).contiguous()
+                outs.append(out)
+
+        return tuple(outs)
+
+    def train(self, mode=True):
+        """Convert the model into training mode while keep layers freezed."""
+        super(SwinTransformer, self).train(mode)
+        self._freeze_stages()
+
+def swin_v1_t():
+    model = SwinTransformer(embed_dim=96, depths=[2, 2, 6, 2], num_heads=[3, 6, 12, 24], window_size=7)
+    return model
+
+def swin_v1_s():
+    model = SwinTransformer(embed_dim=96, depths=[2, 2, 18, 2], num_heads=[3, 6, 12, 24], window_size=7)
+    return model
+
+def swin_v1_b():
+    model = SwinTransformer(embed_dim=128, depths=[2, 2, 18, 2], num_heads=[4, 8, 16, 32], window_size=12)
+    return model
+
+def swin_v1_l():
+    model = SwinTransformer(embed_dim=192, depths=[2, 2, 18, 2], num_heads=[6, 12, 24, 48], window_size=12)
+    return model
+
+
+
+### models/modules/deform_conv.py
+
+import torch
+import torch.nn as nn
+from torchvision.ops import deform_conv2d
+
+
+class DeformableConv2d(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 stride=1,
+                 padding=1,
+                 bias=False):
+
+        super(DeformableConv2d, self).__init__()
+        
+        assert type(kernel_size) == tuple or type(kernel_size) == int
+
+        kernel_size = kernel_size if type(kernel_size) == tuple else (kernel_size, kernel_size)
+        self.stride = stride if type(stride) == tuple else (stride, stride)
+        self.padding = padding
+        
+        self.offset_conv = nn.Conv2d(in_channels,
+                                     2 * kernel_size[0] * kernel_size[1],
+                                     kernel_size=kernel_size,
+                                     stride=stride,
+                                     padding=self.padding,
+                                     bias=True)
+
+        nn.init.constant_(self.offset_conv.weight, 0.)
+        nn.init.constant_(self.offset_conv.bias, 0.)
+        
+        self.modulator_conv = nn.Conv2d(in_channels,
+                                     1 * kernel_size[0] * kernel_size[1],
+                                     kernel_size=kernel_size,
+                                     stride=stride,
+                                     padding=self.padding,
+                                     bias=True)
+
+        nn.init.constant_(self.modulator_conv.weight, 0.)
+        nn.init.constant_(self.modulator_conv.bias, 0.)
+
+        self.regular_conv = nn.Conv2d(in_channels,
+                                      out_channels=out_channels,
+                                      kernel_size=kernel_size,
+                                      stride=stride,
+                                      padding=self.padding,
+                                      bias=bias)
+
+    def forward(self, x):
+        #h, w = x.shape[2:]
+        #max_offset = max(h, w)/4.
+
+        offset = self.offset_conv(x)#.clamp(-max_offset, max_offset)
+        modulator = 2. * torch.sigmoid(self.modulator_conv(x))
+        
+        x = deform_conv2d(
+            input=x,
+            offset=offset,
+            weight=self.regular_conv.weight,
+            bias=self.regular_conv.bias,
+            padding=self.padding,
+            mask=modulator,
+            stride=self.stride,
+        )
+        return x
+
+
+
+
+### utils.py
+
+import torch.nn as nn
+
+
+def build_act_layer(act_layer):
+    if act_layer == 'ReLU':
+        return nn.ReLU(inplace=True)
+    elif act_layer == 'SiLU':
+        return nn.SiLU(inplace=True)
+    elif act_layer == 'GELU':
+        return nn.GELU()
+
+    raise NotImplementedError(f'build_act_layer does not support {act_layer}')
+
+
+def build_norm_layer(dim,
+                     norm_layer,
+                     in_format='channels_last',
+                     out_format='channels_last',
+                     eps=1e-6):
+    layers = []
+    if norm_layer == 'BN':
+        if in_format == 'channels_last':
+            layers.append(to_channels_first())
+        layers.append(nn.BatchNorm2d(dim))
+        if out_format == 'channels_last':
+            layers.append(to_channels_last())
+    elif norm_layer == 'LN':
+        if in_format == 'channels_first':
+            layers.append(to_channels_last())
+        layers.append(nn.LayerNorm(dim, eps=eps))
+        if out_format == 'channels_first':
+            layers.append(to_channels_first())
+    else:
+        raise NotImplementedError(
+            f'build_norm_layer does not support {norm_layer}')
+    return nn.Sequential(*layers)
+
+
+class to_channels_first(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 3, 1, 2)
+
+
+class to_channels_last(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.permute(0, 2, 3, 1)
+
+
+
+### dataset.py
+
+_class_labels_TR_sorted = (
+    'Airplane, Ant, Antenna, Archery, Axe, BabyCarriage, Bag, BalanceBeam, Balcony, Balloon, Basket, BasketballHoop, Beatle, Bed, Bee, Bench, Bicycle, '
+    'BicycleFrame, BicycleStand, Boat, Bonsai, BoomLift, Bridge, BunkBed, Butterfly, Button, Cable, CableLift, Cage, Camcorder, Cannon, Canoe, Car, '
+    'CarParkDropArm, Carriage, Cart, Caterpillar, CeilingLamp, Centipede, Chair, Clip, Clock, Clothes, CoatHanger, Comb, ConcretePumpTruck, Crack, Crane, '
+    'Cup, DentalChair, Desk, DeskChair, Diagram, DishRack, DoorHandle, Dragonfish, Dragonfly, Drum, Earphone, Easel, ElectricIron, Excavator, Eyeglasses, '
+    'Fan, Fence, Fencing, FerrisWheel, FireExtinguisher, Fishing, Flag, FloorLamp, Forklift, GasStation, Gate, Gear, Goal, Golf, GymEquipment, Hammock, '
+    'Handcart, Handcraft, Handrail, HangGlider, Harp, Harvester, Headset, Helicopter, Helmet, Hook, HorizontalBar, Hydrovalve, IroningTable, Jewelry, Key, '
+    'KidsPlayground, Kitchenware, Kite, Knife, Ladder, LaundryRack, Lightning, Lobster, Locust, Machine, MachineGun, MagazineRack, Mantis, Medal, MemorialArchway, '
+    'Microphone, Missile, MobileHolder, Monitor, Mosquito, Motorcycle, MovingTrolley, Mower, MusicPlayer, MusicStand, ObservationTower, Octopus, OilWell, '
+    'OlympicLogo, OperatingTable, OutdoorFitnessEquipment, Parachute, Pavilion, Piano, Pipe, PlowHarrow, PoleVault, Punchbag, Rack, Racket, Rifle, Ring, Robot, '
+    'RockClimbing, Rope, Sailboat, Satellite, Scaffold, Scale, Scissor, Scooter, Sculpture, Seadragon, Seahorse, Seal, SewingMachine, Ship, Shoe, ShoppingCart, '
+    'ShoppingTrolley, Shower, Shrimp, Signboard, Skateboarding, Skeleton, Skiing, Spade, SpeedBoat, Spider, Spoon, Stair, Stand, Stationary, SteeringWheel, '
+    'Stethoscope, Stool, Stove, StreetLamp, SweetStand, Swing, Sword, TV, Table, TableChair, TableLamp, TableTennis, Tank, Tapeline, Teapot, Telescope, Tent, '
+    'TobaccoPipe, Toy, Tractor, TrafficLight, TrafficSign, Trampoline, TransmissionTower, Tree, Tricycle, TrimmerCover, Tripod, Trombone, Truck, Trumpet, Tuba, '
+    'UAV, Umbrella, UnevenBars, UtilityPole, VacuumCleaner, Violin, Wakesurfing, Watch, WaterTower, WateringPot, Well, WellLid, Wheel, Wheelchair, WindTurbine, Windmill, WineGlass, WireWhisk, Yacht'
+)
+class_labels_TR_sorted = _class_labels_TR_sorted.split(', ')
+
+
+### models/backbones/build_backbones.py
+
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+from torchvision.models import vgg16, vgg16_bn, VGG16_Weights, VGG16_BN_Weights, resnet50, ResNet50_Weights
+# from models.pvt_v2 import pvt_v2_b0, pvt_v2_b1, pvt_v2_b2, pvt_v2_b5
+# from models.swin_v1 import swin_v1_t, swin_v1_s, swin_v1_b, swin_v1_l
+# from config import Config
+
+
+config = Config()
+
+def build_backbone(bb_name, pretrained=True, params_settings=''):
+    if bb_name == 'vgg16':
+        bb_net = list(vgg16(pretrained=VGG16_Weights.DEFAULT if pretrained else None).children())[0]
+        bb = nn.Sequential(OrderedDict({'conv1': bb_net[:4], 'conv2': bb_net[4:9], 'conv3': bb_net[9:16], 'conv4': bb_net[16:23]}))
+    elif bb_name == 'vgg16bn':
+        bb_net = list(vgg16_bn(pretrained=VGG16_BN_Weights.DEFAULT if pretrained else None).children())[0]
+        bb = nn.Sequential(OrderedDict({'conv1': bb_net[:6], 'conv2': bb_net[6:13], 'conv3': bb_net[13:23], 'conv4': bb_net[23:33]}))
+    elif bb_name == 'resnet50':
+        bb_net = list(resnet50(pretrained=ResNet50_Weights.DEFAULT if pretrained else None).children())
+        bb = nn.Sequential(OrderedDict({'conv1': nn.Sequential(*bb_net[0:3]), 'conv2': bb_net[4], 'conv3': bb_net[5], 'conv4': bb_net[6]}))
+    else:
+        bb = eval('{}({})'.format(bb_name, params_settings))
+        if pretrained:
+            bb = load_weights(bb, bb_name)
+    return bb
+
+def load_weights(model, model_name):
+    save_model = torch.load(config.weights[model_name], map_location='cpu')
+    model_dict = model.state_dict()
+    state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model.items() if k in model_dict.keys()}
+    # to ignore the weights with mismatched size when I modify the backbone itself.
+    if not state_dict:
+        save_model_keys = list(save_model.keys())
+        sub_item = save_model_keys[0] if len(save_model_keys) == 1 else None
+        state_dict = {k: v if v.size() == model_dict[k].size() else model_dict[k] for k, v in save_model[sub_item].items() if k in model_dict.keys()}
+        if not state_dict or not sub_item:
+            print('Weights are not successully loaded. Check the state dict of weights file.')
+            return None
+        else:
+            print('Found correct weights in the "{}" item of loaded state_dict.'.format(sub_item))
+    model_dict.update(state_dict)
+    model.load_state_dict(model_dict)
+    return model
+
+
+
+### models/modules/decoder_blocks.py
+
+import torch
+import torch.nn as nn
+# from models.aspp import ASPP, ASPPDeformable
+# from config import Config
+
+
+# config = Config()
+
+
+class BasicDecBlk(nn.Module):
+    def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
+        super(BasicDecBlk, self).__init__()
+        inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
+        self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
+        self.relu_in = nn.ReLU(inplace=True)
+        if config.dec_att == 'ASPP':
+            self.dec_att = ASPP(in_channels=inter_channels)
+        elif config.dec_att == 'ASPPDeformable':
+            self.dec_att = ASPPDeformable(in_channels=inter_channels)
+        self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
+        self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity()
+        self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
+
+    def forward(self, x):
+        x = self.conv_in(x)
+        x = self.bn_in(x)
+        x = self.relu_in(x)
+        if hasattr(self, 'dec_att'):
+            x = self.dec_att(x)
+        x = self.conv_out(x)
+        x = self.bn_out(x)
+        return x
+
+
+class ResBlk(nn.Module):
+    def __init__(self, in_channels=64, out_channels=None, inter_channels=64):
+        super(ResBlk, self).__init__()
+        if out_channels is None:
+            out_channels = in_channels
+        inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
+
+        self.conv_in = nn.Conv2d(in_channels, inter_channels, 3, 1, padding=1)
+        self.bn_in = nn.BatchNorm2d(inter_channels) if config.batch_size > 1 else nn.Identity()
+        self.relu_in = nn.ReLU(inplace=True)
+
+        if config.dec_att == 'ASPP':
+            self.dec_att = ASPP(in_channels=inter_channels)
+        elif config.dec_att == 'ASPPDeformable':
+            self.dec_att = ASPPDeformable(in_channels=inter_channels)
+
+        self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, padding=1)
+        self.bn_out = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
+        
+        self.conv_resi = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
+
+    def forward(self, x):
+        _x = self.conv_resi(x)
+        x = self.conv_in(x)
+        x = self.bn_in(x)
+        x = self.relu_in(x)
+        if hasattr(self, 'dec_att'):
+            x = self.dec_att(x)
+        x = self.conv_out(x)
+        x = self.bn_out(x)
+        return x + _x
+
+
+
+### models/modules/lateral_blocks.py
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+
+# from config import Config
+
+
+# config = Config()
+
+
+class BasicLatBlk(nn.Module):
+    def __init__(self, in_channels=64, out_channels=64, inter_channels=64):
+        super(BasicLatBlk, self).__init__()
+        inter_channels = in_channels // 4 if config.dec_channels_inter == 'adap' else 64
+        self.conv = nn.Conv2d(in_channels, out_channels, 1, 1, 0)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return x
+
+
+
+### models/modules/aspp.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+# from models.deform_conv import DeformableConv2d
+# from config import Config
+
+
+# config = Config()
+
+
+class _ASPPModule(nn.Module):
+    def __init__(self, in_channels, planes, kernel_size, padding, dilation):
+        super(_ASPPModule, self).__init__()
+        self.atrous_conv = nn.Conv2d(in_channels, planes, kernel_size=kernel_size,
+                                            stride=1, padding=padding, dilation=dilation, bias=False)
+        self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity()
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.atrous_conv(x)
+        x = self.bn(x)
+
+        return self.relu(x)
+
+
+class ASPP(nn.Module):
+    def __init__(self, in_channels=64, out_channels=None, output_stride=16):
+        super(ASPP, self).__init__()
+        self.down_scale = 1
+        if out_channels is None:
+            out_channels = in_channels
+        self.in_channelster = 256 // self.down_scale
+        if output_stride == 16:
+            dilations = [1, 6, 12, 18]
+        elif output_stride == 8:
+            dilations = [1, 12, 24, 36]
+        else:
+            raise NotImplementedError
+
+        self.aspp1 = _ASPPModule(in_channels, self.in_channelster, 1, padding=0, dilation=dilations[0])
+        self.aspp2 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[1], dilation=dilations[1])
+        self.aspp3 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[2], dilation=dilations[2])
+        self.aspp4 = _ASPPModule(in_channels, self.in_channelster, 3, padding=dilations[3], dilation=dilations[3])
+
+        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
+                                             nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False),
+                                             nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
+                                             nn.ReLU(inplace=True))
+        self.conv1 = nn.Conv2d(self.in_channelster * 5, out_channels, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
+        self.relu = nn.ReLU(inplace=True)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, x):
+        x1 = self.aspp1(x)
+        x2 = self.aspp2(x)
+        x3 = self.aspp3(x)
+        x4 = self.aspp4(x)
+        x5 = self.global_avg_pool(x)
+        x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
+        x = torch.cat((x1, x2, x3, x4, x5), dim=1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        return self.dropout(x)
+
+
+##################### Deformable
+class _ASPPModuleDeformable(nn.Module):
+    def __init__(self, in_channels, planes, kernel_size, padding):
+        super(_ASPPModuleDeformable, self).__init__()
+        self.atrous_conv = DeformableConv2d(in_channels, planes, kernel_size=kernel_size,
+                                            stride=1, padding=padding, bias=False)
+        self.bn = nn.BatchNorm2d(planes) if config.batch_size > 1 else nn.Identity()
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.atrous_conv(x)
+        x = self.bn(x)
+
+        return self.relu(x)
+
+
+class ASPPDeformable(nn.Module):
+    def __init__(self, in_channels, out_channels=None, parallel_block_sizes=[1, 3, 7]):
+        super(ASPPDeformable, self).__init__()
+        self.down_scale = 1
+        if out_channels is None:
+            out_channels = in_channels
+        self.in_channelster = 256 // self.down_scale
+
+        self.aspp1 = _ASPPModuleDeformable(in_channels, self.in_channelster, 1, padding=0)
+        self.aspp_deforms = nn.ModuleList([
+            _ASPPModuleDeformable(in_channels, self.in_channelster, conv_size, padding=int(conv_size//2)) for conv_size in parallel_block_sizes
+        ])
+
+        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
+                                             nn.Conv2d(in_channels, self.in_channelster, 1, stride=1, bias=False),
+                                             nn.BatchNorm2d(self.in_channelster) if config.batch_size > 1 else nn.Identity(),
+                                             nn.ReLU(inplace=True))
+        self.conv1 = nn.Conv2d(self.in_channelster * (2 + len(self.aspp_deforms)), out_channels, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels) if config.batch_size > 1 else nn.Identity()
+        self.relu = nn.ReLU(inplace=True)
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, x):
+        x1 = self.aspp1(x)
+        x_aspp_deforms = [aspp_deform(x) for aspp_deform in self.aspp_deforms]
+        x5 = self.global_avg_pool(x)
+        x5 = F.interpolate(x5, size=x1.size()[2:], mode='bilinear', align_corners=True)
+        x = torch.cat((x1, *x_aspp_deforms, x5), dim=1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        return self.dropout(x)
+
+
+
+### models/refinement/refiner.py
+
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import vgg16, vgg16_bn
+from torchvision.models import resnet50
+
+# from config import Config
+# from dataset import class_labels_TR_sorted
+# from models.build_backbone import build_backbone
+# from models.decoder_blocks import BasicDecBlk
+# from models.lateral_blocks import BasicLatBlk
+# from models.ing import *
+# from models.stem_layer import StemLayer
+
+
+class RefinerPVTInChannels4(nn.Module):
+    def __init__(self, in_channels=3+1):
+        super(RefinerPVTInChannels4, self).__init__()
+        self.config = Config()
+        self.epoch = 1
+        self.bb = build_backbone(self.config.bb, params_settings='in_channels=4')
+
+        lateral_channels_in_collection = {
+            'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
+            'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
+            'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
+        }
+        channels = lateral_channels_in_collection[self.config.bb]
+        self.squeeze_module = BasicDecBlk(channels[0], channels[0])
+
+        self.decoder = Decoder(channels)
+
+        if 0:
+            for key, value in self.named_parameters():
+                if 'bb.' in key:
+                    value.requires_grad = False
+
+    def forward(self, x):
+        if isinstance(x, list):
+            x = torch.cat(x, dim=1)
+        ########## Encoder ##########
+        if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
+            x1 = self.bb.conv1(x)
+            x2 = self.bb.conv2(x1)
+            x3 = self.bb.conv3(x2)
+            x4 = self.bb.conv4(x3)
+        else:
+            x1, x2, x3, x4 = self.bb(x)
+
+        x4 = self.squeeze_module(x4)
+
+        ########## Decoder ##########
+
+        features = [x, x1, x2, x3, x4]
+        scaled_preds = self.decoder(features)
+
+        return scaled_preds
+
+
+class Refiner(nn.Module):
+    def __init__(self, in_channels=3+1):
+        super(Refiner, self).__init__()
+        self.config = Config()
+        self.epoch = 1
+        self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
+        self.bb = build_backbone(self.config.bb)
+
+        lateral_channels_in_collection = {
+            'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
+            'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
+            'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
+        }
+        channels = lateral_channels_in_collection[self.config.bb]
+        self.squeeze_module = BasicDecBlk(channels[0], channels[0])
+
+        self.decoder = Decoder(channels)
+
+        if 0:
+            for key, value in self.named_parameters():
+                if 'bb.' in key:
+                    value.requires_grad = False
+
+    def forward(self, x):
+        if isinstance(x, list):
+            x = torch.cat(x, dim=1)
+        x = self.stem_layer(x)
+        ########## Encoder ##########
+        if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
+            x1 = self.bb.conv1(x)
+            x2 = self.bb.conv2(x1)
+            x3 = self.bb.conv3(x2)
+            x4 = self.bb.conv4(x3)
+        else:
+            x1, x2, x3, x4 = self.bb(x)
+
+        x4 = self.squeeze_module(x4)
+
+        ########## Decoder ##########
+
+        features = [x, x1, x2, x3, x4]
+        scaled_preds = self.decoder(features)
+
+        return scaled_preds
+
+
+class Decoder(nn.Module):
+    def __init__(self, channels):
+        super(Decoder, self).__init__()
+        self.config = Config()
+        DecoderBlock = eval('BasicDecBlk')
+        LateralBlock = eval('BasicLatBlk')
+
+        self.decoder_block4 = DecoderBlock(channels[0], channels[1])
+        self.decoder_block3 = DecoderBlock(channels[1], channels[2])
+        self.decoder_block2 = DecoderBlock(channels[2], channels[3])
+        self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2)
+
+        self.lateral_block4 = LateralBlock(channels[1], channels[1])
+        self.lateral_block3 = LateralBlock(channels[2], channels[2])
+        self.lateral_block2 = LateralBlock(channels[3], channels[3])
+
+        if self.config.ms_supervision:
+            self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0)
+            self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0)
+            self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0)
+        self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0))
+
+    def forward(self, features):
+        x, x1, x2, x3, x4 = features
+        outs = []
+        p4 = self.decoder_block4(x4)
+        _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
+        _p3 = _p4 + self.lateral_block4(x3)
+
+        p3 = self.decoder_block3(_p3)
+        _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
+        _p2 = _p3 + self.lateral_block3(x2)
+
+        p2 = self.decoder_block2(_p2)
+        _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
+        _p1 = _p2 + self.lateral_block2(x1)
+
+        _p1 = self.decoder_block1(_p1)
+        _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
+        p1_out = self.conv_out1(_p1)
+
+        if self.config.ms_supervision:
+            outs.append(self.conv_ms_spvn_4(p4))
+            outs.append(self.conv_ms_spvn_3(p3))
+            outs.append(self.conv_ms_spvn_2(p2))
+        outs.append(p1_out)
+        return outs
+
+
+class RefUNet(nn.Module):
+    # Refinement
+    def __init__(self, in_channels=3+1):
+        super(RefUNet, self).__init__()
+        self.encoder_1 = nn.Sequential(
+            nn.Conv2d(in_channels, 64, 3, 1, 1),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.encoder_2 = nn.Sequential(
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.encoder_3 = nn.Sequential(
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.encoder_4 = nn.Sequential(
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)
+        #####
+        self.decoder_5 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+        #####
+        self.decoder_4 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder_3 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder_2 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder_1 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1)
+
+        self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+
+    def forward(self, x):
+        outs = []
+        if isinstance(x, list):
+            x = torch.cat(x, dim=1)
+        hx = x
+
+        hx1 = self.encoder_1(hx)
+        hx2 = self.encoder_2(hx1)
+        hx3 = self.encoder_3(hx2)
+        hx4 = self.encoder_4(hx3)
+
+        hx = self.decoder_5(self.pool4(hx4))
+        hx = torch.cat((self.upscore2(hx), hx4), 1)
+
+        d4 = self.decoder_4(hx)
+        hx = torch.cat((self.upscore2(d4), hx3), 1)
+
+        d3 = self.decoder_3(hx)
+        hx = torch.cat((self.upscore2(d3), hx2), 1)
+
+        d2 = self.decoder_2(hx)
+        hx = torch.cat((self.upscore2(d2), hx1), 1)
+
+        d1 = self.decoder_1(hx)
+
+        x = self.conv_d0(d1)
+        outs.append(x)
+        return outs
+
+
+
+### models/stem_layer.py
+
+import torch.nn as nn
+# from utils import build_act_layer, build_norm_layer
+
+
+class StemLayer(nn.Module):
+    r""" Stem layer of InternImage
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+    """
+
+    def __init__(self,
+                 in_channels=3+1,
+                 inter_channels=48,
+                 out_channels=96,
+                 act_layer='GELU',
+                 norm_layer='BN'):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels,
+                               inter_channels,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.norm1 = build_norm_layer(
+            inter_channels, norm_layer, 'channels_first', 'channels_first'
+        )
+        self.act = build_act_layer(act_layer)
+        self.conv2 = nn.Conv2d(inter_channels,
+                               out_channels,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.norm2 = build_norm_layer(
+            out_channels, norm_layer, 'channels_first', 'channels_first'
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.norm2(x)
+        return x
+
+
+### models/birefnet.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from kornia.filters import laplacian
+from transformers import PreTrainedModel
+
+# from config import Config
+# from dataset import class_labels_TR_sorted
+# from models.build_backbone import build_backbone
+# from models.decoder_blocks import BasicDecBlk, ResBlk, HierarAttDecBlk
+# from models.lateral_blocks import BasicLatBlk
+# from models.aspp import ASPP, ASPPDeformable
+# from models.ing import *
+# from models.refiner import Refiner, RefinerPVTInChannels4, RefUNet
+# from models.stem_layer import StemLayer
+from .BiRefNet_config import BiRefNetConfig
+
+
+class BiRefNet(
+    PreTrainedModel
+):
+    config_class = BiRefNetConfig
+    def __init__(self, bb_pretrained=True, config=BiRefNetConfig()):
+        super(BiRefNet, self).__init__(config)
+        bb_pretrained = config.bb_pretrained
+        self.config = Config()
+        self.epoch = 1
+        self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained)
+
+        channels = self.config.lateral_channels_in_collection
+
+        if self.config.auxiliary_classification:
+            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+            self.cls_head = nn.Sequential(
+                nn.Linear(channels[0], len(class_labels_TR_sorted))
+            )
+
+        if self.config.squeeze_block:
+            self.squeeze_module = nn.Sequential(*[
+                eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0])
+                for _ in range(eval(self.config.squeeze_block.split('_x')[1]))
+            ])
+
+        self.decoder = Decoder(channels)
+
+        if self.config.ender:
+            self.dec_end = nn.Sequential(
+                nn.Conv2d(1, 16, 3, 1, 1),
+                nn.Conv2d(16, 1, 3, 1, 1),
+                nn.ReLU(inplace=True),
+            )
+
+        # refine patch-level segmentation
+        if self.config.refine:
+            if self.config.refine == 'itself':
+                self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
+            else:
+                self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1'))
+
+        if self.config.freeze_bb:
+            # Freeze the backbone...
+            print(self.named_parameters())
+            for key, value in self.named_parameters():
+                if 'bb.' in key and 'refiner.' not in key:
+                    value.requires_grad = False
+
+        self.post_init()
+
+    
+    def forward_enc(self, x):
+        if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
+            x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3)
+        else:
+            x1, x2, x3, x4 = self.bb(x)
+            if self.config.mul_scl_ipt == 'cat':
+                B, C, H, W = x.shape
+                x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True))
+                x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1)
+                x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1)
+                x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1)
+                x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1)
+            elif self.config.mul_scl_ipt == 'add':
+                B, C, H, W = x.shape
+                x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True))
+                x1 = x1 + F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)
+                x2 = x2 + F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)
+                x3 = x3 + F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)
+                x4 = x4 + F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)
+        class_preds = self.cls_head(self.avgpool(x4).view(x4.shape[0], -1)) if self.training and self.config.auxiliary_classification else None
+        if self.config.cxt:
+            x4 = torch.cat(
+                (
+                    *[
+                        F.interpolate(x1, size=x4.shape[2:], mode='bilinear', align_corners=True),
+                        F.interpolate(x2, size=x4.shape[2:], mode='bilinear', align_corners=True),
+                        F.interpolate(x3, size=x4.shape[2:], mode='bilinear', align_corners=True),
+                    ][-len(self.config.cxt):],
+                    x4
+                ),
+                dim=1
+            )
+        return (x1, x2, x3, x4), class_preds
+
+    def forward_ori(self, x):
+        ########## Encoder ##########
+        (x1, x2, x3, x4), class_preds = self.forward_enc(x)
+        if self.config.squeeze_block:
+            x4 = self.squeeze_module(x4)
+        ########## Decoder ##########
+        features = [x, x1, x2, x3, x4]
+        if self.training and self.config.out_ref:
+            features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5))
+        scaled_preds = self.decoder(features)
+        return scaled_preds, class_preds
+
+    def forward(self, x):
+        scaled_preds, class_preds = self.forward_ori(x)
+        class_preds_lst = [class_preds]
+        return [scaled_preds, class_preds_lst] if self.training else scaled_preds
+
+
+class Decoder(nn.Module):
+    def __init__(self, channels):
+        super(Decoder, self).__init__()
+        self.config = Config()
+        DecoderBlock = eval(self.config.dec_blk)
+        LateralBlock = eval(self.config.lat_blk)
+
+        if self.config.dec_ipt:
+            self.split = self.config.dec_ipt_split
+            N_dec_ipt = 64
+            DBlock = SimpleConvs
+            ic = 64
+            ipt_cha_opt = 1
+            self.ipt_blk5 = DBlock(2**10*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic)
+            self.ipt_blk4 = DBlock(2**8*3 if self.split else 3, [N_dec_ipt, channels[0]//8][ipt_cha_opt], inter_channels=ic)
+            self.ipt_blk3 = DBlock(2**6*3 if self.split else 3, [N_dec_ipt, channels[1]//8][ipt_cha_opt], inter_channels=ic)
+            self.ipt_blk2 = DBlock(2**4*3 if self.split else 3, [N_dec_ipt, channels[2]//8][ipt_cha_opt], inter_channels=ic)
+            self.ipt_blk1 = DBlock(2**0*3 if self.split else 3, [N_dec_ipt, channels[3]//8][ipt_cha_opt], inter_channels=ic)
+        else:
+            self.split = None
+
+        self.decoder_block4 = DecoderBlock(channels[0]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[1])
+        self.decoder_block3 = DecoderBlock(channels[1]+([N_dec_ipt, channels[0]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[2])
+        self.decoder_block2 = DecoderBlock(channels[2]+([N_dec_ipt, channels[1]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3])
+        self.decoder_block1 = DecoderBlock(channels[3]+([N_dec_ipt, channels[2]//8][ipt_cha_opt] if self.config.dec_ipt else 0), channels[3]//2)
+        self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2+([N_dec_ipt, channels[3]//8][ipt_cha_opt] if self.config.dec_ipt else 0), 1, 1, 1, 0))
+
+        self.lateral_block4 = LateralBlock(channels[1], channels[1])
+        self.lateral_block3 = LateralBlock(channels[2], channels[2])
+        self.lateral_block2 = LateralBlock(channels[3], channels[3])
+
+        if self.config.ms_supervision:
+            self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0)
+            self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0)
+            self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0)
+
+            if self.config.out_ref:
+                _N = 16
+                self.gdt_convs_4 = nn.Sequential(nn.Conv2d(channels[1], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
+                self.gdt_convs_3 = nn.Sequential(nn.Conv2d(channels[2], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
+                self.gdt_convs_2 = nn.Sequential(nn.Conv2d(channels[3], _N, 3, 1, 1), nn.BatchNorm2d(_N) if self.config.batch_size > 1 else nn.Identity(), nn.ReLU(inplace=True))
+
+                self.gdt_convs_pred_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                self.gdt_convs_pred_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                self.gdt_convs_pred_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                
+                self.gdt_convs_attn_4 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                self.gdt_convs_attn_3 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+                self.gdt_convs_attn_2 = nn.Sequential(nn.Conv2d(_N, 1, 1, 1, 0))
+
+    def get_patches_batch(self, x, p):
+        _size_h, _size_w = p.shape[2:]
+        patches_batch = []
+        for idx in range(x.shape[0]):
+            columns_x = torch.split(x[idx], split_size_or_sections=_size_w, dim=-1)
+            patches_x = []
+            for column_x in columns_x:
+                patches_x += [p.unsqueeze(0) for p in torch.split(column_x, split_size_or_sections=_size_h, dim=-2)]
+            patch_sample = torch.cat(patches_x, dim=1)
+            patches_batch.append(patch_sample)
+        return torch.cat(patches_batch, dim=0)
+
+    def forward(self, features):
+        if self.training and self.config.out_ref:
+            outs_gdt_pred = []
+            outs_gdt_label = []
+            x, x1, x2, x3, x4, gdt_gt = features
+        else:
+            x, x1, x2, x3, x4 = features
+        outs = []
+
+        if self.config.dec_ipt:
+            patches_batch = self.get_patches_batch(x, x4) if self.split else x
+            x4 = torch.cat((x4, self.ipt_blk5(F.interpolate(patches_batch, size=x4.shape[2:], mode='bilinear', align_corners=True))), 1)
+        p4 = self.decoder_block4(x4)
+        m4 = self.conv_ms_spvn_4(p4) if self.config.ms_supervision else None
+        if self.config.out_ref:
+            p4_gdt = self.gdt_convs_4(p4)
+            if self.training:
+                # >> GT:
+                m4_dia = m4
+                gdt_label_main_4 = gdt_gt * F.interpolate(m4_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
+                outs_gdt_label.append(gdt_label_main_4)
+                # >> Pred:
+                gdt_pred_4 = self.gdt_convs_pred_4(p4_gdt)
+                outs_gdt_pred.append(gdt_pred_4)
+            gdt_attn_4 = self.gdt_convs_attn_4(p4_gdt).sigmoid()
+            # >> Finally:
+            p4 = p4 * gdt_attn_4
+        _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
+        _p3 = _p4 + self.lateral_block4(x3)
+
+        if self.config.dec_ipt:
+            patches_batch = self.get_patches_batch(x, _p3) if self.split else x
+            _p3 = torch.cat((_p3, self.ipt_blk4(F.interpolate(patches_batch, size=x3.shape[2:], mode='bilinear', align_corners=True))), 1)
+        p3 = self.decoder_block3(_p3)
+        m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision else None
+        if self.config.out_ref:
+            p3_gdt = self.gdt_convs_3(p3)
+            if self.training:
+                # >> GT:
+                # m3 --dilation--> m3_dia
+                # G_3^gt * m3_dia --> G_3^m, which is the label of gradient
+                m3_dia = m3
+                gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
+                outs_gdt_label.append(gdt_label_main_3)
+                # >> Pred:
+                # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx
+                # F_3^G --sigmoid--> A_3^G
+                gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt)
+                outs_gdt_pred.append(gdt_pred_3)
+            gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid()
+            # >> Finally:
+            # p3 = p3 * A_3^G
+            p3 = p3 * gdt_attn_3
+        _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
+        _p2 = _p3 + self.lateral_block3(x2)
+
+        if self.config.dec_ipt:
+            patches_batch = self.get_patches_batch(x, _p2) if self.split else x
+            _p2 = torch.cat((_p2, self.ipt_blk3(F.interpolate(patches_batch, size=x2.shape[2:], mode='bilinear', align_corners=True))), 1)
+        p2 = self.decoder_block2(_p2)
+        m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision else None
+        if self.config.out_ref:
+            p2_gdt = self.gdt_convs_2(p2)
+            if self.training:
+                # >> GT:
+                m2_dia = m2
+                gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
+                outs_gdt_label.append(gdt_label_main_2)
+                # >> Pred:
+                gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt)
+                outs_gdt_pred.append(gdt_pred_2)
+            gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid()
+            # >> Finally:
+            p2 = p2 * gdt_attn_2
+        _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
+        _p1 = _p2 + self.lateral_block2(x1)
+
+        if self.config.dec_ipt:
+            patches_batch = self.get_patches_batch(x, _p1) if self.split else x
+            _p1 = torch.cat((_p1, self.ipt_blk2(F.interpolate(patches_batch, size=x1.shape[2:], mode='bilinear', align_corners=True))), 1)
+        _p1 = self.decoder_block1(_p1)
+        _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
+
+        if self.config.dec_ipt:
+            patches_batch = self.get_patches_batch(x, _p1) if self.split else x
+            _p1 = torch.cat((_p1, self.ipt_blk1(F.interpolate(patches_batch, size=x.shape[2:], mode='bilinear', align_corners=True))), 1)
+        p1_out = self.conv_out1(_p1)
+
+        if self.config.ms_supervision:
+            outs.append(m4)
+            outs.append(m3)
+            outs.append(m2)
+        outs.append(p1_out)
+        return outs if not (self.config.out_ref and self.training) else ([outs_gdt_pred, outs_gdt_label], outs)
+
+
+class SimpleConvs(nn.Module):
+    def __init__(
+        self, in_channels: int, out_channels: int, inter_channels=64
+    ) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels, inter_channels, 3, 1, 1)
+        self.conv_out = nn.Conv2d(inter_channels, out_channels, 3, 1, 1)
+
+    def forward(self, x):
+        return self.conv_out(self.conv1(x))

rmbg/config.json

@@ -0,0 +1,20 @@
+{
+  "_name_or_path": "ZhengPeng7/BiRefNet",
+  "architectures": [
+    "BiRefNet"
+  ],
+  "auto_map": {
+    "AutoConfig": "BiRefNet_config.BiRefNetConfig",
+    "AutoModelForImageSegmentation": "birefnet.BiRefNet"
+  },
+  "custom_pipelines": {
+    "image-segmentation": {
+      "pt": [
+        "AutoModelForImageSegmentation"
+      ],
+      "tf": [],
+      "type": "image"
+    }
+  },
+  "bb_pretrained": false
+}

rmbg/model.safetensors

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

texturing_pipeline.json

@@ -0,0 +1,64 @@
+{
+    "name": "Trellis2TexturingPipeline",
+    "args": {
+        "models": {
+            "shape_slat_encoder": "ckpts/shape_enc_next_dc_f16c32_fp16",
+            "tex_slat_decoder": "ckpts/tex_dec_next_dc_f16c32_fp16",
+            "tex_slat_flow_model_512": "ckpts/slat_flow_imgshape2tex_dit_1_3B_512_bf16",
+            "tex_slat_flow_model_1024": "ckpts/slat_flow_imgshape2tex_dit_1_3B_1024_bf16"
+        },
+        "shape_slat_normalization": {
+            "mean": [
+                0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+            ],
+            "std": [
+                5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+            ]
+        },
+        "tex_slat_sampler": {
+            "name": "FlowEulerGuidanceIntervalSampler",
+            "args": {
+                "sigma_min": 1e-5
+            },
+            "params": {
+                "steps": 12,
+                "guidance_strength": 1.0,
+                "guidance_rescale": 0.0,
+                "guidance_interval": [0.6, 0.9],
+                "rescale_t": 3.0
+            }
+        },
+        "tex_slat_normalization": {
+            "mean": [
+                3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+                0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+                -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+                1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+            ],
+            "std": [
+                2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+                2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+                2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+                2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+            ]
+        },
+        "image_cond_model": {
+            "name": "DinoV3FeatureExtractor",
+            "args": {
+                "model_name": "facebook/dinov3-vitl16-pretrain-lvd1689m"
+            }
+        },
+        "rembg_model": {
+            "name": "BiRefNet",
+            "args": {
+                "model_name": "briaai/RMBG-2.0"
+            }
+        }
+    }
+}