Upload handler.py

handler.py

@@ -42,10 +42,10 @@ if IS_COMPILE:
     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
+    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
@@ -207,6 +207,25 @@ def load_pipeline_para(repo_id: str, dtype: torch.dtype) -> Any:
     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()
+    apply_cache_on_pipe(pipe, residual_diff_threshold=0.12)
+    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")
+    #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")
+    #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
+
 class EndpointHandler:
     def __init__(self, path=""):
         repo_id = "NoMoreCopyrightOrg/flux-dev-8step" if IS_TURBO else "NoMoreCopyrightOrg/flux-dev"
@@ -216,7 +235,7 @@ class EndpointHandler:
         print_vram()
         print("Loading pipeline...")
         if IS_AUTOQ: self.pipeline = load_pipeline_autoquant(repo_id, dtype)
-        elif IS_COMPILE: self.pipeline = load_pipeline_para(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)
         if not IS_COMPILE: