Add live previews (now for realz) (#18)

- Add live previews (now for realz) (2ed4418f4c92b7cc2acc65a4b655a3a0c28eadd4)
- Update app.py (aaff7099c16d675defc700eb3004d4015af257a0)
- Update loras.json (0a8b917725e831072aefd745f37bcea1888d4a84)

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("/")

live_preview_helpers.py

@@ -59,6 +59,7 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
     return_dict: bool = True,
     joint_attention_kwargs: Optional[Dict[str, Any]] = None,
     max_sequence_length: int = 512,
+    good_vae: Optional[Any] = None,
 ):
     height = height or self.default_sample_size * self.vae_scale_factor
     width = width or self.default_sample_size * self.vae_scale_factor
@@ -156,10 +157,10 @@ def flux_pipe_call_that_returns_an_iterable_of_images(
         yield self.image_processor.postprocess(image, output_type=output_type)[0]
         torch.cuda.empty_cache()
 
-    # Final image
-    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-    latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
-    image = self.vae.decode(latents, return_dict=False)[0]
+    # Final image using good_vae
+    latents = self._unpack_latents(latents, height, width, good_vae.config.vae_scale_factor)
+    latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
+    image = good_vae.decode(latents, return_dict=False)[0]
     self.maybe_free_model_hooks()
     torch.cuda.empty_cache()
-    return self.image_processor.postprocess(image, output_type=output_type)[0], 0    
+    yield self.image_processor.postprocess(image, output_type=output_type)[0]

loras.json

@@ -7,10 +7,10 @@
     "aspect": "portrait"
   },
   {
-    "image": "https://huggingface.co/alvdansen/frosting_lane_flux/resolve/main/images/content%20-%202024-08-11T005936.346.jpeg",
-    "title": "Frosting Lane Flux",
-    "repo": "alvdansen/frosting_lane_flux",
-    "trigger_word": ""
+    "image": "https://huggingface.co/alvdansen/softpasty-flux-dev/resolve/main/images/ComfyUI_00814_%20(2).png",
+    "title": "SoftPasty",
+    "repo": "alvdansen/softpasty-flux-dev",
+    "trigger_word": "araminta_illus illustration style"
   },
   {
     "image": "https://huggingface.co/AIWarper/RubberCore1920sCartoonStyle/resolve/main/images/Rub_00006_.png",
@@ -19,12 +19,32 @@
     "trigger_word": "RU883R style",
     "trigger_position": "prepend"
   },
+  {
+    "image": "https://huggingface.co/mgwr/Cine-Aesthetic/resolve/main/images/00030-1333633802.png",
+    "title": "Cine Aesthetic",
+    "repo": "mgwr/Cine-Aesthetic",
+    "trigger_word": "mgwr/cine",
+    "trigger_position": "prepend"
+  },
+  {
+    "image": "https://huggingface.co/Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration/resolve/main/images/example3.png",
+    "title": "Blended Realistic Illustration",
+    "repo": "Shakker-Labs/FLUX.1-dev-LoRA-blended-realistic-illustration",
+    "trigger_word": "artistic style blends reality and illustration elements"
+  },
   {
     "image": "https://github.com/XLabs-AI/x-flux/blob/main/assets/readme/examples/picture-6-rev1.png?raw=true",
     "title": "flux-Realism",
     "repo": "XLabs-AI/flux-RealismLora",
     "trigger_word": ""
   },
+  {
+    "image": "https://huggingface.co/multimodalart/vintage-ads-flux/resolve/main/samples/j_XNU6Oe0mgttyvf9uPb3_dc244dd3d6c246b4aff8351444868d66.png",
+    "title": "Vintage Ads",
+    "repo":"multimodalart/vintage-ads-flux",
+    "trigger_word": "a vintage ad of",
+    "trigger_position": "prepend"
+  },
   {
     "image": "https://huggingface.co/nerijs/animation2k-flux/resolve/main/images/Q8-oVxNnXvZ9HNrgbNpGw_02762aaaba3b47859ee5fe9403a371e3.png",
     "title": "animation2k",
@@ -49,6 +69,12 @@
     "repo": "alvdansen/flux-koda",
     "trigger_word": "flmft style"
   },
+  {
+    "image": "https://huggingface.co/alvdansen/frosting_lane_flux/resolve/main/images/content%20-%202024-08-11T005936.346.jpeg",
+    "title": "Frosting Lane Flux",
+    "repo": "alvdansen/frosting_lane_flux",
+    "trigger_word": ""
+  },
   {
     "image": "https://pbs.twimg.com/media/GU7NsZPa8AA4Ddl?format=jpg&name=4096x4096",
     "title": "Half Illustration",
@@ -92,6 +118,13 @@
     "weights": "anime_lora.safetensors",
     "trigger_word": ", anime"
   },
+  {
+    "image": "https://github.com/XLabs-AI/x-flux/blob/main/assets/readme/examples/result_14.png?raw=true",
+    "title": "80s Cyberpunk",
+    "repo": "fofr/flux-80s-cyberpunk",
+    "trigger_word": "style of 80s cyberpunk",
+    "aspect": "portrait"
+  },
   {
     "image": "https://huggingface.co/kudzueye/Boreal/resolve/main/images/ComfyUI_00845_.png",
     "title": "Boreal",