Upload 19 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+32-photo.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,5 +1,152 @@
----
-license: other
-license_name: bria-3.2
-license_link: LICENSE
----
+---
+library_name: diffusers
+inference: false
+tags:
+- text-to-image
+- legal liability
+- Non Commercial Use
+extra_gated_description: >-
+  Bria AI Model weights are open source for non commercial use only, per the
+  provided [license](https://creativecommons.org/licenses/by-nc/4.0/deed.en).
+extra_gated_heading: Fill in this form to immediatly access the model for non commercial use
+extra_gated_fields:
+  Name: text
+  Email: text
+  Company/Org name: text
+  Company Website URL: text
+  Discord user: text
+  I agree to BRIA’s Privacy policy, Terms & conditions, and acknowledge Non commercial use to be Personal use / Academy / Non profit (direct or indirect): checkbox
+license: other
+license_name: bria-3.2
+license_link: https://creativecommons.org/licenses/by-nc/4.0/deed.en
+---
+
+
+
+# TL;DR
+
+Bria 3.2 is the next-generation commercial-ready text-to-image model.
+**With just 4 billion parameters**, it provides exceptional aesthetics and text rendering, evaluated to provide **on par results to leading open-source models, and outperforming other licensed models**.
+
+In addition to being **built entirely on licensed data**, 3.2 provides several advantages for enterprise and commercial use:
+* Efficient Compute - the model is X3 smaller than the equivalent models in the market (4B parameters vs 12B parameters other open source models)
+* Architecture Consistency: Same architecture as 3.1—ideal for users looking to upgrade without disruption.
+* Fine-tuning Speedup: 2x faster fine-tuning on L40S and A100.
+
+
+[CLICK HERE FOR A DEMO](https://huggingface.co/spaces/briaai/BRIA-3.2-API)
+
+
+
+![](32-photo.jpg)
+
+
+
+# BRIA 3.2: Training data and Commercial Licensing
+
+BRIA 3.2 is our latest text-to-image model explicitly designed for commercial applications. 
+This model combines technological innovation with ethical responsibility and legal security, setting a new standard in the AI industry. 
+Bria AI licenses the foundation model with full legal liability coverage. 
+Our dataset does not contain copyrighted materials, such as fictional characters, logos, trademarks, public figures, harmful content, or privacy-infringing content.
+
+For more information, please visit our [website](https://bria.ai/).
+
+Join our [Discord community](https://discord.gg/Nxe9YW9zHS) for more information, tutorials, tools, and to connect with other users!
+
+### For Commercial License : click [Here](https://bria.ai/contact-us?hsCtaAttrib=114250296256).
+
+
+# What's New vs pervious models:
+
+- **Improved Aesthetics**:
+
+  - 65% user preference for BRIA 3.2 over BRIA 3.1.
+  - 76% user preference for BRIA 3.2 over BRIA 2.3.
+
+- **Superior Text Rendering**: The model is optimized to generate short text consists of 1-6 words. OCR Score improvement from 5% (3.1) to 70% (3.2).
+
+- **Consistent Prompt Alignment**: Maintains high-quality textual description adherence.
+
+
+
+### Get Access
+Bria 3.2 is available everywhere you build, either as source-code and weights, ComfyUI nodes or API endpoints.
+
+- **API Endpoint**: [Bria.ai](https://docs.bria.ai/image-generation/endpoints/text-to-image-base) , [Fal.ai](https://fal.ai/models/bria/text-to-image/3.2), [Replicate](https://replicate.com/bria/image-3.2)
+- **ComfyUI**: [Use it in workflows](https://github.com/Bria-AI/ComfyUI-BRIA-API)
+- **GitHub**: [github.com/Bria-AI/BRIA-3.2](https://github.com/Bria-AI/BRIA-3.2)
+- **Interested in BRIA 3.2 source code and weights for commercial use?** Purchase is required to license BRIA 3.2 got commercial use, ensuring royalty management with our data partners and full liability coverage.
+- Are you a startup or a student? We encourage you to apply for our [Startup Program](https://pages.bria.ai/the-visual-generative-ai-platform-for-builders-startups-plan?_gl=1*cqrl81*_ga*MTIxMDI2NzI5OC4xNjk5NTQ3MDAz*_ga_WRN60H46X4*MTcwOTM5OTMzNC4yNzguMC4xNzA5Mzk5MzM0LjYwLjAuMA..) to request access. This program are designed to support emerging businesses and academic pursuits with our cutting-edge technology.
+- By submitting the form above, you agree to BRIA’s [Privacy policy](https://bria.ai/privacy-policy/) and [Terms & conditions](https://bria.ai/terms-and-conditions/).
+
+
+
+# Key Features
+
+- **Architecture**: 4B parameter, rectified flow transformer based model with T5 text encoder.
+
+- **Legally Compliant**: Offers full legal liability coverage for copyright and privacy infringements. Thanks to training on 100% licensed data from leading data partners, we ensure the ethical use of content.
+
+- **Patented Attribution Engine**: Our attribution engine is our way to compensate our data partners, powered by our proprietary and patented algorithms. 
+
+- **Enterprise-Ready**: Specifically designed for business applications, Bria AI 3.2 delivers high-quality, compliant imagery for a variety of commercial needs.
+
+- **Customizable Technology**: Provides access to source code and weights for extensive customization, catering to specific business requirements.
+
+### Model Description
+
+- **Developed by:** BRIA AI
+- **Model type:** Latent diffusion text-to-image model
+- **Resources for more information:** [BRIA AI](https://bria.ai/)
+
+
+
+
+### Code example using Diffusers 
+
+
+```python
+pip install diffusers, hf_hub_download
+```
+
+
+
+```python
+from huggingface_hub import hf_hub_download
+import os
+
+try:
+    local_dir = os.path.dirname(__file__)
+except:
+    local_dir = '.'
+    
+hf_hub_download(repo_id="briaai/BRIA-3.2", filename='pipeline_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.2", filename='transformer_bria.py', local_dir=local_dir)
+hf_hub_download(repo_id="briaai/BRIA-3.2", filename='bria_utils.py', local_dir=local_dir)
+
+import torch
+from pipeline_bria import BriaPipeline, BriaTransformer2DModel
+
+# trust_remote_code = True - allows loading a transformer which is not present at the transformers library(from transformer/bria_transformer.py)
+pipe = BriaPipeline.from_pretrained("briaai/BRIA-3.2", torch_dtype=torch.bfloat16,trust_remote_code=True)
+pipe.to(device="cuda")
+
+prompt = "A portrait of a Beautiful and playful ethereal singer, golden designs, highly detailed, blurry background"
+negative_prompt = "Logo,Watermark,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
+
+images = pipe(prompt=prompt, negative_prompt=negative_prompt, height=1024, width=1024).images[0]
+```
+
+
+### Some tips for using our text-to-image model at inference: 
+
+
+1. Using negative prompt is recommended.
+2. For Fine-tuning, use zeros instead of null text embedding.
+3. We support multiple aspect ratios, yet resolution should overall consists approximately `1024*1024=1M` pixels, for example:
+`((1024,1024), (1280, 768), (1344, 768), (832, 1216), (1152, 832), (1216, 832), (960,1088)`
+
+4. Use 30-50 steps (higher is better)
+5. Use `guidance_scale` of  5.0
+
+

model_index.json

@@ -0,0 +1,25 @@
+{
+  "_class_name": "BriaPipeline",
+  "_diffusers_version": "0.33.1",
+  "_name_or_path": "briaai/BRIA-3.1",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler"
+  ],
+  "text_encoder": [
+    "transformers",
+    "T5EncoderModel"
+  ],
+  "tokenizer": [
+    "transformers",
+    "T5TokenizerFast"
+  ],
+  "transformer": [
+    "transformer_bria",
+    "BriaTransformer2DModel"
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKL"
+  ]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,17 @@
+{
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.33.1",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "invert_sigmas": false,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "shift_terminal": null,
+  "time_shift_type": "exponential",
+  "use_beta_sigmas": false,
+  "use_dynamic_shifting": true,
+  "use_exponential_sigmas": false,
+  "use_karras_sigmas": false
+}

text_encoder/config.json

@@ -0,0 +1,31 @@
+{
+  "architectures": [
+    "T5EncoderModel"
+  ],
+  "classifier_dropout": 0.0,
+  "d_ff": 10240,
+  "d_kv": 64,
+  "d_model": 4096,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "gelu_new",
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "num_decoder_layers": 24,
+  "num_heads": 64,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float16",
+  "transformers_version": "4.52.0",
+  "use_cache": true,
+  "vocab_size": 32128
+}

text_encoder/model-00001-of-00003.safetensors

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

text_encoder/model-00002-of-00003.safetensors

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

text_encoder/model-00003-of-00003.safetensors

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

text_encoder/model.safetensors.index.json

@@ -0,0 +1,226 @@
+{
+  "metadata": {
+    "total_size": 11537887232
+  },
+  "weight_map": {
+    "encoder.block.0.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.0.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.1.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.10.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.11.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.12.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.13.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.14.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.15.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.16.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.17.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.18.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wo.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.19.layer.1.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.2.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.k.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.o.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.q.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.v.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.0.layer_norm.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00003.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wo.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.20.layer.1.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.k.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.o.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.q.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.v.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.0.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wi_0.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wi_1.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wo.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.21.layer.1.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.k.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.o.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.q.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.v.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.0.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wi_0.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wi_1.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wo.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.22.layer.1.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.k.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.o.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.q.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.v.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.0.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wi_0.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wi_1.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wo.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.23.layer.1.layer_norm.weight": "model-00003-of-00003.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.3.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.4.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.5.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.6.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.7.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.8.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.k.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.o.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.q.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.v.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.0.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wo.weight": "model-00001-of-00003.safetensors",
+    "encoder.block.9.layer.1.layer_norm.weight": "model-00001-of-00003.safetensors",
+    "encoder.final_layer_norm.weight": "model-00003-of-00003.safetensors",
+    "shared.weight": "model-00001-of-00003.safetensors"
+  }
+}

tokenizer/special_tokens_map.json

@@ -0,0 +1,125 @@
+{
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer/spiece.model

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

tokenizer/tokenizer_config.json

@@ -0,0 +1,941 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "32000": {
+      "content": "<extra_id_99>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32001": {
+      "content": "<extra_id_98>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32002": {
+      "content": "<extra_id_97>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32003": {
+      "content": "<extra_id_96>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32004": {
+      "content": "<extra_id_95>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32005": {
+      "content": "<extra_id_94>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32006": {
+      "content": "<extra_id_93>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32007": {
+      "content": "<extra_id_92>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32008": {
+      "content": "<extra_id_91>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32009": {
+      "content": "<extra_id_90>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32010": {
+      "content": "<extra_id_89>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32011": {
+      "content": "<extra_id_88>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32012": {
+      "content": "<extra_id_87>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32013": {
+      "content": "<extra_id_86>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32014": {
+      "content": "<extra_id_85>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32015": {
+      "content": "<extra_id_84>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32016": {
+      "content": "<extra_id_83>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32017": {
+      "content": "<extra_id_82>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32018": {
+      "content": "<extra_id_81>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32019": {
+      "content": "<extra_id_80>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32020": {
+      "content": "<extra_id_79>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32021": {
+      "content": "<extra_id_78>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32022": {
+      "content": "<extra_id_77>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32023": {
+      "content": "<extra_id_76>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32024": {
+      "content": "<extra_id_75>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32025": {
+      "content": "<extra_id_74>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32026": {
+      "content": "<extra_id_73>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32027": {
+      "content": "<extra_id_72>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32028": {
+      "content": "<extra_id_71>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32029": {
+      "content": "<extra_id_70>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32030": {
+      "content": "<extra_id_69>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32031": {
+      "content": "<extra_id_68>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32032": {
+      "content": "<extra_id_67>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32033": {
+      "content": "<extra_id_66>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32034": {
+      "content": "<extra_id_65>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32035": {
+      "content": "<extra_id_64>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32036": {
+      "content": "<extra_id_63>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32037": {
+      "content": "<extra_id_62>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32038": {
+      "content": "<extra_id_61>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32039": {
+      "content": "<extra_id_60>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32040": {
+      "content": "<extra_id_59>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32041": {
+      "content": "<extra_id_58>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32042": {
+      "content": "<extra_id_57>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32043": {
+      "content": "<extra_id_56>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32044": {
+      "content": "<extra_id_55>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32045": {
+      "content": "<extra_id_54>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32046": {
+      "content": "<extra_id_53>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32047": {
+      "content": "<extra_id_52>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32048": {
+      "content": "<extra_id_51>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32049": {
+      "content": "<extra_id_50>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32050": {
+      "content": "<extra_id_49>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32051": {
+      "content": "<extra_id_48>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32052": {
+      "content": "<extra_id_47>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32053": {
+      "content": "<extra_id_46>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32054": {
+      "content": "<extra_id_45>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32055": {
+      "content": "<extra_id_44>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32056": {
+      "content": "<extra_id_43>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32057": {
+      "content": "<extra_id_42>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32058": {
+      "content": "<extra_id_41>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32059": {
+      "content": "<extra_id_40>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32060": {
+      "content": "<extra_id_39>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32061": {
+      "content": "<extra_id_38>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32062": {
+      "content": "<extra_id_37>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32063": {
+      "content": "<extra_id_36>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32064": {
+      "content": "<extra_id_35>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32065": {
+      "content": "<extra_id_34>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32066": {
+      "content": "<extra_id_33>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32067": {
+      "content": "<extra_id_32>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32068": {
+      "content": "<extra_id_31>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32069": {
+      "content": "<extra_id_30>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32070": {
+      "content": "<extra_id_29>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32071": {
+      "content": "<extra_id_28>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32072": {
+      "content": "<extra_id_27>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32073": {
+      "content": "<extra_id_26>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32074": {
+      "content": "<extra_id_25>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32075": {
+      "content": "<extra_id_24>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32076": {
+      "content": "<extra_id_23>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32077": {
+      "content": "<extra_id_22>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32078": {
+      "content": "<extra_id_21>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32079": {
+      "content": "<extra_id_20>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32080": {
+      "content": "<extra_id_19>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32081": {
+      "content": "<extra_id_18>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32082": {
+      "content": "<extra_id_17>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32083": {
+      "content": "<extra_id_16>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32084": {
+      "content": "<extra_id_15>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32085": {
+      "content": "<extra_id_14>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32086": {
+      "content": "<extra_id_13>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32087": {
+      "content": "<extra_id_12>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32088": {
+      "content": "<extra_id_11>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32089": {
+      "content": "<extra_id_10>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32090": {
+      "content": "<extra_id_9>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32091": {
+      "content": "<extra_id_8>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32092": {
+      "content": "<extra_id_7>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32093": {
+      "content": "<extra_id_6>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32094": {
+      "content": "<extra_id_5>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32095": {
+      "content": "<extra_id_4>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32096": {
+      "content": "<extra_id_3>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32097": {
+      "content": "<extra_id_2>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32098": {
+      "content": "<extra_id_1>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "32099": {
+      "content": "<extra_id_0>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<extra_id_0>",
+    "<extra_id_1>",
+    "<extra_id_2>",
+    "<extra_id_3>",
+    "<extra_id_4>",
+    "<extra_id_5>",
+    "<extra_id_6>",
+    "<extra_id_7>",
+    "<extra_id_8>",
+    "<extra_id_9>",
+    "<extra_id_10>",
+    "<extra_id_11>",
+    "<extra_id_12>",
+    "<extra_id_13>",
+    "<extra_id_14>",
+    "<extra_id_15>",
+    "<extra_id_16>",
+    "<extra_id_17>",
+    "<extra_id_18>",
+    "<extra_id_19>",
+    "<extra_id_20>",
+    "<extra_id_21>",
+    "<extra_id_22>",
+    "<extra_id_23>",
+    "<extra_id_24>",
+    "<extra_id_25>",
+    "<extra_id_26>",
+    "<extra_id_27>",
+    "<extra_id_28>",
+    "<extra_id_29>",
+    "<extra_id_30>",
+    "<extra_id_31>",
+    "<extra_id_32>",
+    "<extra_id_33>",
+    "<extra_id_34>",
+    "<extra_id_35>",
+    "<extra_id_36>",
+    "<extra_id_37>",
+    "<extra_id_38>",
+    "<extra_id_39>",
+    "<extra_id_40>",
+    "<extra_id_41>",
+    "<extra_id_42>",
+    "<extra_id_43>",
+    "<extra_id_44>",
+    "<extra_id_45>",
+    "<extra_id_46>",
+    "<extra_id_47>",
+    "<extra_id_48>",
+    "<extra_id_49>",
+    "<extra_id_50>",
+    "<extra_id_51>",
+    "<extra_id_52>",
+    "<extra_id_53>",
+    "<extra_id_54>",
+    "<extra_id_55>",
+    "<extra_id_56>",
+    "<extra_id_57>",
+    "<extra_id_58>",
+    "<extra_id_59>",
+    "<extra_id_60>",
+    "<extra_id_61>",
+    "<extra_id_62>",
+    "<extra_id_63>",
+    "<extra_id_64>",
+    "<extra_id_65>",
+    "<extra_id_66>",
+    "<extra_id_67>",
+    "<extra_id_68>",
+    "<extra_id_69>",
+    "<extra_id_70>",
+    "<extra_id_71>",
+    "<extra_id_72>",
+    "<extra_id_73>",
+    "<extra_id_74>",
+    "<extra_id_75>",
+    "<extra_id_76>",
+    "<extra_id_77>",
+    "<extra_id_78>",
+    "<extra_id_79>",
+    "<extra_id_80>",
+    "<extra_id_81>",
+    "<extra_id_82>",
+    "<extra_id_83>",
+    "<extra_id_84>",
+    "<extra_id_85>",
+    "<extra_id_86>",
+    "<extra_id_87>",
+    "<extra_id_88>",
+    "<extra_id_89>",
+    "<extra_id_90>",
+    "<extra_id_91>",
+    "<extra_id_92>",
+    "<extra_id_93>",
+    "<extra_id_94>",
+    "<extra_id_95>",
+    "<extra_id_96>",
+    "<extra_id_97>",
+    "<extra_id_98>",
+    "<extra_id_99>"
+  ],
+  "clean_up_tokenization_spaces": true,
+  "eos_token": "</s>",
+  "extra_ids": 100,
+  "extra_special_tokens": {},
+  "legacy": true,
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "T5Tokenizer",
+  "unk_token": "<unk>"
+}

transformer/config.json

@@ -0,0 +1,20 @@
+  {
+  "_class_name": "BriaTransformer2DModel",
+  "_diffusers_version": "0.33.1",
+  "attention_head_dim": 96,
+  "axes_dims_rope": [
+    0,
+    48,
+    48
+  ],
+  "guidance_embeds": false,
+  "in_channels": 16,
+  "joint_attention_dim": 4096,
+  "num_attention_heads": 24,
+  "num_layers": 8,
+  "num_single_layers": 28,
+  "patch_size": 1,
+  "pooled_projection_dim": null,
+  "rope_theta": 10000,
+  "time_theta": 10000
+}

transformer/diffusion_pytorch_model.safetensors

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

transformer/transformer_bria.py

@@ -0,0 +1,315 @@
+from typing import Any, Dict, List, Optional, Union
+import numpy as np
+import torch
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous
+from diffusers.utils import USE_PEFT_BACKEND, is_torch_version, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.embeddings import TimestepEmbedding, get_timestep_embedding
+from diffusers.models.transformers.transformer_flux import FluxSingleTransformerBlock, FluxTransformerBlock
+from bria_utils import FluxPosEmbed as EmbedND
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+class Timesteps(nn.Module):
+    def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1,time_theta=10000):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+        self.time_theta=time_theta
+
+    def forward(self, timesteps):
+        t_emb = get_timestep_embedding(
+            timesteps,
+            self.num_channels,
+            flip_sin_to_cos=self.flip_sin_to_cos,
+            downscale_freq_shift=self.downscale_freq_shift,
+            scale=self.scale,
+            max_period=self.time_theta
+        )
+        return t_emb
+    
+class TimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim, time_theta):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0,time_theta=time_theta)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+        
+    def forward(self, timestep, dtype):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=dtype))  # (N, D)
+        return timesteps_emb
+    
+"""
+Based on FluxPipeline with several changes:
+- no pooled embeddings
+- We use zero padding for prompts
+- No guidance embedding since this is not a distilled version
+"""
+class BriaTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):
+    """
+    The Transformer model introduced in Flux.
+
+    Reference: https://blackforestlabs.ai/announcing-black-forest-labs/
+
+    Parameters:
+        patch_size (`int`): Patch size to turn the input data into small patches.
+        in_channels (`int`, *optional*, defaults to 16): The number of channels in the input.
+        num_layers (`int`, *optional*, defaults to 18): The number of layers of MMDiT blocks to use.
+        num_single_layers (`int`, *optional*, defaults to 18): The number of layers of single DiT blocks to use.
+        attention_head_dim (`int`, *optional*, defaults to 64): The number of channels in each head.
+        num_attention_heads (`int`, *optional*, defaults to 18): The number of heads to use for multi-head attention.
+        joint_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
+        pooled_projection_dim (`int`): Number of dimensions to use when projecting the `pooled_projections`.
+        guidance_embeds (`bool`, defaults to False): Whether to use guidance embeddings.
+    """
+
+    _supports_gradient_checkpointing = True
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 1,
+        in_channels: int = 64,
+        num_layers: int = 19,
+        num_single_layers: int = 38,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 4096,
+        pooled_projection_dim: int = None,
+        guidance_embeds: bool = False,
+        axes_dims_rope: List[int] = [16, 56, 56],
+        rope_theta = 10000,
+        time_theta = 10000
+    ):
+        super().__init__()
+        self.out_channels = in_channels
+        self.inner_dim = self.config.num_attention_heads * self.config.attention_head_dim
+
+        self.pos_embed = EmbedND(theta=rope_theta, axes_dim=axes_dims_rope)
+        
+        
+        self.time_embed = TimestepProjEmbeddings(
+            embedding_dim=self.inner_dim,time_theta=time_theta
+        )
+
+        # if pooled_projection_dim:
+        #     self.pooled_text_embed = PixArtAlphaTextProjection(pooled_projection_dim, embedding_dim=self.inner_dim, act_fn="silu")
+    
+        if guidance_embeds:
+            self.guidance_embed = TimestepProjEmbeddings(embedding_dim=self.inner_dim)
+    
+        self.context_embedder = nn.Linear(self.config.joint_attention_dim, self.inner_dim)
+        self.x_embedder = torch.nn.Linear(self.config.in_channels, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                FluxTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_layers)
+            ]
+        )
+
+        self.single_transformer_blocks = nn.ModuleList(
+            [
+                FluxSingleTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=self.config.num_attention_heads,
+                    attention_head_dim=self.config.attention_head_dim,
+                )
+                for i in range(self.config.num_single_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if hasattr(module, "gradient_checkpointing"):
+            module.gradient_checkpointing = value
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        pooled_projections: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
+        guidance: torch.Tensor = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+        controlnet_block_samples = None,
+        controlnet_single_block_samples=None,
+
+    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
+        """
+        The [`FluxTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
+                from the embeddings of input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
+                A list of tensors that if specified are added to the residuals of transformer blocks.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if joint_attention_kwargs is not None:
+            joint_attention_kwargs = joint_attention_kwargs.copy()
+            lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+        hidden_states = self.x_embedder(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype)
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype)
+        else:
+            guidance = None
+
+        # temb = (
+        #     self.time_text_embed(timestep, pooled_projections)
+        #     if guidance is None
+        #     else self.time_text_embed(timestep, guidance, pooled_projections)
+        # )
+
+        temb = self.time_embed(timestep,dtype=hidden_states.dtype)
+
+        # if pooled_projections:
+        #     temb+=self.pooled_text_embed(pooled_projections)
+        
+        if guidance:
+            temb+=self.guidance_embed(guidance,dtype=hidden_states.dtype)
+
+        encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+        if len(txt_ids.shape)==2:
+            ids = torch.cat((txt_ids, img_ids), dim=0)
+        else:
+            ids = torch.cat((txt_ids, img_ids), dim=1)
+        image_rotary_emb = self.pos_embed(ids)
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    encoder_hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+            
+            # controlnet residual
+            if controlnet_block_samples is not None:
+                interval_control = len(self.transformer_blocks) / len(controlnet_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
+
+
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        for index_block, block in enumerate(self.single_transformer_blocks):
+            if self.training and self.gradient_checkpointing:
+
+                def create_custom_forward(module, return_dict=None):
+                    def custom_forward(*inputs):
+                        if return_dict is not None:
+                            return module(*inputs, return_dict=return_dict)
+                        else:
+                            return module(*inputs)
+
+                    return custom_forward
+
+                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+                hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    temb,
+                    image_rotary_emb,
+                    **ckpt_kwargs,
+                )
+
+            else:
+                hidden_states = block(
+                    hidden_states=hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                )
+            
+            # controlnet residual
+            if controlnet_single_block_samples is not None:
+                interval_control = len(self.single_transformer_blocks) / len(controlnet_single_block_samples)
+                interval_control = int(np.ceil(interval_control))
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                    + controlnet_single_block_samples[index_block // interval_control]
+                )
+
+        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

vae/config.json

@@ -0,0 +1,38 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.33.1",
+  "_name_or_path": "/home/ubuntu/.cache/huggingface/hub/models--briaai--BRIA-3.1/snapshots/82e82697e2a3a67726e099f5d2ba15e4520ff93c/vae",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": false,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "latents_mean": null,
+  "latents_std": null,
+  "layers_per_block": 2,
+  "mid_block_add_attention": true,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 512,
+  "scaling_factor": 0.13025,
+  "shift_factor": null,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ],
+  "use_post_quant_conv": true,
+  "use_quant_conv": true
+}

vae/diffusion_pytorch_model.safetensors

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