Create generate.py

generate.py

@@ -0,0 +1,38 @@
+import torch
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
+import numpy as np
+
+MODEL_ID = "facebook/musicgen-small"
+device = "cpu"
+
+print(f"[*] Initializing Engine: Loading {MODEL_ID}...")
+processor = AutoProcessor.from_pretrained(MODEL_ID)
+model = MusicgenForConditionalGeneration.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
+model.to(device)
+print("[+] Engine Ready.")
+
+def generate_music(prompt, duration):
+    if not prompt:
+        return None
+    
+    print(f"[#] Generating: {prompt} ({duration}s)")
+    try:
+        duration = min(int(duration), 30)
+        inputs = processor(text=[prompt], padding=True, return_tensors="pt").to(device)
+        max_tokens = int(duration * 50)
+        
+        with torch.no_grad():
+            audio_values = model.generate(
+                **inputs, 
+                max_new_tokens=max_tokens,
+                do_sample=True,
+                guidance_scale=3.0
+            )
+
+        sampling_rate = model.config.audio_encoder.sampling_rate
+        audio_data = audio_values[0, 0].cpu().numpy()
+        return sampling_rate, audio_data
+    except Exception as e:
+        print(f"[!] Error: {str(e)}")
+        return None, None
+