Update src/pipeline.py

src/pipeline.py

@@ -26,7 +26,7 @@ import numpy as np
 import torch.nn as nn
 import torch.nn.functional as F
 from torchao.quantization import quantize_, float8_weight_only, int4_weight_only
-from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe
+from para_attn.first_block_cache.diffusers_adapters import apply_cache_on_pipe as cacher
 import os
 os.environ['PYTORCH_CUDA_ALLOC_CONF']="expandable_segments:True"
 os.environ["TOKENIZERS_PARALLELISM"] = "True"
@@ -34,6 +34,7 @@ torch._dynamo.config.suppress_errors = True
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.enabled = True
 torch.backends.cudnn.benchmark = True
+torch.cuda.set_per_process_memory_fraction(0.99)
 
 # globals
 Pipeline = None
@@ -67,7 +68,7 @@ def load_pipeline() -> Pipeline:
     )
     # pipeline.vae = torch.compile(vae)
     pipeline.to("cuda")
-    pipeline = apply_cache_on_pipe(pipeline,residual_diff_threshold=0.56)
+    pipeline = cacher(pipeline,residual_diff_threshold=0.56)
     quantize_(pipeline.vae, int4_weight_only())
     
     warmup_ = "controllable varied focus thai warriors entertainment claude still goat gang gang yeah"