.

src/lyric_generator.py

@@ -1,4 +1,5 @@
 from transformers import pipeline
+import torch
 from typing import Dict, List, Optional
 
 class LyricGenerator:
@@ -9,10 +10,13 @@ class LyricGenerator:
         Args:
             model_name: The name of the pre-trained model to use
         """
+        # Check if CUDA is available
+        device = 0 if torch.cuda.is_available() else -1
+        
         self.generator = pipeline(
             "text-generation",
             model=model_name,
-            device=0 if pipeline.device.type == "cuda" else -1
+            device=device
         )
         
         # Genre-specific prompts to guide generation