From c564f828212e3950837d5d1af7e99d0d1697f317 Mon Sep 17 00:00:00 2001
From: Roberts Slisans <roberts.slisans@gmail.com>
Date: Sat, 15 Feb 2025 12:42:04 +0200
Subject: [PATCH] fix cfg_scale type in generate_diffusion_cond

---
 stable_audio_tools/inference/generation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/stable_audio_tools/inference/generation.py b/stable_audio_tools/inference/generation.py
index 9944dd86..4e81871f 100644
--- a/stable_audio_tools/inference/generation.py
+++ b/stable_audio_tools/inference/generation.py
@@ -91,7 +91,7 @@ def generate_diffusion_uncond(
 def generate_diffusion_cond(
         model,
         steps: int = 250,
-        cfg_scale=6,
+        cfg_scale: float = 6,
         conditioning: dict = None,
         conditioning_tensors: tp.Optional[dict] = None,
         negative_conditioning: dict = None,
@@ -427,4 +427,4 @@ def build_mask(sample_size, mask_args):
     if marination > 0:        
         mask = mask * (1-marination) 
     #print(mask)
-    return mask
\ No newline at end of file
+    return mask