Update app.py

app.py

@@ -5,7 +5,9 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
+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
@@ -16,11 +18,18 @@ with open('loras.json', 'r') as f:
     loras = json.load(f)
 
 # Initialize the base model
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 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)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to(device)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to(device)
 
 MAX_SEED = 2**32-1
 
+pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -61,10 +70,9 @@ def update_selection(evt: gr.SelectData, width, height):
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
     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,
             num_inference_steps=steps,
             guidance_scale=cfg_scale,
@@ -72,13 +80,14 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
             height=height,
             generator=generator,
             joint_attention_kwargs={"scale": lora_scale},
-        ).images[0]
-    return image
+            output_type="pil",
+            good_vae=good_vae,
+        ):
+            yield img
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
-
     selected_lora = loras[selected_index]
     lora_path = selected_lora["repo"]
     trigger_word = selected_lora["trigger_word"]
@@ -92,24 +101,31 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
             prompt_mash = f"{trigger_word} {prompt}"
     else:
         prompt_mash = prompt
+
     # Load LoRA weights
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
         if "weights" in selected_lora:
             pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"])
-            #pipe.fuse_lora()
         else:
             pipe.load_lora_weights(lora_path)
-            #pipe.fuse_lora()
+
     # Set random seed for reproducibility
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
+
+    image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
     
-    image = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    # Consume the generator to get the final image
+    final_image = None
+    for image in image_generator:
+        final_image = image
+        yield image, seed  # Yield intermediate images and seed
+
     pipe.to("cpu")
-    #pipe.unfuse_lora()
     pipe.unload_lora_weights()
-    return image, seed  
+    
+    return final_image, seed  # Return the final image and seed
 
 def get_huggingface_safetensors(link):
   split_link = link.split("/")