Upload handler.py

handler.py

@@ -8,6 +8,7 @@ from PIL import Image
 import torch
 from torchao.quantization import quantize_, autoquant, int8_dynamic_activation_int8_weight, int8_dynamic_activation_int4_weight
 from huggingface_hub import hf_hub_download
+import gc
 
 IS_COMPILE = False
 IS_TURBO = True
@@ -59,6 +60,7 @@ def load_pipeline_turbo(repo_id: str, dtype: torch.dtype) -> Any:
     pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd")
     pipe.set_adapters(["hyper-sd"], adapter_weights=[0.125])
     pipe.fuse_lora()
+    pipe.unload_lora_weights()
     pipe.transformer.fuse_qkv_projections()
     pipe.vae.fuse_qkv_projections()
     weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
@@ -73,6 +75,7 @@ def load_pipeline_turbo_compile(repo_id: str, dtype: torch.dtype) -> Any:
     pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd")
     pipe.set_adapters(["hyper-sd"], adapter_weights=[0.125])
     pipe.fuse_lora()
+    pipe.unload_lora_weights()
     pipe.transformer.fuse_qkv_projections()
     pipe.vae.fuse_qkv_projections()
     weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
@@ -88,11 +91,13 @@ def load_pipeline_turbo_compile(repo_id: str, dtype: torch.dtype) -> Any:
 
 class EndpointHandler:
     def __init__(self, path=""):
-        repo_id = "NoMoreCopyrightOrg/flux-dev-8step" if IS_TURBO else "camenduru/FLUX.1-dev-diffusers"
+        repo_id = "NoMoreCopyrightOrg/flux-dev-8step" if IS_TURBO else "NoMoreCopyrightOrg/flux-dev"
         #dtype = torch.bfloat16
         dtype = torch.float16 # for older nVidia GPUs
         if IS_COMPILE: load_pipeline_compile(repo_id, dtype)
         else: self.pipeline = load_pipeline_stable(repo_id, dtype)
+        gc.collect()
+        torch.cuda.empty_cache()
 
     def __call__(self, data: Dict[str, Any]) -> Image.Image:
         logger.info(f"Received incoming request with {data=}")