Upload handler.py

handler.py

@@ -4,10 +4,8 @@ from typing import Any, Dict
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
 from PIL import Image
 import torch
-from accelerate import PartialState
-distributed_state = PartialState()
 
-IS_COMPILE = False
+IS_COMPILE = True
 
 if IS_COMPILE:
     import torch._dynamo
@@ -16,8 +14,10 @@ if IS_COMPILE:
 #from huggingface_inference_toolkit.logging import logger
 
 def compile_pipeline(pipe) -> Any:
+    pipe.transformer.to(memory_format=torch.channels_last)
     pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=False, dynamic=False, backend="inductor")
-    pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=False, dynamic=False, backend="inductor")
+    #pipe.vae.to(memory_format=torch.channels_last)
+    #pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=False, dynamic=False, backend="inductor")
     return pipe
 
 class EndpointHandler:
@@ -30,9 +30,7 @@ class EndpointHandler:
         #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)
         self.pipeline.transformer.fuse_qkv_projections()
-        self.pipeline.transformer.to(memory_format=torch.channels_last)
         self.pipeline.vae.fuse_qkv_projections()
-        self.pipeline.vae.to(memory_format=torch.channels_last)
         if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
         self.pipeline.to(distributed_state.device)