Upload handler.py

handler.py

@@ -5,8 +5,11 @@ from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, Torch
 from PIL import Image
 import torch
 
-#import torch._dynamo
-#torch._dynamo.config.suppress_errors = True
+IS_COMPILE = True
+
+if IS_COMPILE:
+    import torch._dynamo
+    torch._dynamo.config.suppress_errors = True
 
 #from huggingface_inference_toolkit.logging import logger
 
@@ -18,7 +21,6 @@ def compile_pipeline(pipe) -> Any:
 
 class EndpointHandler:
     def __init__(self, path=""):
-        is_compile = False
         repo_id = "camenduru/FLUX.1-dev-diffusers"
         #repo_id = "NoMoreCopyright/FLUX.1-dev-test"
         dtype = torch.bfloat16
@@ -26,7 +28,7 @@ class EndpointHandler:
         vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
         #transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype, quantization_config=quantization_config).to("cuda")
         self.pipeline = FluxPipeline.from_pretrained(repo_id, vae=vae, torch_dtype=dtype, quantization_config=quantization_config)
-        if is_compile: self.pipeline = compile_pipeline(self.pipeline)
+        if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
         self.pipeline.to("cuda")
 
     #@torch.inference_mode()
@@ -45,9 +47,9 @@ class EndpointHandler:
 
         parameters = data.pop("parameters", {})
 
-        num_inference_steps = parameters.get("num_inference_steps", 30)
+        num_inference_steps = parameters.get("num_inference_steps", 28)
         width = parameters.get("width", 1024)
-        height = parameters.get("height", 768)
+        height = parameters.get("height", 1024)
         guidance_scale = parameters.get("guidance_scale", 3.5)
 
         # seed generator (seed cannot be provided as is but via a generator)