Update app.py

app.py

@@ -6,36 +6,53 @@ from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepSchedu
 import os
 from uuid import uuid4
 
+# Check if CUDA is available and set device accordingly
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")  # Print device information
+
 # Load model on startup
-model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
-vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
-scheduler = UniPCMultistepScheduler(
-    prediction_type='flow_prediction',
-    use_flow_sigmas=True,
-    num_train_timesteps=1000,
-    flow_shift=5.0
-)
-pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
-pipe.scheduler = scheduler
-pipe.to("cuda")
+try:
+    print("Loading model...")
+    model_id = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+    vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32)
+    scheduler = UniPCMultistepScheduler(
+        prediction_type='flow_prediction',
+        use_flow_sigmas=True,
+        num_train_timesteps=1000,
+        flow_shift=5.0
+    )
+    pipe = WanPipeline.from_pretrained(model_id, vae=vae, torch_dtype=torch.bfloat16)
+    pipe.scheduler = scheduler
+    pipe.to(device)  # Move model to CUDA if available, otherwise CPU
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise e
 
 # Define the generation function
 def generate_video(prompt, negative_prompt="", height=720, width=1280, num_frames=81, guidance_scale=5.0):
-    output = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        guidance_scale=guidance_scale,
-    ).frames[0]
-
-    output_filename = f"{uuid4()}.mp4"
-    output_path = os.path.join("outputs", output_filename)
-    os.makedirs("outputs", exist_ok=True)
-    export_to_video(output, output_path, fps=16)
-
-    return output_path  # Gradio returns this as downloadable file/video
+    try:
+        print(f"Generating video for prompt: {prompt}")
+        output = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            guidance_scale=guidance_scale,
+        ).frames[0]
+
+        output_filename = f"{uuid4()}.mp4"
+        output_path = os.path.join("outputs", output_filename)
+        os.makedirs("outputs", exist_ok=True)
+        export_to_video(output, output_path, fps=16)
+
+        print(f"Video generated successfully: {output_path}")
+        return output_path  # Gradio returns this as a downloadable file/video
+
+    except Exception as e:
+        print(f"Error generating video: {e}")
+        return None  # Return None in case of error
 
 # Gradio Interface with API support
 iface = gr.Interface(
@@ -54,4 +71,9 @@ iface = gr.Interface(
     api=True  # This enables the API
 )
 
-iface.launch(share=True)  # `share=True` will allow others to access your app via a public link
+try:
+    print("Launching Gradio interface...")
+    iface.launch(share=True)  # `share=True` will allow others to access your app via a public link
+    print("Gradio interface launched successfully.")
+except Exception as e:
+    print(f"Error launching Gradio interface: {e}")