Update

app.py

@@ -18,7 +18,7 @@ transform = get_transform(image_size=image_size)
 tag2text_model = tag2text(pretrained="tag2text_swin_14m.pth", image_size=image_size, vit='swin_b').eval().to(device)
 
 
-def generate_music(raw_image):
+def generate_music(raw_image, audio_length):
     raw_image = Image.fromarray(raw_image)
     image = transform(raw_image).unsqueeze(0).to(device)
     res = inference_tag2text(image, tag2text_model)
@@ -37,8 +37,11 @@ def generate_music(raw_image):
         return_tensors="pt",
     )
 
-    audio_values = model.generate(**inputs, max_new_tokens=256)
     sampling_rate = model.audio_encoder.config.sampling_rate
+    frame_rate = model.audio_encoder.config.frame_rate
+    max_new_tokens = int(frame_rate * audio_length)
+    audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens)
+
     target_dtype = np.int16
     max_range = np.iinfo(target_dtype).max
     audio_values = audio_values[0, 0].numpy()
@@ -49,7 +52,10 @@ iface = gr.Interface(
     fn=generate_music,
     title=title,
     description=description,
-    inputs=gr.Image(label="Painting"),
+    inputs=[
+        gr.Image(label="Painting"),
+        gr.Slider(5, 30, value=15, step=1, label="Audio length(sec)")
+    ],
     outputs=gr.Audio(label='Generated Music'))
 
 iface.launch()