Upload handler.py

handler.py

@@ -1,267 +1,52 @@
-# https://github.com/sayakpaul/diffusers-torchao
-# https://github.com/pytorch/ao/releases
-# https://developer.nvidia.com/cuda-gpus
-
 import os
 from typing import Any, Dict
-import gc
-import time
 from PIL import Image
-from huggingface_hub import hf_hub_download
 import torch
-from torchao.quantization import quantize_, autoquant, int8_dynamic_activation_int4_weight, float8_dynamic_activation_float8_weight, int8_dynamic_activation_int8_weight, int8_weight_only
-from torchao.quantization.quant_api import PerRow
-from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
-from transformers import T5EncoderModel
-from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+from diffusers import FluxPipeline
 from huggingface_inference_toolkit.logging import logger
-
-import subprocess
-subprocess.run("pip list", shell=True)
-
-print("Device Name:", torch.cuda.get_device_name())
-print("Device Capability:", torch.cuda.get_device_capability())
-
-IS_TURBO = False
-IS_4BIT = False
-IS_PARA = True
-IS_MGPU = False
-IS_LVRAM = False
-IS_COMPILE = True
-IS_WARM = True
-IS_QUANT = True
-IS_AUTOQ = False
-IS_CC90 = True if torch.cuda.get_device_capability() >= (9, 0) else False
-IS_CC89 = True if torch.cuda.get_device_capability() >= (8, 9) else False
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+from torchao.quantization import autoquant
+import time
 
 # Set high precision for float32 matrix multiplications.
 # This setting optimizes performance on NVIDIA GPUs with Ampere architecture (e.g., A100, RTX 30 series) or newer.
 torch.set_float32_matmul_precision("high")
 
