.

src/classifier.py

@@ -7,17 +7,34 @@ from typing import Union, Tuple, List
 
 class MusicGenreClassifier:
     def __init__(self):
-        # Initialize both audio and text classification pipelines
-        self.text_classifier = pipeline(
-            "zero-shot-classification",
-            model="facebook/bart-large-mnli"
-        )
-        
-        # For audio classification, we'll use MIT's music classification model
-        self.audio_classifier = pipeline(
-            "audio-classification",
-            model="mit/ast-finetuned-audioset-10-10-0.4593"
-        )
+        try:
+            # Initialize both audio and text classification pipelines with auto device mapping
+            self.text_classifier = pipeline(
+                "zero-shot-classification",
+                model="facebook/bart-large-mnli",
+                device_map="auto"
+            )
+            
+            # For audio classification, we'll use MIT's music classification model
+            self.audio_classifier = pipeline(
+                "audio-classification",
+                model="mit/ast-finetuned-audioset-10-10-0.4593",
+                device_map="auto"
+            )
+        except Exception as e:
+            print(f"Warning: GPU initialization failed, falling back to CPU. Error: {str(e)}")
+            # Fall back to CPU if GPU initialization fails
+            self.text_classifier = pipeline(
+                "zero-shot-classification",
+                model="facebook/bart-large-mnli",
+                device="cpu"
+            )
+            
+            self.audio_classifier = pipeline(
+                "audio-classification",
+                model="mit/ast-finetuned-audioset-10-10-0.4593",
+                device="cpu"
+            )
         
         # Define standard genres for classification
         self.genres = [

src/lyric_generator.py

@@ -10,14 +10,25 @@ 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=device
-        )
+        try:
+            # Try to use CUDA if available
+            if torch.cuda.is_available():
+                device = "cuda"
+            else:
+                device = "cpu"
+            
+            self.generator = pipeline(
+                "text-generation",
+                model=model_name,
+                device_map="auto"  # Let transformers handle device mapping
+            )
+        except Exception as e:
+            print(f"Warning: GPU initialization failed, falling back to CPU. Error: {str(e)}")
+            self.generator = pipeline(
+                "text-generation",
+                model=model_name,
+                device="cpu"
+            )
         
         # Genre-specific prompts to guide generation
         self.genre_prompts = {