Update app.py

app.py

@@ -30,16 +30,16 @@ import diffusers
 
 # init
 dtype = torch.bfloat16
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = "cuda"
+
 print(device)
 base_model = "black-forest-labs/FLUX.1-dev"
 
 # load pipe
-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)
 
-txt2img_pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype).to(device)
+txt2img_pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=dtype)
 txt2img_pipe.__class__.load_lora_into_transformer = classmethod(load_lora_into_transformer)
+txt2img_pipe = txt2img_pipe.to(device)
 
 
 MAX_SEED = 2**32 - 1
@@ -108,6 +108,7 @@ def generate_random_4_digit_string():
     return ''.join(random.choices(string.digits, k=4))
 
 @spaces.GPU(duration=120)
+
 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")