Update app.py

app.py

@@ -13,12 +13,16 @@ import numpy as np
 import gradio as gr
 import spaces
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
 pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
 flow_shift = 1.0 #5.0  1.0 for image, 5.0 for 720P, 3.0 for 480P
 pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=flow_shift)
 
+pipe.to(device)
+
 # Configure DDIMScheduler with a beta schedule
 # pipe.scheduler = DDIMScheduler.from_config(
 #     pipe.scheduler.config,
@@ -82,7 +86,7 @@ def generate(prompt, negative_prompt, width=1024, height=1024, num_inference_ste
         pipe.set_adapters([DEFAULT_LORA_NAME], adapter_weights=[1.0])
 
 
-    pipe.to("cuda")
+    #pipe.to("cuda")
     # apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold=0.2))
     apply_cache_on_pipe(
         pipe,