musicgpt 0.2.1

Generate music samples from natural language prompt locally with your own computer
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use std::collections::HashMap;

pub struct MusicGenInputs {
    inputs: HashMap<String, ort::DynValue>,
    pub use_cache_branch: bool,
}

impl MusicGenInputs {
    pub fn new() -> Self {
        Self {
            inputs: HashMap::new(),
            use_cache_branch: false,
        }
    }
    pub fn encoder_attention_mask<T, E>(&mut self, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs
            .insert("encoder_attention_mask".to_string(), v.try_into()?);
        Ok(())
    }

    pub fn input_ids<T, E>(&mut self, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs.insert("input_ids".to_string(), v.try_into()?);
        Ok(())
    }

    pub fn encoder_hidden_states<T, E>(&mut self, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs
            .insert("encoder_hidden_states".to_string(), v.try_into()?);
        Ok(())
    }
    
    pub fn remove_encoder_hidden_states(&mut self) {
        self.inputs.remove("encoder_hidden_states");
    }

    pub fn past_key_value_decoder_key<T, E>(&mut self, i: usize, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs
            .insert(format!("past_key_values.{i}.decoder.key"), v.try_into()?);
        Ok(())
    }

    pub fn past_key_value_decoder_value<T, E>(&mut self, i: usize, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs
            .insert(format!("past_key_values.{i}.decoder.value"), v.try_into()?);
        Ok(())
    }

    pub fn past_key_value_encoder_key<T, E>(&mut self, i: usize, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs
            .insert(format!("past_key_values.{i}.encoder.key"), v.try_into()?);
        Ok(())
    }

    pub fn past_key_value_encoder_value<T, E>(&mut self, i: usize, v: T) -> Result<(), E>
    where
        ort::DynValue: TryFrom<T, Error = E>,
    {
        self.inputs
            .insert(format!("past_key_values.{i}.encoder.value"), v.try_into()?);
        Ok(())
    }

    pub fn use_cache_branch(&mut self, value: bool) {
        self.use_cache_branch = value;
        self.inputs.insert(
            "use_cache_branch".to_string(),
            ort::Tensor::from_array(([1], vec![value]))
                .unwrap()
                .into_dyn(),
        );
    }

    pub fn ort(&self) -> ort::SessionInputs {
        ort::SessionInputs::ValueMap(
            self.inputs
                .iter()
                .map(|e| (e.0.to_string().into(), e.1.view().into()))
                .collect::<Vec<_>>(),
        )
    }
}