Upload app.py

app.py

@@ -64,7 +64,7 @@ def change_base_model(repo_id: str, cn_on: bool): # , progress=gr.Progress(track
             #clear_cache()
             #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 = 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
@@ -74,7 +74,7 @@ def change_base_model(repo_id: str, cn_on: bool): # , progress=gr.Progress(track
             #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, vae=taef1).to(device)
+            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
@@ -149,6 +149,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
     try:
         #good_vae.to("cuda")
         #taef1.to("cuda")
+        pipe.vae = taef1
         pipe.to("cuda")
         generator = torch.Generator(device="cuda").manual_seed(seed)