testing out caching

app.py

@@ -19,12 +19,14 @@ from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
-    AutoencoderKL,
+    # AutoencoderKL,
 )
+from huggingface_hub import cached_download, hf_hub_url
 from controlnet_aux_local import NormalBaeDetector
 
 MAX_SEED = np.iinfo(np.int32).max
 API_KEY = os.environ.get("API_KEY", None)
+os.environ['HF_HOME'] = '/data/.huggingface'
 
 print("CUDA version:", torch.version.cuda)
 print("loading everything")
@@ -37,34 +39,75 @@ class Preprocessor:
         self.model = None
         self.name = ""
 
+    # def load(self, name: str) -> None:
+    #     if name == self.name:
+    #         return
+    #     elif name == "NormalBae":
+    #         print("Loading NormalBae")
+    #         self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+    #         torch.cuda.empty_cache()
+    #         self.name = name
+    #     else:
+    #         raise ValueError
+    #     return
+
     def load(self, name: str) -> None:
         if name == self.name:
             return
         elif name == "NormalBae":
             print("Loading NormalBae")
-            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to("cuda")
+            model_file = cached_download(hf_hub_url(self.MODEL_ID, filename="NormalBaeDetector.pth"))
+            self.model = NormalBaeDetector.from_pretrained(model_file).to("cuda")
             torch.cuda.empty_cache()
             self.name = name
         else:
             raise ValueError
         return
-
+    
     def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
         return self.model(image, **kwargs)
 
 torch.cuda.max_memory_allocated(device="cuda")
+
+# # Controlnet Normal
+# model_id = "lllyasviel/control_v11p_sd15_normalbae"
+# print("initializing controlnet")
+# controlnet = ControlNetModel.from_pretrained(
+#     model_id,
+#     torch_dtype=torch.float16,
+#     attn_implementation="flash_attention_2",
+# ).to("cuda")
+
 # Controlnet Normal
 model_id = "lllyasviel/control_v11p_sd15_normalbae"
 print("initializing controlnet")
+controlnet_file = cached_download(hf_hub_url(model_id, filename="diffusion_pytorch_model.safetensors"))
 controlnet = ControlNetModel.from_pretrained(
-    model_id,
+    controlnet_file,
     torch_dtype=torch.float16,
     attn_implementation="flash_attention_2",
 ).to("cuda")
 
+# # Scheduler
+# scheduler = DPMSolverMultistepScheduler.from_pretrained(
+#     "runwayml/stable-diffusion-v1-5",
+#     solver_order=2,
+#     subfolder="scheduler",
+#     use_karras_sigmas=True,
+#     final_sigmas_type="sigma_min",
+#     algorithm_type="sde-dpmsolver++",
+#     prediction_type="epsilon",
+#     thresholding=False,
+#     denoise_final=True,
+#     device_map="cuda",
+#     torch_dtype=torch.float16,
+# )
+
 # Scheduler
+scheduler_repo = "runwayml/stable-diffusion-v1-5"
+scheduler_file = cached_download(hf_hub_url(scheduler_repo, filename="scheduler/scheduler_config.json", subfolder="scheduler"))
 scheduler = DPMSolverMultistepScheduler.from_pretrained(
-    "runwayml/stable-diffusion-v1-5",
+    scheduler_file,
     solver_order=2,
     subfolder="scheduler",
     use_karras_sigmas=True,
@@ -86,18 +129,30 @@ vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/v
 # vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
 # vae.to(memory_format=torch.channels_last) 
 
+# print('loading pipe')
+# pipe = StableDiffusionControlNetPipeline.from_single_file(
+#     base_model_url,
+#     safety_checker=None,
+#     # load_safety_checker=True,
+#     controlnet=controlnet,
+#     scheduler=scheduler,
+#     # vae=vae,
+#     torch_dtype=torch.float16,
+# )
+
+# Stable Diffusion Pipeline
+base_model_repo = "Lykon/AbsoluteReality"
+base_model_file = cached_download(hf_hub_url(base_model_repo, filename="AbsoluteReality_1.8.1_pruned.safetensors"))
+
 print('loading pipe')
 pipe = StableDiffusionControlNetPipeline.from_single_file(
-    base_model_url,
+    base_model_file,
     safety_checker=None,
-    # load_safety_checker=True,
     controlnet=controlnet,
     scheduler=scheduler,
-    # vae=vae,
     torch_dtype=torch.float16,
 )
 
-
 print("loading preprocessor")
 preprocessor = Preprocessor()
 preprocessor.load("NormalBae")