From 4990a8a375d0078de246341ba8d707ecf4e3043d Mon Sep 17 00:00:00 2001
From: GiusTex <112352961+GiusTex@users.noreply.github.com>
Date: Sun, 17 Nov 2024 00:34:44 +0100
Subject: [PATCH] fix "variant=fp16" error

---
 utils.py | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/utils.py b/utils.py
index 42271eb..8eebb2a 100644
--- a/utils.py
+++ b/utils.py
@@ -12,6 +12,7 @@ from .pipeline_fill_sd_xl import StableDiffusionXLFillPipeline
 from diffusers import AutoencoderKL, TCDScheduler
 from diffusers.models.model_loading_utils import load_state_dict
 from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
+from diffusers import UNet2DConditionModel
 
 
 def get_first_folder_list(folder_name: str) -> tuple[list[str], dict[str, float], float]:
@@ -128,21 +129,26 @@ def loadVaeModel(vae_path, device, dtype, enable_vae_slicing, enable_vae_tiling)
     return vae
 
 
+def loadUnetModel(model_path, device, dtype):
+    unet = UNet2DConditionModel.from_pretrained(model_path, subfolder="unet", use_safetensors=True)
+    unet.to(device, dtype)
+    return unet
+
+
 def diffuserOutpaintSamples(model_path, controlnet_model, diffuser_outpaint_cnet_image, dtype, controlnet_path, 
                             prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds, 
                             device, steps, controlnet_strength, guidance_scale, 
                             keep_model_device):
     
     controlnet_model = loadControlnetModel(device, dtype, controlnet_path)
-    
+    unet = loadUnetModel(model_path, device, dtype)
+                                
     with open(f"{model_path}/scheduler/scheduler_config.json", "r") as f:
         scheduler_config = json.load(f)
     scheduler = TCDScheduler.from_config(scheduler_config)
     
-    pipe = StableDiffusionXLFillPipeline.from_pretrained(
-        model_path,
-        torch_dtype=dtype,
-        variant="fp16",
+    pipe = StableDiffusionXLFillPipeline(
+        unet,
         scheduler=scheduler,
     )
     if not keep_model_device: