Update app.py

app.py

@@ -40,6 +40,8 @@ os.putenv("TOKENIZERS_PARALLELISM","False")
 model_id = "Skywork/SkyReels-V1-Hunyuan-I2V"
 base_model_id = "hunyuanvideo-community/HunyuanVideo"
 
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 def init_predictor():
     global pipe
     text_encoder = LlamaModel.from_pretrained(
@@ -59,7 +61,7 @@ def init_predictor():
         text_encoder=text_encoder,
         torch_dtype=torch.bfloat16,
     ).to("cpu")
-    pipe.to(torch.device('cuda'))
+    pipe.to(device)
 
 @spaces.GPU(duration=60)
 def generate(segment, image, prompt, size, guidance_scale, num_inference_steps, frames, seed, progress=gr.Progress(track_tqdm=True) ):
@@ -69,7 +71,7 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = pipe.encode_prompt(
         prompt=prompt, prompt_2=prompt, device=device
         )
-        pipe.scheduler.set_timesteps(num_inference_steps, device=torch.device('cuda'))
+        pipe.scheduler.set_timesteps(num_inference_steps, device=device)
         timesteps = pipe.scheduler.timesteps
         all_timesteps_cpu = timesteps.cpu()
         timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)