Update app.py

app.py

@@ -49,18 +49,31 @@ css="""
 def infer(prompt, seed=42, randomize_seed=False, width=896, height=1152, num_inference_steps=16, progress=gr.Progress(track_tqdm=True), guidancescale=3.5):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+        
     generator = torch.Generator().manual_seed(seed)
-    image = pipe(
-            prompt = prompt, 
-            width = width,
-            height = height,
-            num_inference_steps = num_inference_steps, 
-            generator = generator,
-            guidance_scale=guidancescale,
+    
+    #image = pipe(
+    #        prompt = prompt, 
+    #        width = width,
+    #        height = height,
+    #        num_inference_steps = num_inference_steps, 
+    #        generator = generator,
+    #        guidance_scale=guidancescale,
+    #        output_type="pil",
+    #).images[0] 
+    #return image, seed
+    for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
             output_type="pil",
-    ).images[0] 
-    # image = good_vae.decode(image, return_dict=False)[0]
-    return image, seed
+            good_vae=good_vae,
+        ):
+            yield img, seed
+    
 
 def encode_image(image_path):
     """Encode the image to base64."""