Move FluxPipeline to GPU and back to CPU during image generation; remove unused AutoencoderKL reference

app.py

@@ -21,7 +21,6 @@ from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_imag
 llm_client = Client("Qwen/Qwen2.5-72B-Instruct")
 
 pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16)
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 def generate_t2i_prompt(item_name):
@@ -72,6 +71,7 @@ def preprocess_pil_image(image: Image.Image) -> Tuple[str, Image.Image]:
 @spaces.GPU
 def generate_item_image(object_t2i_prompt):
     trial_id = ""
+    pipe.to("cuda")
     for image in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         prompt=object_t2i_prompt,
         guidance_scale=3.5,
@@ -80,7 +80,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(
@@ -94,6 +93,7 @@ def generate_item_image(object_t2i_prompt):
 	# 	api_name="/infer"
     # )[0]
     # image = Image.open(img_path)
+    pipe.to("cpu")
     trial_id, processed_image = preprocess_pil_image(image)
     yield trial_id, processed_image