Update app.py

app.py

@@ -5,40 +5,57 @@ import spaces
 import torch
 import os
 import time
-from diffusers import FluxPipeline
+from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
+from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 
-MAX_SEED = np.iinfo(np.int32).max
-pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cuda")
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
 torch.cuda.empty_cache()
+
+MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
 
 @spaces.GPU(duration=75)
 def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
-    image = pipe(
-        prompt,
-        height=height,
-        width=width,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        max_sequence_length=512,
-        generator=torch.Generator().manual_seed(seed)
-    ).images[0]
-    image.save("flux-dev.png")
+    generator = torch.Generator().manual_seed(seed)
+    
+    # Get the final image from the generator
+    final_img = None
+    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",
+            good_vae=good_vae,
+        ):
+            final_img = img
     
-    if image is not None:
+    # Save image to /tmp directory
+    if final_img is not None:
+        # Create a unique filename with timestamp
         timestamp = int(time.time())
         filename = f"flux_generated_{timestamp}_{seed}.png"
         
         # Save to /tmp directory (as requested)
         tmp_path = os.path.join("/tmp", filename)
         os.makedirs("/tmp", exist_ok=True)
-        image.save(tmp_path)
+        final_img.save(tmp_path)
         
-        return image, seed, tmp_path
+        return final_img, seed, tmp_path
+    
+    return final_img, seed, ""
     
 examples = [
     "a tiny astronaut hatching from an egg on the moon",