Update app.py

app.py

@@ -1,9 +1,146 @@
-import gradio as gr
-from VitsModelSplit.vits_model import VitsModel
+import torch
+import torch.onnx
+import onnx
 from VitsModelSplit.vits_model_only_d import  Vits_models_only_decoder
-def greet(name):
-   
-    model = Vits_models_only_decoder.from_pretrained("wasmdashai/vits-ar-sa-huba",token=name)
-    return "Hello " + name + "!!"
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch(share=True)
+from VitsModelSplit.vits_model import VitsModel
+import gradio as gr
+
+class OnnxModelConverter:
+    def __init__(self):
+        self.model = None
+    def download_file(self,file_path):
+        ff= gr.File(value=file_path, visible=True)
+        file_url = ff.value['url']
+        return file_url
+
+    def convert(self, model_name, token, onnx_filename, conversion_type):
+        """
+        Main function to handle different types of model conversions.
+
+        Args:
+            model_name (str): Name of the model to convert.
+            token (str): Access token for loading the model.
+            onnx_filename (str): Desired filename for the ONNX output.
+            conversion_type (str): Type of conversion ('decoder', 'only_decoder', or 'full_model').
+
+        Returns:
+            str: The path to the generated ONNX file.
+        """
+        if conversion_type == "decoder":
+            return self.convert_decoder(model_name, token, onnx_filename)
+        elif conversion_type == "only_decoder":
+            return self.convert_only_decoder(model_name, token, onnx_filename)
+        elif conversion_type == "full_model":
+            return self.convert_full_model(model_name, token, onnx_filename)
+        else:
+            raise ValueError("Invalid conversion type. Choose from 'decoder', 'only_decoder', or 'full_model'.")
+
+    def convert_decoder(self, model_name, token, onnx_filename):
+        """
+        Converts only the decoder part of the Vits model to ONNX format.
+
+        Args:
+            model_name (str): Name of the model to convert.
+            token (str): Access token for loading the model.
+            onnx_filename (str): Desired filename for the ONNX output.
+
+        Returns:
+            str: The path to the generated ONNX file.
+        """
+        model = VitsModel.from_pretrained(model_name, token=token)
+        onnx_file = f"/tmp/{onnx_filename}.onnx"
+        vocab_size = model.text_encoder.embed_tokens.weight.size(0)
+        example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)
+
+        torch.onnx.export(
+            model,
+            example_input,
+            onnx_file,
+            opset_version=11,
+            input_names=['input'],
+            output_names=['output'],
+            dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
+        )
+
+        return onnx_file
+
+
+    def convert_only_decoder(self, model_name, token, onnx_filename):
+        """
+        Converts only the decoder part of the Vits model to ONNX format.
+
+        Args:
+            model_name (str): Name of the model to convert.
+            token (str): Access token for loading the model.
+            onnx_filename (str): Desired filename for the ONNX output.
+
+        Returns:
+            str: The path to the generated ONNX file.
+        """
+        model = Vits_models_only_decoder.from_pretrained(model_name, token=token)
+        onnx_file = f"/tmp/{onnx_filename}.onnx"
+
+        vocab_size = model.text_encoder.embed_tokens.weight.size(0)
+        example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)
+
+        torch.onnx.export(
+            model,
+            example_input,
+            onnx_file,
+            opset_version=11,
+            input_names=['input'],
+            output_names=['output'],
+            dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
+        )
+
+        return onnx_file
+
+    def convert_full_model(self, model_name, token, onnx_filename):
+        """
+        Converts the full Vits model (including encoder and decoder) to ONNX format.
+
+        Args:
+            model_name (str): Name of the model to convert.
+            token (str): Access token for loading the model.
+            onnx_filename (str): Desired filename for the ONNX output.
+
+        Returns:
+            str: The path to the generated ONNX file.
+        """
+        model = VitsModel.from_pretrained(model_name, token=token)
+        onnx_file = f"/tmp/{onnx_filename}.onnx"
+
+        vocab_size = model.text_encoder.embed_tokens.weight.size(0)
+        example_input = torch.randint(0, vocab_size, (1, 100), dtype=torch.long)
+
+        torch.onnx.export(
+            model,
+            example_input,
+            onnx_file,
+            opset_version=11,
+            input_names=['input'],
+            output_names=['output'],
+            dynamic_axes={'input': {0: 'batch_size', 1: 'sequence_length'}, 'output': {0: 'batch_size'}}
+        )
+
+        return onnx_file
+    def starrt(self):
+        with gr.Blocks() as demo:
+            with gr.Row():
+                  with gr.Column():
+                    text_n_model=gr.Textbox(label="name model")
+                    text_n_token=gr.Textbox(label="token")
+                    text_n_onxx=gr.Textbox(label="name model onxx")
+                    choice = gr.Dropdown(choices=["decoder", "only_decoder", "full_model"], label="My Dropdown")
+
+                  with gr.Column():
+
+                    btn=gr.Button("convert")
+                    label=gr.Label("return  name model onxx")
+                    btn.click(self.convert,[text_n_model,text_n_token,text_n_onxx,choice],[outputs=gr.File(label="Download File")])
+                    #choice.change(fn=function_change, inputs=choice, outputs=label)
+        return demo
+c=OnnxModelConverter()
+cc=c.starrt()
+cc.launch(share=True)
+