Update handler.py

handler.py

@@ -4,21 +4,16 @@ from typing import Any, Dict
 from PIL import Image
 import torch
 from diffusers import FluxPipeline
-from huggingface_inference_toolkit.logging import logger
+from transformers import logger
 import time
-IS_TURBO=True
-class EndpointHandler:
-    def __init__(self, path=""):
-        #dtype = torch.float16 # for older nVidia GPUs
-
-        self.pipeline =FluxPipeline.from_pretrained("NoMoreCopyrightOrg/flux-dev-8step", torch_dtype=torch.bfloat16)
-
-
-        self.pipeline.enable_model_cpu_offload()  # save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
 
-        prompt = "A cat holding a sign that says hello world"
 
-        torch.cuda.empty_cache()
+class EndpointHandler:
+    def __init__(self):
+        self.pipe = FluxPipeline.from_pretrained(
+            "NoMoreCopyrightOrg/flux-dev",
+            torch_dtype=torch.bfloat16,
+        ).to("cuda")
 
     def __call__(self, data: Dict[str, Any]) -> Image.Image:
         logger.info(f"Received incoming request with {data=}")
@@ -35,7 +30,7 @@ class EndpointHandler:
 
         parameters = data.pop("parameters", {})
 
-        num_inference_steps = parameters.get("num_inference_steps", 8 if IS_TURBO else 28)
+        num_inference_steps = parameters.get("num_inference_steps", 28)
         width = parameters.get("width", 1024)
         height = parameters.get("height", 1024)
         guidance_scale = parameters.get("guidance_scale", 3.5)
@@ -43,8 +38,8 @@ class EndpointHandler:
         # seed generator (seed cannot be provided as is but via a generator)
         seed = parameters.get("seed", 0)
         generator = torch.manual_seed(seed)
-        start_time=time.time()
-        result = self.pipeline(  # type: ignore
+        start_time = time.time()
+        result = self.pipe(  # type: ignore
             prompt,
             height=height,
             width=width,
@@ -57,5 +52,3 @@ class EndpointHandler:
         time_taken = end_time - start_time
         print(f"Time taken: {time_taken:.2f} seconds")
         return result
-
-