Upload handler.py

handler.py

@@ -4,9 +4,11 @@ from typing import Any, Dict
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
 from PIL import Image
 import torch
-from torchao.quantization import quantize_, autoquant
+from torchao.quantization import quantize_, autoquant, int8_dynamic_activation_int8_weight
+from huggingface_hub import hf_hub_download
 
-IS_COMPILE = True
+IS_COMPILE = False
+IS_TURBO = True
 
 if IS_COMPILE:
     import torch._dynamo
@@ -49,13 +51,44 @@ def load_pipeline_autoquant(repo_id: str, dtype: torch.dtype) -> Any:
     pipe.to("cuda")
     return pipe
 
+def load_pipeline_turbo(repo_id: str, dtype: torch.dtype) -> Any:
+    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype).to("cuda")
+    pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd")
+    pipe.set_adapters(["hyper-sd"], adapter_weights=[0.125])
+    pipe.fuse_lora()
+    pipe.transformer.fuse_qkv_projections()
+    pipe.vae.fuse_qkv_projections()
+    quantize_(pipe.transformer, int8_dynamic_activation_int8_weight(), device="cuda")
+    quantize_(pipe.vae, int8_dynamic_activation_int8_weight(), device="cuda")
+    pipe.to("cuda")
+    return pipe
+
+def load_pipeline_turbo_compile(repo_id: str, dtype: torch.dtype) -> Any:
+    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype).to("cuda")
+    pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd")
+    pipe.set_adapters(["hyper-sd"], adapter_weights=[0.125])
+    pipe.fuse_lora()
+    pipe.transformer.fuse_qkv_projections()
+    pipe.vae.fuse_qkv_projections()
+    quantize_(pipe.transformer, int8_dynamic_activation_int8_weight(), device="cuda")
+    quantize_(pipe.vae, int8_dynamic_activation_int8_weight(), device="cuda")
+    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
+
 class EndpointHandler:
     def __init__(self, path=""):
         repo_id = "camenduru/FLUX.1-dev-diffusers"
         dtype = torch.bfloat16
-        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)
+        if IS_COMPILE:
+            if IS_TURBO: self.pipeline = load_pipeline_turbo_compile(repo_id, dtype)
+            else: self.pipeline = load_pipeline_compile(repo_id, dtype)
+        else:
+            if IS_TURBO: self.pipeline = load_pipeline_turbo(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=}")
@@ -72,7 +105,7 @@ class EndpointHandler:
 
         parameters = data.pop("parameters", {})
 
-        num_inference_steps = parameters.get("num_inference_steps", 28)
+        num_inference_steps = parameters.get("num_inference_steps", 8 if IS_TURBO else 28)
         width = parameters.get("width", 1024)
         height = parameters.get("height", 1024)
         guidance_scale = parameters.get("guidance_scale", 3.5)