Upload handler.py

handler.py

@@ -5,10 +5,10 @@ from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, Torch
 from PIL.Image import Image
 import torch
 
-#import torch._dynamo
-#torch._dynamo.config.suppress_errors = True
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
 
-#from huggingface_inference_toolkit.logging import logger
+from huggingface_inference_toolkit.logging import logger
 
 def compile_pipeline(pipe) -> Any:
     pipe.transformer.fuse_qkv_projections()
@@ -17,10 +17,10 @@ def compile_pipeline(pipe) -> Any:
     return pipe
 
 class EndpointHandler:
-    def __init__(self, path="NoMoreCopyright/FLUX.1-dev-test"):
-        is_compile = False
-        repo_id = "camenduru/FLUX.1-dev-diffusers"
-        #repo_id = "NoMoreCopyright/FLUX.1-dev-test"
+    def __init__(self, path=""):
+        is_compile = True
+        #repo_id = "camenduru/FLUX.1-dev-diffusers"
+        repo_id = "NoMoreCopyright/FLUX.1-dev-test"
         dtype = torch.bfloat16
         quantization_config = TorchAoConfig("int4dq")
         vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
@@ -29,9 +29,9 @@ class EndpointHandler:
         if is_compile: self.pipeline = compile_pipeline(self.pipeline)
         self.pipeline.to("cuda")
 
-    #@torch.inference_mode()
+    @torch.inference_mode()
     def __call__(self, data: Dict[str, Any]) -> Image:
-        #logger.info(f"Received incoming request with {data=}")
+        logger.info(f"Received incoming request with {data=}")
 
         if "inputs" in data and isinstance(data["inputs"], str):
             prompt = data.pop("inputs")