Upload 2 files

handler.py

@@ -1,27 +1,35 @@
 import os
 from typing import Any, Dict
 
-from diffusers import FluxPipeline, FluxTransformer2DModel
-from torchao.quantization import int8_weight_only, quantize_
+from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
 from PIL.Image import Image
 import torch
 
-from huggingface_inference_toolkit.logging import logger
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
+
+#from huggingface_inference_toolkit.logging import logger
+
+def compile_pipeline(pipe):
+    pipe.transformer.to(memory_format=torch.channels_last)
+    pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead", fullgraph=False, dynamic=False, backend="inductor")
+    return pipe
 
 class EndpointHandler:
     def __init__(self, **kwargs: Any) -> None:  # type: ignore
+        is_compile = False
         repo_id = "camenduru/FLUX.1-dev-diffusers"
         dtype = torch.bfloat16
-        transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype)
-        quantize_(transformer, int8_weight_only(), device="cuda")
-        transformer.to(memory_format=torch.channels_last)
-        transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
-        self.pipeline = FluxPipeline.from_pretrained(repo_id, transformer=transformer, torch_dtype=torch.bfloat16).to("cuda")
-        self.pipeline.vae.to(memory_format=torch.channels_last)
-        self.pipeline.vae.decode = torch.compile(self.pipeline.vae.decode, mode="max-autotune", fullgraph=True)
-
+        quantization_config = TorchAoConfig("int4dq")
+        vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
+        #transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype).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)
+        self.pipeline.to("cuda")
+
+    @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")

requirements.txt

@@ -1,7 +1,10 @@
 torch
+torchvision
 diffusers
 peft
 accelerate
 transformers
 numpy
-Pillow
+scipy
+Pillow
+triton