Upload app.py

app.py

@@ -4,7 +4,8 @@ import json
 import logging
 import torch
 from PIL import Image
-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 diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel
 from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_download
 import copy
@@ -21,35 +22,48 @@ from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_
 from tagger.tagger import predict_tags_wd, compose_prompt_to_copy
 from tagger.fl2flux import predict_tags_fl2_flux
 
+# Load LoRAs from JSON file
+with open('loras.json', 'r') as f:
+    loras = json.load(f)
+
+dtype = torch.bfloat16
+#dtype = torch.float8_e4m3fn
+device = "cuda" if torch.cuda.is_available() else "cpu"
 # Initialize the base model
 base_model = models[0]
 controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union'
 #controlnet_model_union_repo = 'InstantX/FLUX.1-dev-Controlnet-Union-alpha'
-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)
 controlnet_union = None
 controlnet = None
 last_model = models[0]
 last_cn_on = False
 
+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)
+
 # https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Union
 # https://huggingface.co/spaces/jiuface/FLUX.1-dev-Controlnet-Union
+@spaces.GPU()
 def change_base_model(repo_id: str, cn_on: bool):
     global pipe
     global controlnet_union
     global controlnet
     global last_model
     global last_cn_on
-    dtype = torch.bfloat16
-    #dtype = torch.float8_e4m3fn
     try:
         if (repo_id == last_model and cn_on is last_cn_on) or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(visible=True)
         if cn_on:
             #progress(0, desc=f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
             print(f"Loading model: {repo_id} / Loading ControlNet: {controlnet_model_union_repo}")
             clear_cache()
-            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype)
-            controlnet = FluxMultiControlNetModel([controlnet_union])
-            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype)
+            controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union_repo, torch_dtype=dtype).to(device)
+            controlnet = FluxMultiControlNetModel([controlnet_union]).to(device)
+            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype).to(device)
+            #pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
             last_model = repo_id
             last_cn_on = cn_on
             #progress(1, desc=f"Model loaded: {repo_id} / ControlNet Loaded: {controlnet_model_union_repo}")
@@ -58,7 +72,8 @@ def change_base_model(repo_id: str, cn_on: bool):
             #progress(0, desc=f"Loading model: {repo_id}")
             print(f"Loading model: {repo_id}")
             clear_cache()
-            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype)
+            pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=dtype, vae=taef1).to(device)
+            pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
             last_model = repo_id
             last_cn_on = cn_on
             #progress(1, desc=f"Model loaded: {repo_id}")
@@ -70,12 +85,6 @@ def change_base_model(repo_id: str, cn_on: bool):
 
 change_base_model.zerogpu = True
 
-# Load LoRAs from JSON file
-with open('loras.json', 'r') as f:
-    loras = json.load(f)
-
-MAX_SEED = 2**32-1
-
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -118,9 +127,13 @@ 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, cn_on, progress=gr.Progress(track_tqdm=True)):
     global pipe
+    global taef1
+    global good_vae
     global controlnet
     global controlnet_union
     try:
+        good_vae.to("cuda")
+        taef1.to("cuda")
         pipe.to("cuda")
         generator = torch.Generator(device="cuda").manual_seed(seed)
         
@@ -129,7 +142,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
             modes, images, scales = get_control_params()
             if not cn_on or len(modes) == 0:
                 progress(0, desc="Start Inference.")
-                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,
@@ -137,12 +150,15 @@ 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]
+                    output_type="pil",
+                    good_vae=good_vae,
+                ):
+                    yield img
             else:
                 progress(0, desc="Start Inference with ControlNet.")
                 if controlnet is not None: controlnet.to("cuda")
                 if controlnet_union is not None: controlnet_union.to("cuda")
-                image = pipe(
+                for img in pipe(
                     prompt=prompt_mash,
                     control_image=images,
                     control_mode=modes,
@@ -153,23 +169,35 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
                     controlnet_conditioning_scale=scales,
                     generator=generator,
                     joint_attention_kwargs={"scale": lora_scale},
-                ).images[0]
+                ).images:
+                    yield img
     except Exception as e:
         print(e)
         raise gr.Error(f"Inference Error: {e}")
-    return image
 
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
               lora_scale, lora_json, cn_on, progress=gr.Progress(track_tqdm=True)):
     global pipe
+    global taef1
+    global good_vae
+    global controlnet
+    global controlnet_union
     if selected_index is None and not is_valid_lora(lora_json):
         gr.Info("LoRA isn't selected.")
     #    raise gr.Error("You must select a LoRA before proceeding.")
     progress(0, desc="Preparing Inference.")
 
+    with calculateDuration("Unloading LoRA"):
+        try:
+            pipe.unfuse_lora()
+            pipe.unload_lora_weights()
+        except Exception as e:
+            print(e)
+
     prompt_mash = prompt
     if is_valid_lora(lora_json):
-        with calculateDuration("Loading LoRA weights"):
+        # Load External LoRA weights
+        with calculateDuration("Loading External LoRA weights"):
             fuse_loras(pipe, lora_json)
             trigger_word = get_trigger_word(lora_json)
             prompt_mash = f"{prompt} {trigger_word}"
@@ -200,17 +228,28 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
             seed = random.randint(0, MAX_SEED)
     
     progress(0, desc="Running Inference.")
-
-    image = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, progress)
-    if is_valid_lora(lora_json):
-        pipe.unfuse_lora()
-        pipe.unload_lora_weights()
-    if selected_index is not None: pipe.unload_lora_weights()
-    pipe.to("cpu")
-    if controlnet is not None: controlnet.to("cpu")
-    if controlnet_union is not None: controlnet_union.to("cpu")
-    clear_cache()
-    return image, seed  
+    image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, cn_on, 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)
+    #if is_valid_lora(lora_json):
+    #    pipe.unfuse_lora()
+    #    pipe.unload_lora_weights()
+    #if selected_index is not None: pipe.unload_lora_weights()
+    #pipe.to("cpu")
+    #good_vae.to("cpu")
+    #taef1.to("cpu")
+    #if controlnet is not None: controlnet.to("cpu")
+    #if controlnet_union is not None: controlnet_union.to("cpu")
+    #clear_cache()
+    #return final_image, seed  # Return the final image and seed
 
 def get_huggingface_safetensors(link):
   split_link = link.split("/")
@@ -343,6 +382,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
                 custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
                 deselect_lora_button = gr.Button("Deselect LoRA", variant="secondary")
             with gr.Column():
+                progress_bar = gr.Markdown(elem_id="progress",visible=False)
                 result = gr.Image(label="Generated Image", format="png", show_share_button=False)
         with gr.Group():
             model_name = gr.Dropdown(label="Base Model", info="You can enter a huggingface model repo_id to want to use.", choices=models, value=models[0], allow_custom_value=True)
@@ -450,7 +490,7 @@ with gr.Blocks(theme='Nymbo/Nymbo_Theme', fill_width=True, css=css) as app:
         fn=run_lora,
         inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height,
                  lora_scale, lora_repo_json, cn_on], 
-        outputs=[result, seed],
+        outputs=[result, seed, progress_bar],
         queue=True,
         show_api=True,
     )