Update README.md

README.md

@@ -7,7 +7,93 @@ tags: []
 
 <!-- Provide a quick summary of what the model is/does. -->
 
-
+## Code to create model
+```py
+import torch
+from transformers import MimiConfig, MimiModel, AutoProcessor
+
+model_id = 'kyutai/mimi'
+config = MimiConfig.from_pretrained(
+    model_id,
+    intermediate_size=64,
+    hidden_size=16,
+    num_hidden_layers=2,
+    num_key_value_heads=2,
+    upsample_groups=16,
+    num_filters=8,
+    codebook_dim=8,
+    vector_quantization_hidden_dimension=8,
+    codebook_size=32,
+)
+
+# Create model and randomize all weights
+model = MimiModel(config)
+
+torch.manual_seed(0) # Set for reproducibility
+for name, param in model.named_parameters():
+    param.data = torch.randn_like(param)
+
+processor = AutoProcessor.from_pretrained(model_id)
+```
+
+## ONNX conversion code
+```py
+import torch
+import torch.nn as nn
+from transformers import MimiModel
+
+class MimiEncoder(nn.Module):
+    def __init__(self, model):
+        super(MimiEncoder, self).__init__()
+        self.model = model
+
+    def forward(self, input_values, padding_mask=None):
+        return self.model.encode(input_values, padding_mask=padding_mask).audio_codes
+
+class MimiDecoder(nn.Module):
+    def __init__(self, model):
+        super(MimiDecoder, self).__init__()
+        self.model = model
+
+    def forward(self, audio_codes, padding_mask=None):
+        return self.model.decode(audio_codes, padding_mask=padding_mask).audio_values
+
+model = MimiModel.from_pretrained("hf-internal-testing/tiny-random-MimiModel")
+encoder = MimiEncoder(model)
+decoder = MimiDecoder(model)
+
+dummy_encoder_inputs = torch.randn((5, 1, 82500))
+torch.onnx.export(
+    encoder,
+    dummy_encoder_inputs,
+    "encoder_model.onnx",
+    export_params=True,
+    opset_version=14,
+    do_constant_folding=True,
+    input_names=['input_values'],
+    output_names=['audio_codes'],
+    dynamic_axes={
+        'input_values': {0: 'batch_size', 1: 'num_channels', 2: 'sequence_length'},
+        'audio_codes': {0: 'batch_size', 2: 'codes_length'},
+    },
+)
+
+dummy_decoder_inputs = torch.randint(8, (4, 32, 91))
+torch.onnx.export(
+    decoder,
+    dummy_decoder_inputs,
+    "decoder_model.onnx",
+    export_params=True,
+    opset_version=14,
+    do_constant_folding=True,
+    input_names=['audio_codes'],
+    output_names=['audio_values'],
+    dynamic_axes={
+        'audio_codes': {0: 'batch_size', 2: 'codes_length'},
+        'audio_values': {0: 'batch_size', 1: 'num_channels', 2: 'sequence_length'},
+    },
+)
+```
 
 ## Model Details