Upload handler.py

handler.py

@@ -4,6 +4,7 @@ from typing import Any, Dict
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
 from PIL import Image
 import torch
+from torchao.quantization import quantize_, autoquant
 
 IS_COMPILE = True
 
@@ -11,33 +12,53 @@ if IS_COMPILE:
     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:
+def load_pipeline_stable(repo_id: str, dtype: torch.dtype) -> Any:
+    quantization_config = TorchAoConfig("int8dq")
+    vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
+    pipe = FluxPipeline.from_pretrained(repo_id, vae=vae, torch_dtype=dtype, quantization_config=quantization_config)
+    pipe.transformer.fuse_qkv_projections()
+    pipe.vae.fuse_qkv_projections()
+    pipe.to("cuda")
+    return pipe
+
+def load_pipeline_compile(repo_id: str, dtype: torch.dtype) -> Any:
+    quantization_config = TorchAoConfig("int8dq")
+    vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
+    pipe = FluxPipeline.from_pretrained(repo_id, vae=vae, torch_dtype=dtype, quantization_config=quantization_config)
+    pipe.transformer.fuse_qkv_projections()
+    pipe.vae.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")
+    pipe.to("cuda")
+    return pipe
+
+def load_pipeline_autoquant(repo_id: str, dtype: torch.dtype) -> Any:
+    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype).to("cuda")
+    pipe.transformer.fuse_qkv_projections()
+    pipe.vae.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.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
     pipe.vae.to(memory_format=torch.channels_last)
-    #pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=False, dynamic=False, backend="inductor")
     pipe.vae = torch.compile(pipe.vae, mode="max-autotune", fullgraph=True)
+    pipe.transformer = autoquant(pipe.transformer, error_on_unseen=False)
+    pipe.vae = autoquant(pipe.vae, error_on_unseen=False)
+    pipe.to("cuda")
     return pipe
 
 class EndpointHandler:
     def __init__(self, path=""):
         repo_id = "camenduru/FLUX.1-dev-diffusers"
-        #repo_id = "NoMoreCopyright/FLUX.1-dev-test"
         dtype = torch.bfloat16
-        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")
+        self.pipeline = load_pipeline_autoquant(repo_id, dtype)
+        #if IS_COMPILE: self.pipeline = load_pipeline_compile(repo_id, dtype)
+        #else: self.pipeline = load_pipeline_stable(repo_id, dtype)
 
     def __call__(self, data: Dict[str, Any]) -> Image.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")