.

options/Video_model/Model.py

@@ -12,10 +12,12 @@ svd_path = 'stabilityai/stable-video-diffusion-img2vid-xt-1-1'
 lora_repo_path = 'RED-AIGC/TDD'
 lora_weight_name = 'svd-xt-1-1_tdd_lora_weights.safetensors'
 
+device="cuda" if torch.cuda.is_available() else "cpu"
+
 # if torch.cuda.is_available():
 noise_scheduler = TDDSVDStochasticIterativeScheduler(num_train_timesteps = 250, sigma_min = 0.002, sigma_max = 700.0, sigma_data = 1.0, 
                                                     s_noise = 1.0, rho = 7, clip_denoised = False)
-pipeline = StableVideoDiffusionPipeline.from_pretrained(svd_path, scheduler = noise_scheduler, torch_dtype = torch.float16, variant = "fp16").to('cuda')
+pipeline = StableVideoDiffusionPipeline.from_pretrained(svd_path, scheduler = noise_scheduler, torch_dtype = torch.float16, variant = "fp16").to(device)
 load_lora_weights(pipeline.unet, lora_repo_path, weight_name = lora_weight_name)
 
 # @spaces.GPU
@@ -44,7 +46,7 @@ def Video(
     base_count = len(glob(os.path.join(output_folder, "*.mp4")))
     video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
 
-    with torch.autocast("cuda"):
+    with torch.autocast(device):
         frames = pipeline(
             image, height = height, width = width,
             num_inference_steps = num_inference_steps,