From 3a9f56a7ac7f1134082aa1f54d955cdc8042a93b Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Thu, 13 Jun 2024 09:07:47 -0400
Subject: [PATCH] sd3 allow loading safetensors without clip

---
 modules/model_sd3.py | 44 +++++++++++++++++++++++++++++---------------
 modules/sd_models.py |  2 +-
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/modules/model_sd3.py b/modules/model_sd3.py
index 583d317f3..11d5808c2 100644
--- a/modules/model_sd3.py
+++ b/modules/model_sd3.py
@@ -26,14 +26,37 @@ def load_sd3(te3=None, fn=None, cache_dir=None, config=None):
     repo_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
     model_id = 'stabilityai/stable-diffusion-3-medium-diffusers'
     dtype = torch.float16
+    kwargs = {}
     if fn is not None and fn.endswith('.safetensors') and os.path.exists(fn):
         model_id = fn
         loader = diffusers.StableDiffusion3Pipeline.from_single_file
-        reload_te = True
+        kwargs = {
+            'text_encoder': transformers.CLIPTextModelWithProjection.from_pretrained(
+                repo_id,
+                subfolder='text_encoder',
+                cache_dir=cache_dir,
+                torch_dtype=dtype,
+            ),
+            'text_encoder_2': transformers.CLIPTextModelWithProjection.from_pretrained(
+                repo_id,
+                subfolder='text_encoder_2',
+                cache_dir=cache_dir,
+                torch_dtype=dtype,
+            ),
+            'tokenizer': transformers.CLIPTokenizer.from_pretrained(
+                repo_id,
+                subfolder='tokenizer',
+                cache_dir=cache_dir,
+            ),
+            'tokenizer_2': transformers.CLIPTokenizer.from_pretrained(
+                repo_id,
+                subfolder='tokenizer_2',
+                cache_dir=cache_dir,
+            ),
+        }
     else:
         model_id = repo_id
         loader = diffusers.StableDiffusion3Pipeline.from_pretrained
-        reload_te = False
     if te3 == 'fp16':
         text_encoder_3 = transformers.T5EncoderModel.from_pretrained(
             repo_id,
@@ -47,6 +70,7 @@ def load_sd3(te3=None, fn=None, cache_dir=None, config=None):
             text_encoder_3=text_encoder_3,
             cache_dir=cache_dir,
             config=config,
+            **kwargs,
         )
     elif te3 == 'fp8':
         quantization_config = transformers.BitsAndBytesConfig(load_in_8bit=True)
@@ -64,6 +88,7 @@ def load_sd3(te3=None, fn=None, cache_dir=None, config=None):
             torch_dtype=dtype,
             cache_dir=cache_dir,
             config=config,
+            **kwargs,
         )
     else:
         pipe = loader(
@@ -72,20 +97,9 @@ def load_sd3(te3=None, fn=None, cache_dir=None, config=None):
             text_encoder_3=None,
             cache_dir=cache_dir,
             config=config,
+            **kwargs,
         )
-    if reload_te:
-        pipe.text_encoder = transformers.CLIPTextModelWithProjection.from_pretrained(
-            repo_id,
-            subfolder='text_encoder',
-            cache_dir=cache_dir,
-            torch_dtype=pipe.vae.dtype,
-        )
-        pipe.text_encoder_2 = transformers.CLIPTextModelWithProjection.from_pretrained(
-            repo_id,
-            subfolder='text_encoder_2',
-            cache_dir=cache_dir,
-            torch_dtype=pipe.vae.dtype,
-        )
+    diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING["StableDiffusion3Img2ImgPipeline"] = diffusers.StableDiffusion3Img2ImgPipeline
     return pipe
 
 
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 5a5fcc6ea..7f8d4ed3b 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -590,7 +590,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True):
                     if not shared.native:
                         warn(f'Model detected as Segmind Vega model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion XL'
-                elif size > 5692 and size < 5698:
+                elif size > 5692 and size < 5698 or size > 4134 and size < 4138:
                     if not shared.native:
                         warn(f'Model detected as Stable Diffusion 3 model, but attempting to load using backend=original: {op}={f} size={size} MB')
                     guess = 'Stable Diffusion 3'