Upload folder using huggingface_hub

config.json

@@ -6,11 +6,11 @@
     "AutoConfig": "model.INF5Config",
     "AutoModel": "model.INF5Model"
   },
-  "ckpt_path": "checkpoints/model_last.pt",
+  "ckpt_path": "model_last.pt",
   "model_type": "inf5",
   "remove_sil": true,
   "speed": 1.0,
   "torch_dtype": "float32",
   "transformers_version": "4.49.0",
-  "vocab_path": "checkpoints/vocab.txt"
-}
+  "vocab_path": "vocab.txt"
+}

model.py

@@ -19,51 +19,25 @@ import io
 from pydub import AudioSegment, silence
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-import os
 
 class INF5Config(PretrainedConfig):
     model_type = "inf5"
 
-    def __init__(self, ckpt_path: str = "checkpoints/model_best.pt", vocab_path: str = "checkpoints/vocab.txt", 
+    def __init__(self, ckpt_repo_id: str = None, vocab_repo_id: str = None,
+                 ckpt_filename: str = None, vocab_filename: str = "vocab.txt",
                  speed: float = 1.0, remove_sil: bool = True, **kwargs):
         super().__init__(**kwargs)
-        self.ckpt_path = ckpt_path
-        self.vocab_path = vocab_path
+        # If not specified, use the model's own repo for both
+        self.ckpt_repo_id = ckpt_repo_id
+        self.vocab_repo_id = vocab_repo_id
+        self.ckpt_filename = ckpt_filename
+        self.vocab_filename = vocab_filename
         self.speed = speed
         self.remove_sil = remove_sil
 
 class INF5Model(PreTrainedModel):
     config_class = INF5Config
 
-    # def __init__(self, config):
-    #     super().__init__(config)
-    #     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    #     # Load vocoder
-    #     self.vocoder = torch.compile(load_vocoder(vocoder_name="vocos", is_local=False, device=device))
-
-    #     # Download and load model weights
-    #     # safetensors_path = hf_hub_download(config.name_or_path, filename="model.safetensors")
-    #     # print(f"Loading model weights from {safetensors_path} (safetensors)...")
-    #     # state_dict = load_file(safetensors_path, device=str(device))
-
-    #     # Download vocab.txt from HF Hub
-    #     vocab_path = hf_hub_download(config.name_or_path, filename="checkpoints/vocab.txt")
-
-    #     self.ema_model = torch.compile(load_model(
-    #             DiT,
-    #             dict(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4),
-    #             mel_spec_type="vocos",
-    #             vocab_file=vocab_path,
-    #             device=device
-    #         )
-    #     )
-
-
-    #     # # Load state dict into model
-    #     # self.ema_model.load_state_dict(state_dict, strict=False)
-
-
     def __init__(self, config):
         super().__init__(config)
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -73,10 +47,18 @@ class INF5Model(PreTrainedModel):
             load_vocoder(vocoder_name="vocos", is_local=False, device=device)
         )
 
+        # Determine which repo to use for vocab
+        # Default to the model's own repo if not specified
+        vocab_repo = config.vocab_repo_id or config.name_or_path
+        
         # Download vocab.txt from HF Hub
-        vocab_path = hf_hub_download(config.name_or_path, filename="checkpoints/vocab.txt")
+        vocab_path = hf_hub_download(repo_id=vocab_repo, filename=config.vocab_filename)
+
+        # Determine which repo to use for checkpoint
+        ckpt_repo = config.ckpt_repo_id or config.name_or_path
 
         ckpt_candidates = [
+            "model_last.pt",  # Try this first since it's in your repo
             "checkpoints/model.safetensors",
             "model.safetensors",
             "checkpoints/pytorch_model.bin",
@@ -87,12 +69,15 @@ class INF5Model(PreTrainedModel):
             "checkpoint.pt"
         ]
 
+        # If a specific checkpoint filename is provided, use only that
+        if config.ckpt_filename:
+            ckpt_candidates = [config.ckpt_filename]
+
         ckpt_path = None
-        from huggingface_hub import hf_hub_download, RepositoryNotFoundError, hf_hub_download
 
         for fname in ckpt_candidates:
             try:
-                ckpt_path = hf_hub_download(repo_id=config.name_or_path, filename=fname)
+                ckpt_path = hf_hub_download(repo_id=ckpt_repo, filename=fname)
                 print(f"Found checkpoint on hub: {fname} -> {ckpt_path}")
                 break
             except Exception as e:
