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)
-good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype)
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1)
+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]
@@ -59,9 +59,9 @@ def change_base_model(repo_id: str, cn_on: bool):
             #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, vae=taef1)
+            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, 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
@@ -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, vae=taef1)
+            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