Upload app.py

app.py

@@ -33,9 +33,9 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 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'
-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=torch.bfloat16)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1)
 controlnet_union = None
 controlnet = None
 last_model = models[0]
@@ -61,7 +61,7 @@ def change_base_model(repo_id: str, cn_on: bool):
             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)
+            pipe = FluxControlNetPipeline.from_pretrained(repo_id, controlnet=controlnet, torch_dtype=dtype, vae=taef1)
             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
@@ -71,7 +71,7 @@ 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)
             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
@@ -123,9 +123,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)
         
@@ -172,6 +176,10 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
 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.")
@@ -221,6 +229,8 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, wid
         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()