@@ -106,7 +91,7 @@ class INF5Model(PreTrainedModel):
                 "Could not find a checkpoint file on the Hub. "
                 "Tried: " + ", ".join(ckpt_candidates) + ".\n"
                 "If your checkpoint is stored under a different path or name, "
-                "update ckpt_candidates or pass the path via config (e.g. config.ckpt_path). "
+                "update ckpt_candidates or pass the path via config (e.g. config.ckpt_filename). "
                 "If the file is >5GB, ensure Git LFS is enabled for the repo (hf lfs-enable-largefiles)."
             )
 
@@ -135,7 +120,6 @@ class INF5Model(PreTrainedModel):
             text (str): The text to be synthesized.
             ref_audio_path (str): Path to the reference audio file.
             ref_text (str): The reference text.
-
         Returns:
             np.array: Generated waveform.
         """
@@ -189,41 +173,41 @@ class INF5Model(PreTrainedModel):
 
 
 if __name__ == '__main__':
-    model = INF5Model(INF5Config(ckpt_path="checkpoints/model_best.pt", vocab_path="checkpoints/vocab.txt"))
+    model = INF5Model(INF5Config())
     model.save_pretrained("INF5")
     model.config.save_pretrained("INF5")
-
-    import numpy as np
-    import soundfile as sf
-    from transformers import AutoConfig, AutoModel
     
-    AutoConfig.register("inf5", INF5Config)
-    AutoModel.register(INF5Config, INF5Model)
 
-    model = AutoModel.from_pretrained("INF5")
-    audio = model("नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.", 
-                  ref_audio_path="prompts/PAN_F_HAPPY_00001.wav",
-                  ref_text="भਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।")
+    # import numpy as np
+    # import soundfile as sf
+    # from transformers import AutoConfig, AutoModel
     
-    if audio.dtype == np.int16:
-        audio = audio.astype(np.float32) / 32768.0  
-    sf.write("samples/namaste.wav", np.array(audio, dtype=np.float32), samplerate=24000)
-
-    from huggingface_hub import HfApi
-
-    repo_id = "svp19/INF5"  # Change to your HF repo
-
-    # Upload model directory to HF
-    api = HfApi()
-    api.upload_folder(
-        folder_path="INF5",
-        repo_id=repo_id,
-        repo_type="model"
-    )
-    print(f"Model pushed to https://huggingface.co/{repo_id} 🚀")
-
-    print("Verify Upload")
-    from transformers import AutoModel
-    model = AutoModel.from_pretrained(repo_id)
-    print("Success")
+    # AutoConfig.register("inf5", INF5Config)
+    # AutoModel.register(INF5Config, INF5Model)
 
+    # model = AutoModel.from_pretrained("INF5")
+    # audio = model("नमस्ते! संगीत की तरह जीवन भी खूबसूरत होता है, बस इसे सही ताल में जीना आना चाहिए.", 
+    #               ref_audio_path="prompts/PAN_F_HAPPY_00001.wav",
+    #               ref_text="भਹੰਪੀ ਵਿੱਚ ਸਮਾਰਕਾਂ ਦੇ ਭਵਨ ਨਿਰਮਾਣ ਕਲਾ ਦੇ ਵੇਰਵੇ ਗੁੰਝਲਦਾਰ ਅਤੇ ਹੈਰਾਨ ਕਰਨ ਵਾਲੇ ਹਨ, ਜੋ ਮੈਨੂੰ ਖੁਸ਼ ਕਰਦੇ  ਹਨ।")
+    
+    # if audio.dtype == np.int16:
+    #     audio = audio.astype(np.float32) / 32768.0  
+    # sf.write("samples/namaste.wav", np.array(audio, dtype=np.float32), samplerate=24000)
+
+    # from huggingface_hub import HfApi
+
+    # repo_id = "svp19/INF5"  # Change to your HF repo
+
+    # # Upload model directory to HF
+    # api = HfApi()
+    # api.upload_folder(
+    #     folder_path="INF5",
+    #     repo_id=repo_id,
+    #     repo_type="model"
+    # )
+    # print(f"Model pushed to https://huggingface.co/{repo_id}")
+
+    # print("Verify Upload")
+    # from transformers import AutoModel
+    # model = AutoModel.from_pretrained(repo_id)
+    # print("Success")