fix quantized model saving

app.py

@@ -32,7 +32,7 @@ from optimum.intel import (
     OVWeightQuantizationConfig,
 )
 
-def process_model(
+def quantize_model(
     model_id: str,
     dtype: str,
     calibration_dataset: str,
@@ -49,6 +49,8 @@ def process_model(
     new_repo_id = f"{username}/{model_name}-openvino-{dtype}"
     task = TasksManager.infer_task_from_model(model_id)
     library_name = TasksManager.infer_library_from_model(model_id)
+    # task = TasksManager.infer_task_from_model(model_id, token=oauth_token.token)
+    # library_name = TasksManager.infer_library_from_model(model_id, token=oauth_token.token)
 
     if task not in _HEAD_TO_AUTOMODELS:
         raise ValueError(
@@ -82,41 +84,52 @@ def process_model(
     )
 
     api = HfApi(token=oauth_token.token)
-
     if api.repo_exists(new_repo_id) and not overwritte:
         raise Exception(f"Model {new_repo_id} already exist, please set overwritte=True to push on an existing repo")
 
     with TemporaryDirectory() as d:
         folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models"))
         os.makedirs(folder)
+
         try:
             api.snapshot_download(repo_id=model_id, local_dir=folder, allow_patterns=["*.json"])
-
             ov_model = eval(auto_model_class).from_pretrained(
                 model_id,
                 export=export,
-                quantization_config=quantization_config,
+                cache_dir=folder,
+                token=oauth_token.token,
+                quantization_config=quantization_config
             )
             ov_model.save_pretrained(folder)
-
             new_repo_url = api.create_repo(repo_id=new_repo_id, exist_ok=True, private=private_repo)
             new_repo_id = new_repo_url.repo_id
             print("Repo created successfully!", new_repo_url)
 
             folder = Path(folder)
-            folder_parts = len(folder.parts)
-            for file_path in folder.glob("**/*"):
-                name = Path(*file_path.parts[folder_parts:])
-                if not file_path.is_file() or any(part_name.startswith(".") for part_name in name.parts):
+            for dir_name in (
+                "",
+                "vae_encoder",
+                "vae_decoder",
+                "text_encoder",
+                "text_encoder_2",
+                "unet",
+                "tokenizer",
+                "tokenizer_2",
+                "scheduler",
+                "feature_extractor",
+            ):
+                if not (folder / dir_name).is_dir():
                     continue
-                try:
-                    api.upload_file(
-                        path_or_fileobj=file_path,
-                        path_in_repo=str(name),
-                        repo_id=new_repo_id,
-                    )
-                except Exception as e:
-                    raise Exception(f"Error uploading file {file_path}: {e}")
+                for file_path in (folder / dir_name).iterdir():
+                    if file_path.is_file():
+                        try:
+                            api.upload_file(
+                                path_or_fileobj=file_path,
+                                path_in_repo=os.path.join(dir_name, file_path.name),
+                                repo_id=new_repo_id,
+                            )
+                        except Exception as e:
+                            raise Exception(f"Error uploading file {file_path}: {e}")
 
             try:
                 card = ModelCard.load(model_id, token=oauth_token.token)
@@ -220,7 +233,7 @@ overwritte = gr.Checkbox(
     info="Push files on existing repo potentially overwriting existing files",
 )
 interface = gr.Interface(
-    fn=process_model,
+    fn=quantize_model,
     inputs=[
         model_id,
         dtype,

requirements.txt

@@ -1,5 +1,9 @@
 huggingface_hub==0.23.4
-optimum[diffusers]==1.21.2
-optimum-intel[openvino]==1.18.1
-gradio[oauth]>=4.28.0
+gradio[oauth]>=4.37.2
 gradio_huggingfacehub_search==0.0.6
+transformers==4.42.4
+diffusers==0.29.1
+optimum==1.21.2
+optimum-intel==1.18.1
+openvino
+nncf