Update app.py

app.py

@@ -4,21 +4,35 @@ import logging
 import torch
 from PIL import Image
 import spaces
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderTiny, AutoencoderKL
 import copy
 import random
 import time
+from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images
+from huggingface_hub import HfFileSystem, ModelCard
+
+from huggingface_hub import login
+hf_token = os.environ.get("HF_TOKEN")
+login(token=hf_token)
 
 # Load LoRAs from JSON file
 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 = "John6666/real-flux-10b-schnell-fp8-flux"
 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("black-forest-labs/FLUX.1-dev", 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
@@ -49,30 +63,33 @@ def update_selection(evt: gr.SelectData, width, height):
             width = 1024
             height = 768
     return (
-        gr.update(placeholder=new_placeholder),
+        #gr.update(placeholder=new_placeholder),
+        prompt,
         updated_text,
         evt.index,
         width,
         height,
+        lora_scale,
     )
 
 @spaces.GPU(duration=70)
 def generate_image(prompt, trigger_word, 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(
-            prompt=f"{prompt} {trigger_word}",
+        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
+            prompt=prompt_mash,
             num_inference_steps=steps,
             guidance_scale=cfg_scale,
             width=width,
             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:
@@ -81,6 +98,19 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
     selected_lora = loras[selected_index]
     lora_path = selected_lora["repo"]
     trigger_word = selected_lora["trigger_word"]
+    if(trigger_word):
+        if "trigger_position" in selected_lora:
+            if selected_lora["trigger_position"] == "prepend":
+                prompt_mash = f"{trigger_word} {prompt}"
+            else:
+                prompt_mash = f"{prompt} {trigger_word}"
+        else:
+            prompt_mash = f"{trigger_word} {prompt}"
+    else:
+        prompt_mash = prompt
+
+    with calculateDuration("Unloading LoRA"):
+        pipe.unload_lora_weights()
 
     # Load LoRA weights
     with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
@@ -94,10 +124,18 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
     
-    image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress)
-    pipe.to("cpu")
-    pipe.unload_lora_weights()
-    return image, seed  
+    image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
+    # Consume the generator to get the final image
+    final_image = None
+    step_counter = 0
+    for image in image_generator:
+        step_counter+=1
+        final_image = image
+        progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
+        yield image, seed, gr.update(value=progress_bar, visible=True)
+        
+    yield final_image, seed, gr.update(value=progress_bar, visible=False) 
+    
 
 run_lora.zerogpu = True
 
@@ -156,20 +194,21 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as app:
                 with gr.Row():
                     randomize_seed = gr.Checkbox(True, label="Randomize seed")
                     seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True)
-                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3.0, step=0.01, value=0.5)
+                    default_scale = gr.Checkbox(True, label="Use default LoRA scale")
+                    lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3.0, step=0.01, value=0.95)
 
     gallery.select(
         update_selection,
-        inputs=[width, height],
-        outputs=[prompt, selected_info, selected_index, width, height]
+        inputs=[width, height, default_scale, lora_scale],
+        outputs=[prompt, selected_info, selected_index, width, height, lora_scale]
     )
 
     gr.on(
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
         inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale],
-        outputs=[result, seed]
+        outputs=[result, seed, progress_bar]
     )
 
 app.queue(default_concurrency_limit=2).launch(show_error=True)
-app.launch()
+app.launch()