Update app.py

app.py

@@ -1423,7 +1423,7 @@ base_model = "black-forest-labs/FLUX.1-dev"
 #TAEF1 is very tiny autoencoder which uses the same "latent API" as FLUX.1's VAE. FLUX.1 is useful for real-time previewing of the FLUX.1 generation process.#
 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)
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1)
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
                                                       vae=good_vae,
                                                       transformer=pipe.transformer,
@@ -1440,7 +1440,7 @@ controlnet_model = "InstantX/FLUX.1-dev-controlnet-canny"
 controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=dtype)
 pipe_canny = FluxControlNetPipeline.from_pretrained(
     base_model, controlnet=controlnet, torch_dtype=dtype
-).to(device)
+)
 
 MAX_SEED = 2**32-1
 
@@ -1507,6 +1507,7 @@ def generate_canny(image, type="canny"):
 @spaces.GPU(duration=100)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
     generator = torch.Generator(device="cuda").manual_seed(seed)
+    pipe.to('cuda')
     with calculateDuration("Generating image"):
         # Generate image
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -1524,6 +1525,7 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
 
 def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
     generator = torch.Generator(device="cuda").manual_seed(seed)
+    pipe_i2i.to('cuda')
     image_input = load_image(image_input_path)
     final_image = pipe_i2i(
         prompt=prompt_mash,
@@ -1541,8 +1543,9 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
 
 def generate_image_canny(prompt_mash, canny, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
     control_image = load_image(canny)
+    pipe_canny.to('cuda')    
 
-    image = pipe(
+    image = pipe_canny(
         prompt=prompt_mash,
         control_image=control_image,
         controlnet_conditioning_scale=0.6,