x

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 = "city96/t5-v1_1-xxl-encoder-bf16"

src/pipeline.py

@@ -41,7 +41,7 @@ def remove_cache():
     torch.cuda.reset_peak_memory_stats()
 
 
-class InitModel:
+class InitializingModel:
 
     @staticmethod
     def load_text_encoder() -> T5EncoderModel:
@@ -53,16 +53,6 @@ class InitModel:
         )
         return text_encoder.to(memory_format=torch.channels_last)
 
-    @staticmethod
-    def load_vae() -> AutoencoderTiny:
-        print("Loading VAE model...")
-        vae = AutoencoderTiny.from_pretrained(
-            "XiangquiAI/FLUX_Vae_Model",
-            revision="103bcc03998f48ef311c100ee119f1b9942132ab",
-            torch_dtype=torch.bfloat16,
-        )
-        return vae
-
     @staticmethod
     def load_transformer(trans_path: str) -> FluxTransformer2DModel:
         print("Loading transformer model...")
@@ -74,35 +64,46 @@ class InitModel:
         return transformer.to(memory_format=torch.channels_last)
 
 
-
+class CompileTransformerDiffusion:
+    def __init__(self, pipeline, optimize=False):
+        self.pipeline = pipeline
+        self.optimize = optimize
+        if self.optimize:
+            self._compile_model()
+    
+    def _compile_model(self):
+        print("Compiling transformer model for optimized diffusion...")
+        self.pipeline.unet = torch.compile(self.pipeline.unet)
+    
+    def __call__(self, *args, **kwargs):
+        return self.pipeline(*args, **kwargs)
+    
 def load_pipeline() -> Pipeline:
 
 
-    transformer_path = os.path.join(HF_HUB_CACHE, "models--MyApricity--Flux_Transformer_float8/snapshots/66c5f182385555a00ec90272ab711bb6d3c197db")
-    transformer = InitModel.load_transformer(transformer_path)
+    base_transformer_path = os.path.join(HF_HUB_CACHE, "models--MyApricity--Flux_Transformer_float8/snapshots/66c5f182385555a00ec90272ab711bb6d3c197db")
+    base_transformer = InitializingModel.load_transformer(base_transformer_path)
     
-    text_encoder_2 = InitModel.load_text_encoder()
-    vae = InitModel.load_vae()
+    text_encoder_2 = InitializingModel.load_text_encoder()
 
 
     pipeline = DiffusionPipeline.from_pretrained(CHECKPOINT, 
                         revision=REVISION, 
-                        vae=vae, 
-                        transformer=transformer, 
+                        transformer=base_transformer, 
                         text_encoder_2=text_encoder_2, 
                         torch_dtype=torch.bfloat16)
     pipeline.to("cuda")
     try:
         pipeline.disable_vae_slice()
+        compiled_pipeline = CompileTransformerDiffusion(pipeline, optimize=False)
     except:
-        print("Using origin pipeline")
+        print("Stay safe here pipeline")
         
 
     promts_listing = [
-        "melanogen, endosome",
-        "buffer, cutie, buttinsky, prototrophic",
-        "puzzlehead, fistical, must return non duplicate",
-        "apical, polymyodous, tiptilt"
+        "sellate, Tremellales, thro, albescent",
+        "must return non duplicate",
+        "albaspidin, pillmonger, palaeocrystalline"
     ]
     
     for p in promts_listing: