Update gradio_app.py

gradio_app.py

@@ -18,6 +18,8 @@ logging.basicConfig(level=logging.DEBUG)
 
 # Ensure necessary devices are available
 device = "cuda" if torch.cuda.is_available() else "cpu"
+logger.info("device: ", device)
+
 accelerator = Accelerator(mixed_precision='bf16', device_placement=True)
 
 # Model paths (replace these with your actual model paths)
@@ -124,7 +126,7 @@ def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
     info = lora_model.load_state_dict(weights_sd, strict=True)
     logger.info(f"Loaded LoRA weights from {LORA_WEIGHTS_PATH}: {info}")
     lora_model.eval()
-    lora_model.to("cuda")
+    lora_model.to(device)
 
     # Process the seed
     if randomize_seed:
@@ -145,7 +147,7 @@ def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
     logger.debug("Conditional image preprocessed.")
 
     # Encode the image to latents
-    ae.to("cuda")
+    ae.to(device)
     latents = ae.encode(image)
     logger.debug("Image encoded to latents.")
 
@@ -153,8 +155,8 @@ def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
     conditions[prompt] = latents.to("cpu")
 
     ae.to("cpu")
-    clip_l.to("cuda")
-    t5xxl.to("cuda")
+    clip_l.to(device)
+    t5xxl.to(device)
 
     # Encode the prompt
     tokenize_strategy = strategy_flux.FluxTokenizeStrategy(512)
@@ -192,8 +194,7 @@ def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
     clip_l.to("cpu")
     t5xxl.to("cpu")
 
-    torch.cuda.empty_cache()
-    model.to("cuda")
+    model.to(device)
 
     # import pdb
     # pdb.set_trace()
@@ -209,7 +210,7 @@ def infer(prompt, sample_image, frame_num, seed=0, randomize_seed=False):
     x = x.float()
     x = flux_utils.unpack_latents(x, packed_latent_height, packed_latent_width)
     model.to("cpu")
-    ae.to("cuda")
+    ae.to(device)
     with accelerator.autocast(), torch.no_grad():
         x = ae.decode(x)
     logger.debug("Latents decoded into image.")