Update app.py

app.py

@@ -30,7 +30,7 @@ import diffusers
 
 # init
 dtype = torch.bfloat16
-device = "cuda"
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 base_model = "black-forest-labs/FLUX.1-dev"
 
 # load pipe
@@ -42,7 +42,7 @@ txt2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_t
 
 # img2img model
 img2img_pipe = AutoPipelineForImage2Image.from_pretrained(base_model, vae=good_vae, transformer=txt2img_pipe.transformer, text_encoder=txt2img_pipe.text_encoder, tokenizer=txt2img_pipe.tokenizer, text_encoder_2=txt2img_pipe.text_encoder_2, tokenizer_2=txt2img_pipe.tokenizer_2, torch_dtype=dtype)
-# img2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
+img2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
 
 
 MAX_SEED = 2**32 - 1
@@ -152,13 +152,16 @@ def generate_random_4_digit_string():
 def run_lora(prompt, image_url, lora_strings_json, image_strength,  cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
     print("run_lora", prompt, lora_strings_json, cfg_scale, steps, width, height)
     gr.Info("Starting process")
-    txt2img_pipe.to("cuda")
+    
     torch.cuda.empty_cache() 
     img2img_model = False
     orginal_image = None
     if image_url:
         orginal_image = load_image(image_url)
         img2img_model = True
+        img2img_pipe.to(device)
+    else:
+        txt2img_pipe.to(device)
     # Set random seed for reproducibility
     if randomize_seed:
         with calculateDuration("Set random seed"):