Update src/pipeline.py

src/pipeline.py

@@ -9,7 +9,7 @@ from torch import Generator
 
 Pipeline = None
 # Consistent environment variable setting
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:1024"  # More robust memory management
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"  # More robust memory management
 
 ckpt_id = "black-forest-labs/FLUX.1-schnell"
 
@@ -18,7 +18,7 @@ def load_pipeline() -> Pipeline:
     # torch.cuda.empty_cache()
 
     dtype = torch.bfloat16
-    text_encoder = T5EncoderModel.from_pretrained(
+    text_encoder_2 = T5EncoderModel.from_pretrained(
         "city96/t5-v1_1-xxl-encoder-bf16", torch_dtype=dtype
     ).to(memory_format=torch.channels_last)
     vae = AutoencoderKL.from_pretrained(ckpt_id, subfolder="vae", torch_dtype=dtype).to(memory_format=torch.channels_last)
@@ -27,7 +27,7 @@ def load_pipeline() -> Pipeline:
     pipeline = DiffusionPipeline.from_pretrained(
         ckpt_id,
         vae=vae,
-        text_encoder=text_encoder,
+        text_encoder_2=text_encoder_2,
         torch_dtype=dtype,
     )#.to("cuda")