-if IS_COMPILE:
-    import torch._dynamo
-    torch._dynamo.config.suppress_errors = False
-    #torch._dynamo.config.suppress_errors = True
-    torch._inductor.config.disable_progress = False
-    #torch._inductor.config.conv_1x1_as_mm = True
-    #torch._inductor.config.coordinate_descent_tuning = True
-    #torch._inductor.config.coordinate_descent_check_all_directions = True
-    #torch._inductor.config.epilogue_fusion = False
-
-if IS_MGPU:
-    import torch.distributed as dist
-    dist.init_process_group()
-    torch.cuda.set_device(dist.get_rank())
-
-def print_vram():
-    free = torch.cuda.mem_get_info()[0] / (1024 ** 3)
-    total = torch.cuda.mem_get_info()[1] / (1024 ** 3)
-    print(f"VRAM: {total - free:.2f}/{total:.2f}GB")
-
-def pil_to_base64(image: Image.Image, modelname: str, prompt: str, height: int, width: int, steps: int, cfg: float, seed: int) -> str:
-    import base64
-    from io import BytesIO
-    import json
-    from PIL import PngImagePlugin
-    metadata = {"prompt": prompt, "num_inference_steps": steps, "guidance_scale": cfg, "seed": seed, "resolution": f"{width} x {height}",
-                "Model": {"Model": modelname.split("/")[-1]}}
-    info = PngImagePlugin.PngInfo()
-    info.add_text("metadata", json.dumps(metadata))
-    buffered = BytesIO()
-    image.save(buffered, "PNG", pnginfo=info)
-    return base64.b64encode(buffered.getvalue()).decode('ascii')
-
-def load_te2(repo_id: str, dtype: torch.dtype) -> Any:
-    text_encoder_2 = T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", torch_dtype=dtype)
-    if IS_4BIT: quantize_(text_encoder_2,  int8_dynamic_activation_int4_weight())
-    else: quantize_(text_encoder_2, int8_dynamic_activation_int8_weight())
-    return text_encoder_2
-
-def load_pipeline_stable(repo_id: str, dtype: torch.dtype) -> Any:
-    quantization_config = TorchAoConfig("int4dq" if IS_4BIT else "float8dq" if IS_CC90 else "int8wo")
-    vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
-    pipe = FluxPipeline.from_pretrained(repo_id, vae=vae, text_encoder_2=load_te2(repo_id, dtype), torch_dtype=dtype, quantization_config=quantization_config)
-    pipe.transformer.fuse_qkv_projections()
-    pipe.vae.fuse_qkv_projections()
-    return pipe
-
-def load_pipeline_lowvram(repo_id: str, dtype: torch.dtype) -> Any:
-    int4_config = TorchAoConfig("int4dq")
-    float8_config = TorchAoConfig("float8dq")
-    vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
-    transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype, quantization_config=float8_config)
-    pipe = FluxPipeline.from_pretrained(repo_id, vae=None, transformer=None, text_encoder_2=load_te2(repo_id, dtype), torch_dtype=dtype, quantization_config=int4_config)
-    pipe.transformer = transformer
-    pipe.vae = vae
-    #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("int4dq" if IS_4BIT else "float8dq" if IS_CC90 else "int8wo")
-    vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
-    pipe = FluxPipeline.from_pretrained(repo_id, vae=vae, text_encoder_2=load_te2(repo_id, dtype), 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="max-autotune", fullgraph=True)
-    pipe.vae.to(memory_format=torch.channels_last)
-    pipe.vae = torch.compile(pipe.vae, mode="max-autotune", fullgraph=True)
-    return pipe
-
-def load_pipeline_autoquant(repo_id: str, dtype: torch.dtype) -> Any:
-    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype)
-    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="max-autotune", fullgraph=True)
-    pipe.vae.to(memory_format=torch.channels_last)
-    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)
-    return pipe
-
-def load_pipeline_turbo(repo_id: str, dtype: torch.dtype) -> Any:
-    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype)
-    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.unload_lora_weights()
-    pipe.transformer.fuse_qkv_projections()
-    pipe.vae.fuse_qkv_projections()
-    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
-    quantize_(pipe.transformer, weight, device="cuda")
-    quantize_(pipe.vae, weight, device="cuda")
-    return pipe
-
-def load_pipeline_turbo_compile(repo_id: str, dtype: torch.dtype) -> Any:
-    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype)
-    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.unload_lora_weights()
-    pipe.transformer.fuse_qkv_projections()
-    pipe.vae.fuse_qkv_projections()
-    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
-    quantize_(pipe.transformer, weight, device="cuda")
-    quantize_(pipe.vae, weight, device="cuda")
-    pipe.transformer.to(memory_format=torch.channels_last)
-    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="max-autotune", fullgraph=True)
-    return pipe
-
-def load_pipeline_opt(repo_id: str, dtype: torch.dtype) -> Any:
-    quantization_config = TorchAoConfig("int4dq" if IS_4BIT else "float8dq" if IS_CC90 else "int8wo")
-    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
-    transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype)
-    #transformer.fuse_qkv_projections()
-    if IS_CC90: quantize_(transformer, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
-    elif IS_CC89: quantize_(transformer, float8_dynamic_activation_float8_weight(), device="cuda")
-    else: quantize_(transformer, weight, device="cuda")
-    transformer.to(memory_format=torch.channels_last)
-    #transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)
-    transformer = torch.compile(transformer, mode="max-autotune-no-cudagraphs")
-    vae = AutoencoderKL.from_pretrained(repo_id, subfolder="vae", torch_dtype=dtype)
-    #vae.fuse_qkv_projections()
-    if IS_CC90: quantize_(vae, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
-    elif IS_CC89: quantize_(vae, float8_dynamic_activation_float8_weight(), device="cuda")
-    else: quantize_(vae, weight, device="cuda")
-    vae.to(memory_format=torch.channels_last)
-    #vae = torch.compile(vae, mode="max-autotune", fullgraph=True)
-    vae = torch.compile(vae, mode="max-autotune-no-cudagraphs")
-    pipe = FluxPipeline.from_pretrained(repo_id, transformer=None, vae=None, torch_dtype=dtype, quantization_config=quantization_config)
-    quantize_(pipe.text_encoder, int8_dynamic_activation_int8_weight())
-    quantize_(pipe.text_encoder_2, int8_dynamic_activation_int8_weight())
-    pipe.transformer = transformer
-    pipe.vae = vae
-    return pipe
-
-def load_pipeline_para(repo_id: str, dtype: torch.dtype) -> Any:
-    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
-    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype).to("cuda")
-    if IS_PARA:
-        if IS_MGPU:
-            from para_attn.context_parallel import init_context_parallel_mesh
-            from para_attn.context_parallel.diffusers_adapters import parallelize_pipe
-            from para_attn.parallel_vae.diffusers_adapters import parallelize_vae
-            mesh = init_context_parallel_mesh(pipe.device.type, max_ring_dim_size=2)
-            parallelize_pipe(pipe, mesh=mesh)
-            parallelize_vae(pipe.vae, mesh=mesh._flatten())
-        apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
-    quantize_(pipe.text_encoder, int8_dynamic_activation_int8_weight())
-    quantize_(pipe.text_encoder_2, int8_dynamic_activation_int8_weight())
-    if IS_CC90: quantize_(pipe.transformer, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
-    elif IS_CC89: quantize_(pipe.transformer, float8_dynamic_activation_float8_weight(), device="cuda")
-    else: quantize_(pipe.transformer, weight, device="cuda")
-    #pipe.transformer.fuse_qkv_projections()
-    pipe.transformer.to(memory_format=torch.channels_last)
-    #pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True)
-    pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs")
-    if IS_CC90: quantize_(pipe.vae, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
-    elif IS_CC89: quantize_(pipe.vae, float8_dynamic_activation_float8_weight(), device="cuda")
-    else: quantize_(pipe.vae, weight, device="cuda")
-    #pipe.vae.fuse_qkv_projections()
-    pipe.vae.to(memory_format=torch.channels_last)
-    #pipe.vae = torch.compile(pipe.vae, mode="max-autotune", fullgraph=True)
-    pipe.vae = torch.compile(pipe.vae, mode="max-autotune-no-cudagraphs")
-    return pipe
-
-def load_pipeline_fast(repo_id: str, dtype: torch.dtype) -> Any:
-    pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype).to("cuda")
-    pipe.enable_vae_slicing()
-    pipe.enable_vae_tiling()
-    pipe.transformer.fuse_qkv_projections()
-    pipe.vae.fuse_qkv_projections()
-    pipe.transformer.to(memory_format=torch.channels_last)
-    pipe.vae.to(memory_format=torch.channels_last)
-    if IS_QUANT and not IS_AUTOQ:
-        quantize_(pipe.text_encoder, int8_weight_only())
-        quantize_(pipe.text_encoder_2, int8_weight_only())
-        if IS_CC90: quantize_(pipe.transformer, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
-        elif IS_CC89: quantize_(pipe.transformer, float8_dynamic_activation_float8_weight(), device="cuda")
-        else: quantize_(pipe.vae, int8_weight_only())
-        if IS_CC90: quantize_(pipe.vae, float8_dynamic_activation_float8_weight(granularity=PerRow()), device="cuda")
-        elif IS_CC89: quantize_(pipe.vae, float8_dynamic_activation_float8_weight(), device="cuda")
-        else: quantize_(pipe.vae, int8_weight_only())
-    return pipe
+import torch._dynamo
+torch._dynamo.config.suppress_errors = False # for debugging
 
 class EndpointHandler:
     def __init__(self, path=""):
-        repo_id = "NoMoreCopyrightOrg/flux-dev-8step" if IS_TURBO else "NoMoreCopyrightOrg/flux-dev"
-        self.repo_id = repo_id
-        dtype = torch.bfloat16
-        #dtype = torch.float16 # for older nVidia GPUs
-        print_vram()
-        print("Loading pipeline...")
-        if IS_AUTOQ: self.pipeline = load_pipeline_fast(repo_id, dtype)
-        elif IS_COMPILE: self.pipeline = load_pipeline_fast(repo_id, dtype)
-        elif IS_LVRAM and IS_CC89: self.pipeline = load_pipeline_lowvram(repo_id, dtype)
-        else: self.pipeline = load_pipeline_stable(repo_id, dtype)
-        self.pipeline.enable_vae_slicing()
-        self.pipeline.enable_vae_tiling()
-        self.pipeline.to("cuda")
-        if IS_PARA: apply_cache_on_pipe(self.pipeline, residual_diff_threshold=0.12)
-        if IS_COMPILE:
-            print("Compiling pipeline...")
-            self.pipeline.transformer = torch.compile(self.pipeline.transformer, mode="max-autotune-no-cudagraphs")
-            self.pipeline.vae = torch.compile(self.pipeline.vae, mode="max-autotune-no-cudagraphs")
-            if IS_AUTOQ:
-                print("Running autoquant...")
-                self.pipeline.transformer = autoquant(self.pipeline.transformer, error_on_unseen=False)
-                self.pipeline.vae = autoquant(self.pipeline.vae, error_on_unseen=False)
-        gc.collect()
-        torch.cuda.empty_cache()
-        print_vram()
-        if IS_WARM:
-            print("Warming pipeline...")
-            start = time.time()
-            self.pipeline("Hello world!", output_type="pil")
-            end = time.time()
-            print(f'Compiled in {end - start:.3f} sec.')
-
-    def __call__(self, data: Dict[str, Any]) -> str:
+        self.pipe = FluxPipeline.from_pretrained(
+            "NoMoreCopyrightOrg/flux-dev",
+            torch_dtype=torch.bfloat16,
+        ).to("cuda")
+        self.pipe.enable_vae_slicing()
+        self.pipe.enable_vae_tiling()
+        self.pipe.transformer.fuse_qkv_projections()
+        self.pipe.vae.fuse_qkv_projections()
+        self.pipe.transformer.to(memory_format=torch.channels_last)
+        self.pipe.vae.to(memory_format=torch.channels_last)
+        apply_cache_on_pipe(self.pipe, residual_diff_threshold=0.12)
+        self.pipe.transformer = torch.compile(
+            self.pipe.transformer, mode="max-autotune-no-cudagraphs",
+        )
+        self.pipe.vae = torch.compile(
+            self.pipe.vae, mode="max-autotune-no-cudagraphs",
+        )
+        self.pipe.transformer = autoquant(self.pipe.transformer, error_on_unseen=False)
+        self.pipe.vae = autoquant(self.pipe.vae, error_on_unseen=False)
+        self.pipe.text_encoder = autoquant(self.pipe.text_encoder, error_on_unseen=False)
+        self.pipe.text_encoder_2 = autoquant(self.pipe.text_encoder_2, error_on_unseen=False)
+
+        start_time = time.time()
+        print("Start warming-up pipeline")
+        self.pipe("Hello world!") # Warm-up for compiling
+        end_time = time.time()
+        time_taken = end_time - start_time
+        print(f"Time taken: {time_taken:.2f} seconds")
+
+    def __call__(self, data: Dict[str, Any]) -> Image.Image:
         logger.info(f"Received incoming request with {data=}")
 
         if "inputs" in data and isinstance(data["inputs"], str):
@@ -276,7 +61,7 @@ class EndpointHandler:
 
         parameters = data.pop("parameters", {})
 
-        num_inference_steps = parameters.get("num_inference_steps", 8 if IS_TURBO else 28)
+        num_inference_steps = parameters.get("num_inference_steps", 28)
         width = parameters.get("width", 1024)
         height = parameters.get("height", 1024)
         guidance_scale = parameters.get("guidance_scale", 3.5)
@@ -284,20 +69,17 @@ class EndpointHandler:
         # seed generator (seed cannot be provided as is but via a generator)
         seed = parameters.get("seed", 0)
         generator = torch.manual_seed(seed)
-
-        start = time.time()
-        image = self.pipeline(  # type: ignore
+        start_time = time.time()
+        result = self.pipe(  # type: ignore
             prompt,
             height=height,
             width=width,
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             generator=generator,
-            output_type="pil",
         ).images[0]
-        end = time.time()
-        print(f'Elapsed {end - start:.3f} sec. / prompt:"{prompt}" / size:{width}x{height} / steps:{num_inference_steps} / guidance scale:{guidance_scale} / seed:{seed}')
-
-        return pil_to_base64(image, self.repo_id, prompt, height, width, num_inference_steps, guidance_scale, seed)
-
+        end_time = time.time()
+        time_taken = end_time - start_time
+        print(f"Time taken: {time_taken:.2f} seconds")
 
+        return result