Upload 3 files

handler.py

@@ -14,9 +14,10 @@ if IS_COMPILE:
 #from huggingface_inference_toolkit.logging import logger
 
 def compile_pipeline(pipe) -> Any:
-    pipe.transformer.fuse_qkv_projections()
     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.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:
@@ -24,14 +25,15 @@ class EndpointHandler:
         repo_id = "camenduru/FLUX.1-dev-diffusers"
         #repo_id = "NoMoreCopyright/FLUX.1-dev-test"
         dtype = torch.bfloat16
-        quantization_config = TorchAoConfig("int8wo")
+        quantization_config = TorchAoConfig("int8dq")
         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)
+        self.pipeline.transformer.fuse_qkv_projections()
+        self.pipeline.vae.fuse_qkv_projections()
         if IS_COMPILE: self.pipeline = compile_pipeline(self.pipeline)
         self.pipeline.to("cuda")
 
-    #@torch.inference_mode()
     def __call__(self, data: Dict[str, Any]) -> Image.Image:
         #logger.info(f"Received incoming request with {data=}")
 
@@ -63,4 +65,7 @@ class EndpointHandler:
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             generator=generator,
-        ).images[0]
+            output_type="pil",
+        ).images[0]
+
+

requirements.txt

@@ -1,11 +1,15 @@
 huggingface_hub
-torch
+torch==2.4.0
 torchvision
-torchao
+torchao==0.9.0
 diffusers
 peft
 accelerate
 transformers
 numpy
 scipy
-Pillow
+Pillow
+sentencepiece
+protobuf
+pytorch-lightning
+triton