Enhance image generation pipeline by enabling tiling, slicing, and sequential CPU offload in app.py

app.py

@@ -24,9 +24,10 @@ from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_imag
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device)
 pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to(device)
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+pipe.vae.enable_tiling()
+pipe.vae.enable_slicing()
+pipe.enable_sequential_cpu_offload() # offloads modules to CPU on a submodule level (rather than model level)
 torch.cuda.empty_cache()
 
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
@@ -92,7 +93,6 @@ def generate_item_image(object_t2i_prompt):
         height=1024,
         generator=torch.Generator("cpu").manual_seed(0),
         output_type="pil",
-        good_vae=good_vae,
     ):
         yield trial_id, image
     # img_path = t2i_client.predict(