pp

pyproject.toml

@@ -25,7 +25,6 @@ dependencies = [
 [[tool.edge-maxxing.models]]
 repository = "black-forest-labs/FLUX.1-schnell"
 revision = "741f7c3ce8b383c54771c7003378a50191e9efe9"
-exclude = ["transformer", "vae", "text_encoder_2"]
 
 [[tool.edge-maxxing.models]]
 repository = "RichardWilliam/XULF_T5_bf16"

src/pipeline.py

@@ -33,7 +33,36 @@ CHECKPOINT = "black-forest-labs/FLUX.1-schnell"
 REVISION = "741f7c3ce8b383c54771c7003378a50191e9efe9"
 Pipeline = None
 
+class CleanAndOptimization:
+    def __init__(self, model, device="cuda"):
+        self.model = model
+        self.device = device
+        self.cache = {}
+
+    @staticmethod
+    def enhance_performance():
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cudnn.deterministic = False
+        return "Torch backend opt"
+    
+    def preprocess(self, data):
+        return [d[::-1] for d in data]
+
+    def quantize_model(self):
+        self.model = quantize_(self.model, weight_dtype=torch.float16)
+        self.model = int8_weight_only(self.model)
+        return self.model
+
+    def optimize_memory(self):
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.cache.clear()
 
+    def apply_all(self, data):
+        self.optimize_memory()
+        processed = self.preprocess(data)
+        self.quantize_model()
+        return self.enhance_performance()
 
 def t5_mapping_loader(repo_path):
 
@@ -84,6 +113,7 @@ def load_pipeline() -> Pipeline:
     try:
         base_encoder_2 = os.path.join(HF_HUB_CACHE, "models--RichardWilliam--XULF_T5_bf16/snapshots/63a3d9ef7b586655600ac9bd4e4747d038237761")
         text_encoder_2 = t5_mapping_loader(repo_path=base_encoder_2)
+        # opt opt opt opt opt opt opt
     except:
         text_encoder_2 =  T5EncoderModel.from_pretrained("RichardWilliam/XULF_T5_bf16",
                     revision = "63a3d9ef7b586655600ac9bd4e4747d038237761", 
@@ -101,8 +131,10 @@ def load_pipeline() -> Pipeline:
     try:
         torch.cuda.empty_cache()
         gc.collect()
-        # flux_pipeline.enable_sequential_cpu_offload()
+
         flux_pipeline.transformer.enable_cuda_graph()
+        torch_opt = CleanAndOptimization.enhance_performance()
+        print(torch_opt)
     except:
         pass