Update app.py

app.py

@@ -10,9 +10,12 @@ import torchaudio
 from generator import Segment, load_csm_1b
 from watermarking import watermark
 
+# Set a default watermark key to avoid the NoneType error
+if "WATERMARK_KEY" not in os.environ:
+    os.environ["WATERMARK_KEY"] = "0 0 0 0"  # Default placeholder
+
 # Simplified environment variables handling
 gpu_timeout = int(os.getenv("GPU_TIMEOUT", 60))
-CSM_1B_HF_WATERMARK = list(map(int, os.getenv("WATERMARK_KEY", "0 0 0").split(" ")))
 
 SPACE_INTRO_TEXT = """\
 # Sesame CSM 1B
@@ -179,9 +182,12 @@ def _infer(
     audio_tensors = [segment.audio for segment in generated_segments]
     audio_tensor = torch.cat(audio_tensors, dim=0)
 
+    # Get the watermark key from environment
+    watermark_key = list(map(int, os.getenv("WATERMARK_KEY").split(" ")))
+    
     # Watermarking
     audio_tensor, wm_sample_rate = watermark(
-        generator._watermarker, audio_tensor, generator.sample_rate, CSM_1B_HF_WATERMARK
+        generator._watermarker, audio_tensor, generator.sample_rate, watermark_key
     )
     audio_tensor = torchaudio.functional.resample(
         audio_tensor, orig_freq=wm_sample_rate, new_freq=generator.sample_rate