chore: Integrate live preview functionality in image enhancement pipeline

app.py

@@ -3,7 +3,7 @@ from diffusers import FluxPriorReduxPipeline, FluxPipeline
 from diffusers.utils import load_image
 import gradio as gr
 import spaces
-# from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
+from live_preview_helpers import flux_pipe_call_that_returns_an_iterable_of_images
 
 pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained("black-forest-labs/FLUX.1-Redux-dev", revision="refs/pr/8", torch_dtype=torch.bfloat16).to("cuda")
 pipe = FluxPipeline.from_pretrained(
@@ -11,26 +11,36 @@ pipe = FluxPipeline.from_pretrained(
     text_encoder=None,
     text_encoder_2=None,
     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)
+).to("cuda")
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 # pipe.enable_sequential_cpu_offload()
 
-@spaces.GPU(duration=120)
+@spaces.GPU
 def enhance_image(image_path, keep_aspect_ratio=False):
     print(image_path)
     image = load_image(image_path)
     print(image.size)
     width, height = image.size if keep_aspect_ratio else (None, None)
     pipe_prior_output = pipe_prior_redux(image)
-    images = pipe(
-        height=height,
-        width=width,
+    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
         guidance_scale=2.5,
         num_inference_steps=50,
+        width=width,
+        height=height,
         generator=torch.Generator("cpu").manual_seed(0),
-        **pipe_prior_output,
-    ).images
-    return images[0]
+        output_type="pil",
+        **pipe_prior_output
+        ):
+        yield img
+    # images = pipe(
+    #     height=height,
+    #     width=width,
+    #     guidance_scale=2.5,
+    #     num_inference_steps=50,
+    #     generator=torch.Generator("cpu").manual_seed(0),
+    #     **pipe_prior_output,
+    # ).images
+    # return images[0]
 
 
 with gr.Blocks(title="Flux.1 Dev Redux") as demo: