Update src/pipeline.py

src/pipeline.py

@@ -71,16 +71,14 @@ def load_pipeline() -> Pipeline:
     pipeline = FluxPipeline.from_pretrained(ckpt_id, revision=ckpt_revision, transformer=transformer, vae=vae, local_files_only=True, torch_dtype=torch.bfloat16,)
     pipeline.to("cuda")
 
-    # Optimize memory format
-    for component in [pipeline.text_encoder, pipeline.text_encoder_2, pipeline.transformer, pipeline.vae]:
-        component.to(memory_format=torch.channels_last)
+    pipeline.to(memory_format=torch.channels_last)
     
     pipeline.transformer = torch.compile(pipeline.transformer, mode="max-autotune", fullgraph=True)
-    pipeline.vae = torch.compile(pipeline.vae, mode="max-autotune", fullgraph=True)
+    # pipeline.vae = torch.compile(pipeline.vae, mode="max-autotune", fullgraph=True)
         
     for _ in range(2):
         pipeline(prompt="insensible, timbale, pothery, electrovital, centrodesmus", width=1024, height=1024, guidance_scale=0.0, num_inference_steps=4, max_sequence_length=256)
-    empty_cache()
+    # empty_cache()
     return pipeline