Improve SDXL auth error logging

app/colorize_model.py

@@ -156,41 +156,42 @@ class ColorizeModel:
 
     def _load_pipeline(self) -> None:
         controlnet_path = self._download_controlnet()
+        base_kwargs = {"use_auth_token": self.hf_token} if self.hf_token else {}
 
         logger.info("Loading SDXL components...")
-        vae = AutoencoderKL.from_pretrained(
-            self.BASE_MODEL,
-            subfolder="vae",
-            torch_dtype=self.dtype,
-            token=self.hf_token,
-        )
+        vae = AutoencoderKL.from_pretrained(self.BASE_MODEL, subfolder="vae", torch_dtype=self.dtype, token=self.hf_token)
         unet = UNet2DConditionModel.from_config(
             self.BASE_MODEL,
             subfolder="unet",
-            token=self.hf_token,
+            token=self.hf_token if self.hf_token else None,
         )
         lightning_path = hf_hub_download(
             repo_id=self.LIGHTNING_REPO,
             filename=self.LIGHTNING_WEIGHTS,
-            token=self.hf_token,
+            token=self.hf_token if self.hf_token else None,
         )
         unet.load_state_dict(load_file(lightning_path))
 
-        controlnet = ControlNetModel.from_pretrained(
-            controlnet_path,
-            torch_dtype=self.dtype,
-        )
+        controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=self.dtype)
 
-        self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
-            self.BASE_MODEL,
-            vae=vae,
-            unet=unet,
-            controlnet=controlnet,
-            torch_dtype=self.dtype,
-            safety_checker=None,
-            requires_safety_checker=False,
-            token=self.hf_token,
-        )
+        try:
+            self.pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+                self.BASE_MODEL,
+                vae=vae,
+                unet=unet,
+                controlnet=controlnet,
+                torch_dtype=self.dtype,
+                safety_checker=None,
+                requires_safety_checker=False,
+                token=self.hf_token if self.hf_token else None,
+            )
+        except Exception as exc:
+            logger.error("Failed to load base SDXL model: %s", exc)
+            logger.error(
+                "Ensure the account associated with HUGGINGFACE_HUB_TOKEN has accepted "
+                "the license for %s and that the token has access.", self.BASE_MODEL
+            )
+            raise
         self.pipe.set_progress_bar_config(disable=True)
 
         if self.device.type == "cuda":