Update app.py

app.py

@@ -34,97 +34,234 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 #                         discriminator=False,
 #                         duration=False
 # )
-class model_onxx:
-    def __init__(self):
-        self.model=None
-        self.n_onxx=""
-        self.storage_dir = "uploads"  
-        pass
+# class model_onxx:
+#     def __init__(self):
+#         self.model=None
+#         self.n_onxx=""
+#         self.storage_dir = "uploads"  
+#         pass
      
 
  
 
 
+#     def download_file(self,file_path):
+#         ff= gr.File(value=file_path, visible=True) 
+#         file_url = ff.value['url']  
+#         return file_url
+#     def function_change(self,n_model,token,n_onxx,choice):
+#         if choice=="decoder":
+
+#             V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
+#         elif choice=="all only decoder":
+#             V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
+#         else:
+#             V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
+#         return V
+
+#     def install_model(self,n_model,token,n_onxx):
+#         self.n_onxx=n_onxx
+#         self.model= VitsModel.from_pretrained(n_model,token=token)
+#         return self.model
+#     def convert_model_decoder_onxx(self,n_model,token,namemodelonxx):
+#         self.model= VitsModel.from_pretrained(n_model,token=token)
+#         x=f"/tmp/{namemodelonxx}.onnx"
+#         return x
+#     def convert_to_onnx_only_decoder(self,n_model,token,namemodelonxx):
+#           model=VitsModel.from_pretrained(n_model,token=token)
+#           x=f"/tmp/{namemodelonxx}.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,  # The model to be exported
+#               example_input,  # Example input for the model
+#               x,# The filename for the exported ONNX model
+#               opset_version=11,  # Use an appropriate ONNX opset version
+#               input_names=['input'],  # Name of the input layer
+#               output_names=['output'],  # Name of the output layer
+#               dynamic_axes={
+#                   'input': {0: 'batch_size', 1: 'sequence_length'},  # Dynamic axes for variable-length inputs
+#                   'output': {0: 'batch_size'}
+#               }
+#            )
+#           return x
+           
+#     def convert_to_onnx_all(self,n_model,token ,namemodelonxx):
+
+#           model=VitsModel.from_pretrained(n_model,token=token)
+#           x=f"dowload_file/{namemodelonxx}.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,  # The model to be exported
+#               example_input,  # Example input for the model
+#               x,  # The filename for the exported ONNX model
+#               opset_version=11,  # Use an appropriate ONNX opset version
+#               input_names=['input'],  # Name of the input layer
+#               output_names=['output'],  # Name of the output layer
+#               dynamic_axes={
+#                   'input': {0: 'batch_size', 1: 'sequence_length'},  # Dynamic axes for variable-length inputs
+#                   'output': {0: 'batch_size'}
+#               }
+#            )
+#           return x
+#     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", "all anoly decoder", "All"], label="My Dropdown")
+
+#                   with gr.Column():
+
+#                     btn=gr.Button("convert")
+#                     label=gr.Label("return  name model onxx")
+#                     btn.click(self.function_change,[text_n_model,text_n_token,text_n_onxx,choice],[gr.File(label="Download File")])
+#                     #choice.change(fn=function_change, inputs=choice, outputs=label)
+#         #return demo
+# c=model_onxx()
+
+#3333333333333333333333333333
+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 function_change(self,n_model,token,n_onxx,choice):
-        if choice=="decoder":
 
-            V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
-        elif choice=="all only decoder":
-            V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
+    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:
-            V=self.convert_to_onnx_only_decoder(n_model,token,n_onxx)
-        return V
-
-    def install_model(self,n_model,token,n_onxx):
-        self.n_onxx=n_onxx
-        self.model= VitsModel.from_pretrained(n_model,token=token)
-        return self.model
-    def convert_model_decoder_onxx(self,n_model,token,namemodelonxx):
-        self.model= VitsModel.from_pretrained(n_model,token=token)
-        x=f"/tmp/{namemodelonxx}.onnx"
-        return x
-    def convert_to_onnx_only_decoder(self,n_model,token,namemodelonxx):
-          model=VitsModel.from_pretrained(n_model,token=token)
-          x=f"/tmp/{namemodelonxx}.onnx"
+            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'}}
+        )
         
-          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,  # The model to be exported
-              example_input,  # Example input for the model
-              x,# The filename for the exported ONNX model
-              opset_version=11,  # Use an appropriate ONNX opset version
-              input_names=['input'],  # Name of the input layer
-              output_names=['output'],  # Name of the output layer
-              dynamic_axes={
-                  'input': {0: 'batch_size', 1: 'sequence_length'},  # Dynamic axes for variable-length inputs
-                  'output': {0: 'batch_size'}
-              }
-           )
-          return x
-           
-    def convert_to_onnx_all(self,n_model,token ,namemodelonxx):
-
-          model=VitsModel.from_pretrained(n_model,token=token)
-          x=f"dowload_file/{namemodelonxx}.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,  # The model to be exported
-              example_input,  # Example input for the model
-              x,  # The filename for the exported ONNX model
-              opset_version=11,  # Use an appropriate ONNX opset version
-              input_names=['input'],  # Name of the input layer
-              output_names=['output'],  # Name of the output layer
-              dynamic_axes={
-                  'input': {0: 'batch_size', 1: 'sequence_length'},  # Dynamic axes for variable-length inputs
-                  'output': {0: 'batch_size'}
-              }
-           )
-          return x
+        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.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", "all anoly decoder", "All"], label="My Dropdown")
-
+                    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.function_change,[text_n_model,text_n_token,text_n_onxx,choice],[gr.File(label="Download File")])
+                    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=model_onxx()
-#cc=c.starrt()
+        return demo
+c=OnnxModelConverter()
 ###############################################################
 Lst=['input_ids',
  'attention_mask',