Update app.py

app.py

@@ -5,19 +5,23 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 import copy
 import random
 import time
 
+dtype = torch.bfloat16
+
 # Load LoRAs from JSON file
 with open('loras.json', 'r') as f:
     loras = json.load(f)
 
 # Initialize the base model
-base_model = "black-forest-labs/FLUX.1-dev"
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", 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 = 2**32-1
 
@@ -60,19 +64,23 @@ def update_selection(evt: gr.SelectData, width, height):
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device="cuda").manual_seed(seed)
     
     with calculateDuration("Generating image"):
-        # Generate image
-        image = pipe(
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
             prompt=prompt_mash,
+            guidance_scale=guidance_scale,
             num_inference_steps=steps,
-            guidance_scale=cfg_scale,
             width=width,
             height=height,
+            guidance_scale=cfg_scale,
             generator=generator,
+            output_type="pil",
             joint_attention_kwargs={"scale": lora_scale},
-        ).images[0]
+        ):
+            yield img, seed
     return image
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):