Upload folder using huggingface_hub

deploy.bat

@@ -0,0 +1 @@
+F:\Projects\UAI\UAIBrainServer\Code\Env\Scripts\huggingface-cli.exe upload API_SDXLLightning . .

handler.py

@@ -4,6 +4,10 @@ import sys
 rootDir = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(rootDir)
 from imageRequest import ImageRequest
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+import torch
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -12,7 +16,27 @@ class EndpointHandler:
         # self.model= load_model(path)
         self.pipe = None
         self.modelName = ""
-        os.listdir(path)
+        baseReq = ImageRequest()
+        baseReq.model = "SG161222/RealVisXL_V4.0"
+        self.LoadModel(baseReq)
+        
+    def LoadModel(self, request):
+        base = "stabilityai/stable-diffusion-xl-base-1.0"
+        repo = "ByteDance/SDXL-Lightning"
+        ckpt = "sdxl_lightning_8step_unet.safetensors" # Use the correct ckpt for your step setting!
+        if request.model == "default":
+            request.model = base
+        else:
+            base = request.model
+        if self.pipe is None or self.modelName != request.model:
+            # Load model.
+            unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+            unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+            pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+
+            # Ensure sampler uses "trailing" timesteps.
+            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+            self.pipe = pipe
         
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -48,26 +72,9 @@ class EndpointHandler:
         Run SDXL Lightning pipeline
         """
         import torch
-        from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
-        from huggingface_hub import hf_hub_download
-        from safetensors.torch import load_file
-
-        base = "stabilityai/stable-diffusion-xl-base-1.0"
-        repo = "ByteDance/SDXL-Lightning"
-        ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
-        if request.model == "default":
-            request.model = base
-        else:
-            base = request.model
-        if self.pipe is None or self.modelName != request.model:
-            # Load model.
-            unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
-            unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
-            pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+        
 
-            # Ensure sampler uses "trailing" timesteps.
-            pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-            self.pipe = pipe
+        
 
         # Ensure using the same inference steps as the loaded model and CFG set to 0.
         images = pipe(request.prompt, negative_prompt = request.negative_prompt, num_inference_steps=request.steps, guidance_scale=0).images