Update app.py

GEMINI IS A GOD IT KNOWS HOW TO LOAD A KAGGLE INTERFACE (Lulz.)

app.py

@@ -14,6 +14,7 @@ from huggingface_hub.utils import validate_repo_id, HFValidationError
 from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
 from huggingface_hub.utils import HfHubHTTPError
 
+
 # ---------------------- UTILITY FUNCTIONS ----------------------
 # (download_model, create_model_repo, etc. - All unchanged, but included for completeness)
 
@@ -97,13 +98,13 @@ def load_sdxl_checkpoint(checkpoint_path):
     unet_state = OrderedDict()
 
     for key, value in state_dict.items():
-        if key.startswith("first_stage_model."):
+        if key.startswith("first_stage_model."):  # VAE
             vae_state[key.replace("first_stage_model.", "")] = value.to(torch.float16)
-        elif key.startswith("condition_model.model.text_encoder."):
+        elif key.startswith("condition_model.model.text_encoder."):  # First Text Encoder
             text_encoder1_state[key.replace("condition_model.model.text_encoder.", "")] = value.to(torch.float16)
-        elif key.startswith("condition_model.model.text_encoder_2."):
+        elif key.startswith("condition_model.model.text_encoder_2."):  # Second Text Encoder
             text_encoder2_state[key.replace("condition_model.model.text_encoder_2.", "")] = value.to(torch.float16)
-        elif key.startswith("model.diffusion_model."):
+        elif key.startswith("model.diffusion_model."):  # UNet
             unet_state[key.replace("model.diffusion_model.", "")] = value.to(torch.float16)
 
     return text_encoder1_state, text_encoder2_state, vae_state, unet_state
@@ -115,16 +116,22 @@ def build_diffusers_model(text_encoder1_state, text_encoder2_state, vae_state, u
     if not reference_model_path:
         reference_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
 
-    config_text_encoder1 = CLIPTextConfig.from_pretrained(reference_model_path, subfolder="text_encoder")
-    config_text_encoder2 = CLIPTextConfig.from_pretrained(reference_model_path, subfolder="text_encoder_2")
-    config_vae = AutoencoderKL.from_pretrained(reference_model_path, subfolder="vae").config
-    config_unet = UNet2DConditionModel.from_pretrained(reference_model_path, subfolder="unet").config
+    # Load configurations from the reference model
+    config_text_encoder1 = CLIPTextConfig.from_pretrained(
+        reference_model_path, subfolder="text_encoder"
+    )
+    config_text_encoder2 = CLIPTextConfig.from_pretrained(
+       reference_model_path, subfolder="text_encoder_2"
+    )
+    # Use from_pretrained with subfolder for VAE and UNet
+    vae = AutoencoderKL.from_pretrained(reference_model_path, subfolder="vae")  # Corrected
+    unet = UNet2DConditionModel.from_pretrained(reference_model_path, subfolder="unet")  # Corrected
 
+    # Create instances using the configurations
     text_encoder1 = CLIPTextModel(config_text_encoder1)
-    text_encoder2 = CLIPTextModelWithProjection(config_text_encoder2)  # Correct class
-    vae = AutoencoderKL(config=config_vae)
-    unet = UNet2DConditionModel(config=config_unet)
+    text_encoder2 = CLIPTextModelWithProjection(config_text_encoder2)  # Use CLIPTextModelWithProjection
 
+    # Load state dicts with strict=False
     text_encoder1.load_state_dict(text_encoder1_state, strict=False)
     text_encoder2.load_state_dict(text_encoder2_state, strict=False)
     vae.load_state_dict(vae_state, strict=False)
@@ -135,6 +142,7 @@ def build_diffusers_model(text_encoder1_state, text_encoder2_state, vae_state, u
     vae.to(torch.float16).to("cpu")
     unet.to(torch.float16).to("cpu")
 
+
     return text_encoder1, text_encoder2, vae, unet
 
 def convert_and_save_sdxl_to_diffusers(checkpoint_path_or_url, output_path, reference_model_path):