clean

app.py

@@ -72,52 +72,27 @@ def export(
         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=True)
+            api.snapshot_download(repo_id=model_id, local_dir=folder, allow_patterns=["*.json"])
+            ov_model = eval(auto_model_class).from_pretrained(model_id, export=True, cache_dir=folder)
             ov_model.save_pretrained(folder)
-
-            """
-            if not isinstance(ov_model, OVStableDiffusionPipelineBase):
-                model = TasksManager.get_model_from_task(task, model_id)
-                exporter_config_class = TasksManager.get_exporter_config_constructor(
-                    exporter="openvino",
-                    model=model,
-                    task=task,
-                    model_name=model_id,
-                    model_type=model.config.model_type.replace("_", "-"),
-                )
-                openvino_config = exporter_config_class(model.config)
-                inputs = openvino_config.generate_dummy_inputs(framework="pt")
-                ov_outputs = ov_model(**inputs)
-                outputs = model(**inputs)
-
-                for output_name in ov_outputs:
-                    if isinstance(outputs, torch.Tensor) and not torch.allclose(outputs[output_name], ov_outputs[output_name], atol=1e-3):
-                        raise ValueError(
-                            "The exported model does not have the same outputs as the original model. Export interrupted."
-                        )
-            """
             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"):
+                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)