From 56246711915552a1f7bf31e3d99a5508a2b360e3 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Fri, 6 Jun 2025 19:25:26 +0300
Subject: [PATCH] Fix PixArt Sigma Small and Large

---
 CHANGELOG.md            |  1 +
 modules/model_pixart.py | 12 ++++++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 80ed4f77a..4b0b1d531 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@
 - **Fixes**
   - Meissonic with multiple generators  
   - Kandinsky V2.2 invalid attention processor  
+  - PixArt Sigma Small and Large loading  
 
 ## Update for 2025-06-02
 
diff --git a/modules/model_pixart.py b/modules/model_pixart.py
index 58d3dbfa7..6f6d6cf1c 100644
--- a/modules/model_pixart.py
+++ b/modules/model_pixart.py
@@ -1,11 +1,19 @@
 import transformers
 import diffusers
+from huggingface_hub import file_exists
 
 
 def load_pixart(checkpoint_info, diffusers_load_config={}):
     from modules import shared, devices, modelloader, sd_models, model_quant
     modelloader.hf_login()
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    repo_id_tenc = repo_id
+    repo_id_pipe = repo_id
+
+    if not file_exists(repo_id_tenc, "text_encoder/config.json"):
+        repo_id_tenc = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
+    if not file_exists(repo_id_pipe, "model_index.json"):
+        repo_id_pipe = "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS"
 
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
     transformer = diffusers.PixArtTransformer2DModel.from_pretrained(
@@ -17,7 +25,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
     )
     load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     text_encoder = transformers.T5EncoderModel.from_pretrained(
-        repo_id,
+        repo_id_tenc,
         subfolder="text_encoder",
         cache_dir=shared.opts.hfcache_dir,
         **load_args,
@@ -26,7 +34,7 @@ def load_pixart(checkpoint_info, diffusers_load_config={}):
 
     load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
     pipe = diffusers.PixArtSigmaPipeline.from_pretrained(
-        'PixArt-alpha/PixArt-Sigma-XL-2-1024-MS',
+        repo_id_pipe,
         cache_dir=shared.opts.diffusers_dir,
         transformer=transformer,
         text_encoder=text_encoder,