revert

app.py

@@ -19,9 +19,8 @@ from diffusers import (
     ControlNetModel,
     DPMSolverMultistepScheduler,
     StableDiffusionControlNetPipeline,
-    # AutoencoderKL,
+    AutoencoderKL,
 )
-from huggingface_hub import hf_hub_download, hf_hub_url
 from controlnet_aux_local import NormalBaeDetector
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -39,31 +38,18 @@ 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")
-            model_file = hf_hub_download(self.MODEL_ID, filename="scannet.pt")
-            self.model = NormalBaeDetector.from_pretrained(model_file).to("cuda")
+            self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).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)
 
@@ -78,26 +64,9 @@ controlnet = ControlNetModel.from_pretrained(
     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 = hf_hub_download(scheduler_repo, filename="scheduler/scheduler_config.json", subfolder="scheduler")
 scheduler = DPMSolverMultistepScheduler.from_pretrained(
-    scheduler_file,
+    "runwayml/stable-diffusion-v1-5",
     solver_order=2,
     subfolder="scheduler",
     use_karras_sigmas=True,
@@ -110,52 +79,50 @@ scheduler = DPMSolverMultistepScheduler.from_pretrained(
     torch_dtype=torch.float16,
 )
 
+
 # Stable Diffusion Pipeline URL
 # base_model_url = "https://huggingface.co/broyang/hentaidigitalart_v20/blob/main/realcartoon3d_v15.safetensors"
 base_model_url = "https://huggingface.co/Lykon/AbsoluteReality/blob/main/AbsoluteReality_1.8.1_pruned.safetensors"
 vae_url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
 
 # print('loading vae')
-# 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,
-# )
+vae = AutoencoderKL.from_single_file(vae_url, torch_dtype=torch.float16).to("cuda")
+vae.to(memory_format=torch.channels_last) 
 
-# Stable Diffusion Pipeline
-base_model_repo = "Lykon/AbsoluteReality"
-base_model_file = hf_hub_download(base_model_repo, filename="AbsoluteReality_1.8.1_pruned.safetensors")
+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,
+)
 
 print('loading pipe')
 pipe = StableDiffusionControlNetPipeline.from_single_file(
-    base_model_file,
+    base_model_url,
     safety_checker=None,
     controlnet=controlnet,
     scheduler=scheduler,
+    vae=vae,
     torch_dtype=torch.float16,
 )
 
 print("loading preprocessor")
 preprocessor = Preprocessor()
 preprocessor.load("NormalBae")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",)
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
-# pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="EasyNegativeV2.safetensors", token="EasyNegativeV2",)
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="badhandv4.pt", token="badhandv4")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="fcNeg-neg.pt", token="fcNeg-neg")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Ahegao.pt", token="HDA_Ahegao")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Bondage.pt", token="HDA_Bondage")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_pet_play.pt", token="HDA_pet_play")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_unconventional maid.pt", token="HDA_unconventional_maid")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NakedHoodie.pt", token="HDA_NakedHoodie")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_NunDress.pt", token="HDA_NunDress")
+pipe.load_textual_inversion("broyang/hentaidigitalart_v20", weight_name="HDA_Shibari.pt", token="HDA_Shibari")
 pipe.to("cuda")