Upload handler.py

handler.py

@@ -9,7 +9,7 @@ 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, float8_weight_only
+from torchao.quantization import quantize_, autoquant, int8_dynamic_activation_int4_weight, float8_dynamic_activation_float8_weight, int8_dynamic_activation_int8_weight
 from torchao.quantization.quant_api import PerRow
 from diffusers import FluxPipeline, FluxTransformer2DModel, AutoencoderKL, TorchAoConfig
 from transformers import T5EncoderModel
@@ -72,7 +72,7 @@ def pil_to_base64(image: Image.Image, modelname: str, prompt: str, height: int,
 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, float8_weight_only())
+    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:
@@ -128,7 +128,7 @@ def load_pipeline_turbo(repo_id: str, dtype: torch.dtype) -> Any:
     pipe.unload_lora_weights()
     pipe.transformer.fuse_qkv_projections()
     pipe.vae.fuse_qkv_projections()
-    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else float8_weight_only()
+    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
@@ -141,7 +141,7 @@ def load_pipeline_turbo_compile(repo_id: str, dtype: torch.dtype) -> Any:
     pipe.unload_lora_weights()
     pipe.transformer.fuse_qkv_projections()
     pipe.vae.fuse_qkv_projections()
-    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else float8_weight_only()
+    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)
@@ -152,7 +152,7 @@ def load_pipeline_turbo_compile(repo_id: str, dtype: torch.dtype) -> Any:
 
 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 float8_weight_only()
+    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")
@@ -170,14 +170,14 @@ def load_pipeline_opt(repo_id: str, dtype: torch.dtype) -> Any:
     #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, float8_weight_only())
-    quantize_(pipe.text_encoder_2, float8_weight_only())
+    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 float8_weight_only()
+    weight = int8_dynamic_activation_int4_weight() if IS_4BIT else int8_dynamic_activation_int8_weight()
     pipe = FluxPipeline.from_pretrained(repo_id, torch_dtype=dtype)
     if IS_PARA:
         if IS_MGPU:
@@ -188,8 +188,8 @@ def load_pipeline_para(repo_id: str, dtype: torch.dtype) -> Any:
             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, float8_weight_only())
-    quantize_(pipe.text_encoder_2, float8_weight_only())
+    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")