Update app.py

app.py

@@ -19,12 +19,12 @@ with open('loras.json', 'r') as f:
 
 # Initialize the base model
 dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
-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)
+taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to("cpu")
+good_vae = AutoencoderKL.from_pretrained(base_model, subfolder="vae", torch_dtype=dtype).to("cpu")
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype, vae=taef1).to("cpu")
 pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
                                                       vae=good_vae,
                                                       transformer=pipe.transformer,
@@ -33,7 +33,7 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
                                                       text_encoder_2=pipe.text_encoder_2,
                                                       tokenizer_2=pipe.tokenizer_2,
                                                       torch_dtype=dtype
-                                                     )
+                                                     ).to("cpu")
 
 MAX_SEED = 2**32-1
 
@@ -80,8 +80,8 @@ 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, progress):
-    pipe.to("cuda")
-    generator = torch.Generator(device="cuda").manual_seed(seed)
+    pipe.to("cpu")
+    generator = torch.Generator(device="cpu").manual_seed(seed)
     with calculateDuration("Generating image"):
         # Generate image
         for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
@@ -98,8 +98,8 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
             yield img
 
 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")
+    generator = torch.Generator(device="cpu").manual_seed(seed)
+    pipe_i2i.to("cpu")
     image_input = load_image(image_input_path)
     final_image = pipe_i2i(
         prompt=prompt_mash,