From 9069624d5e1ae303adecb87fd7c4d14072a34fde Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Tue, 12 Nov 2024 14:32:45 +0530
Subject: [PATCH 01/25] initial commit

---
 .../components/import_model/spec.yaml         | 105 +++++++++++++++++-
 1 file changed, 103 insertions(+), 2 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 2ff88ce5f8..d6632002f0 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: 0.0.39
+version: 0.0.40
 
 # Pipeline inputs
 inputs:
@@ -58,6 +58,12 @@ inputs:
     optional: false
     description: Framework from which model is imported from.
 
+  MMLab:
+    type: boolean
+    default: false
+    optional: true
+    description: Set to true for MMLab models.
+
   vllm_enabled:
     type: boolean
     description: Enable vllm in the converted model 
@@ -212,13 +218,22 @@ outputs:
     description: Output path for the converted MLflow model
     type: mlflow_model
 
+  mlflow_model_folder_mmd:
+    description: Output path for the converted MLflow model MMD
+    type: mlflow_model
+
   model_registration_details:
     description: Output folder with a file which captures transformations applied above and registration details in JSON file
     type: uri_folder
 
+  model_registration_details_mmd:
+    description: Output folder with a file which captures transformations applied above and registration details in JSON file MMD
+    type: uri_folder
+
 jobs:
   validation_trigger_import:
     component: azureml:validation_trigger_import:0.0.13
+    # component: azureml://registries/azureml/components/validation_trigger_import/versions/0.0.11
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -251,8 +266,15 @@ jobs:
       validation_info:
         type: uri_file
 
+  is_mmlab:
+    type: if_else
+    condition: ${{parent.inputs.MMLab}} 
+    true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
+    false_block: ${{parent.jobs.download_model}}
+
   download_model:
     component: azureml:download_model:0.0.29
+    # component: azureml://registries/azureml/components/download_model/versions/0.0.28
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -270,8 +292,24 @@ jobs:
       model_output:
         type: uri_folder
 
+  mmdetection_image_objectdetection_instancesegmentation_model_import:
+    component: azureml:mmdetection_image_objectdetection_instancesegmentation_model_import/versions/0.0.19
+    compute: ${{parent.inputs.compute}}
+    resources:
+      instance_type: '${{parent.inputs.instance_type}}'
+    identity:
+      type: user_identity
+    inputs:
+      model_family: 'MmDetectionImage'
+      model_name: ${{parent.inputs.model_id}}
+      download_from_source: False
+    outputs:
+      output_dir:
+        type: uri_file
+
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.33
+    component: azureml:convert_model_to_mlflow:0.0.31
+    # component: azureml://registries/azureml/components/convert_model_to_mlflow/versions/0.0.31
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -298,8 +336,37 @@ jobs:
       mlflow_model_folder:
         type: mlflow_model
 
+  convert_model_to_mlflow_mmd:
+    component: azureml:convert_model_to_mlflow:0.0.31
+    # component: azureml://registries/azureml/components/convert_model_to_mlflow/versions/0.0.31
+    compute: ${{parent.inputs.compute}}
+    resources:
+      instance_type: '${{parent.inputs.instance_type}}'
+    identity:
+      type: user_identity
+    inputs:
+      task_name: ${{parent.inputs.task_name}}
+      model_flavor: ${{parent.inputs.model_flavor}}
+      vllm_enabled: ${{parent.inputs.vllm_enabled}}
+      license_file_path: ${{parent.inputs.license_file_path}}
+      model_framework: ${{parent.inputs.model_framework}}
+      model_path:  ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
+      hf_config_args: ${{parent.inputs.hf_config_args}}
+      hf_tokenizer_args:  ${{parent.inputs.hf_tokenizer_args}}
+      hf_model_args: ${{parent.inputs.hf_model_args}}
+      hf_pipeline_args: ${{parent.inputs.hf_pipeline_args}}
+      hf_config_class:  ${{parent.inputs.hf_config_class}}
+      hf_model_class:  ${{parent.inputs.hf_model_class}}
+      hf_tokenizer_class:  ${{parent.inputs.hf_tokenizer_class}}
+      hf_use_experimental_features: ${{parent.inputs.hf_use_experimental_features}}
+      extra_pip_requirements: ${{parent.inputs.extra_pip_requirements}}
+    outputs:
+      mlflow_model_folder:
+        type: mlflow_model
+
   mlflow_model_local_validation:
     component: azureml:mlflow_model_local_validation:0.0.16
+    # component: azureml://registries/azureml/components/mlflow_model_local_validation/versions/0.0.14
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -311,8 +378,23 @@ jobs:
     outputs:
       mlflow_model_folder: ${{parent.outputs.mlflow_model_folder}}
 
+  mlflow_model_local_validation_mmd:
+    component: azureml:mlflow_model_local_validation:0.0.16
+    # component: azureml://registries/azureml/components/mlflow_model_local_validation/versions/0.0.14
+    compute: ${{parent.inputs.compute}}
+    resources:
+      instance_type: '${{parent.inputs.instance_type}}'
+    inputs:
+      model_path: ${{parent.jobs.convert_model_to_mlflow_mmd.outputs.mlflow_model_folder}}
+      test_data_path: ${{parent.inputs.local_validation_test_data}}
+      column_rename_map: ${{parent.inputs.local_validation_column_rename_map}}
+      task_name: ${{parent.inputs.task_name}}
+    outputs:
+      mlflow_model_folder: ${{parent.outputs.mlflow_model_folder_mmd}}
+
   register_model:
     component: azureml:register_model:0.0.19
+    # component: azureml://registries/azureml/components/register_model/versions/0.0.17
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -330,5 +412,24 @@ jobs:
     outputs:
       registration_details_folder: ${{parent.outputs.model_registration_details}}
 
+  register_model_mmd:
+    component: azureml:register_model:0.0.19
+    # component: azureml://registries/azureml/components/register_model/versions/0.0.17
+    compute: ${{parent.inputs.compute}}
+    resources:
+      instance_type: '${{parent.inputs.instance_type}}'
+    identity:
+      type: user_identity
+    inputs:
+      model_name: ${{parent.inputs.custom_model_name}}
+      model_version: ${{parent.inputs.model_version}}
+      model_description: ${{parent.inputs.model_description}}
+      registry_name: ${{parent.inputs.registry_name}}
+      model_metadata: ${{parent.inputs.model_metadata}}
+      model_type: mlflow_model
+      model_path: ${{parent.jobs.mlflow_model_local_validation_mmd.outputs.mlflow_model_folder}}
+    outputs:
+      registration_details_folder: ${{parent.outputs.model_registration_details_mmd}}
+
 tags:
     Preview: ""

From 476a37f1d85a26359c6169a5a90c8e7356845022 Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Wed, 13 Nov 2024 14:22:16 +0530
Subject: [PATCH 02/25] Review comments

---
 .../convert_model_to_mlflow/spec.yaml         |  8 +-
 .../components/import_model/spec.yaml         | 86 +------------------
 .../src/run_model_preprocess.py               |  3 +-
 3 files changed, 12 insertions(+), 85 deletions(-)

diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
index 7d2a12753c..84dae892f0 100644
--- a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
+++ b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -1,7 +1,7 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: convert_model_to_mlflow
-version: 0.0.33
+version: 0.0.34
 type: command
 
 is_deterministic: True
@@ -158,6 +158,12 @@ inputs:
     mode: ro_mount
     optional: false
 
+  model_path_mmd:
+    type: uri_folder
+    description: Path to the MMD model.
+    mode: ro_mount
+    optional: false
+
   license_file_path:
     type: uri_file
     description: Path to the license file
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index d6632002f0..b416f2822e 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -58,12 +58,6 @@ inputs:
     optional: false
     description: Framework from which model is imported from.
 
-  MMLab:
-    type: boolean
-    default: false
-    optional: true
-    description: Set to true for MMLab models.
-
   vllm_enabled:
     type: boolean
     description: Enable vllm in the converted model 
@@ -218,18 +212,10 @@ outputs:
     description: Output path for the converted MLflow model
     type: mlflow_model
 
-  mlflow_model_folder_mmd:
-    description: Output path for the converted MLflow model MMD
-    type: mlflow_model
-
   model_registration_details:
     description: Output folder with a file which captures transformations applied above and registration details in JSON file
     type: uri_folder
 
-  model_registration_details_mmd:
-    description: Output folder with a file which captures transformations applied above and registration details in JSON file MMD
-    type: uri_folder
-
 jobs:
   validation_trigger_import:
     component: azureml:validation_trigger_import:0.0.13
@@ -266,12 +252,6 @@ jobs:
       validation_info:
         type: uri_file
 
-  is_mmlab:
-    type: if_else
-    condition: ${{parent.inputs.MMLab}} 
-    true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
-    false_block: ${{parent.jobs.download_model}}
-
   download_model:
     component: azureml:download_model:0.0.29
     # component: azureml://registries/azureml/components/download_model/versions/0.0.28
@@ -308,8 +288,8 @@ jobs:
         type: uri_file
 
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.31
-    # component: azureml://registries/azureml/components/convert_model_to_mlflow/versions/0.0.31
+    component: azureml:convert_model_to_mlflow:0.0.34
+    # component: azureml://registries/azureml-preview-test1/components/convert_model_to_mlflow/versions/0.0.33-groopa-test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -322,6 +302,7 @@ jobs:
       license_file_path: ${{parent.inputs.license_file_path}}
       model_framework: ${{parent.inputs.model_framework}}
       model_download_metadata: ${{parent.jobs.download_model.outputs.model_download_metadata}}
+      model_path_mmd:  ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
       model_path: ${{parent.jobs.download_model.outputs.model_output}}
       hf_config_args: ${{parent.inputs.hf_config_args}}
       hf_tokenizer_args:  ${{parent.inputs.hf_tokenizer_args}}
@@ -336,34 +317,6 @@ jobs:
       mlflow_model_folder:
         type: mlflow_model
 
-  convert_model_to_mlflow_mmd:
-    component: azureml:convert_model_to_mlflow:0.0.31
-    # component: azureml://registries/azureml/components/convert_model_to_mlflow/versions/0.0.31
-    compute: ${{parent.inputs.compute}}
-    resources:
-      instance_type: '${{parent.inputs.instance_type}}'
-    identity:
-      type: user_identity
-    inputs:
-      task_name: ${{parent.inputs.task_name}}
-      model_flavor: ${{parent.inputs.model_flavor}}
-      vllm_enabled: ${{parent.inputs.vllm_enabled}}
-      license_file_path: ${{parent.inputs.license_file_path}}
-      model_framework: ${{parent.inputs.model_framework}}
-      model_path:  ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
-      hf_config_args: ${{parent.inputs.hf_config_args}}
-      hf_tokenizer_args:  ${{parent.inputs.hf_tokenizer_args}}
-      hf_model_args: ${{parent.inputs.hf_model_args}}
-      hf_pipeline_args: ${{parent.inputs.hf_pipeline_args}}
-      hf_config_class:  ${{parent.inputs.hf_config_class}}
-      hf_model_class:  ${{parent.inputs.hf_model_class}}
-      hf_tokenizer_class:  ${{parent.inputs.hf_tokenizer_class}}
-      hf_use_experimental_features: ${{parent.inputs.hf_use_experimental_features}}
-      extra_pip_requirements: ${{parent.inputs.extra_pip_requirements}}
-    outputs:
-      mlflow_model_folder:
-        type: mlflow_model
-
   mlflow_model_local_validation:
     component: azureml:mlflow_model_local_validation:0.0.16
     # component: azureml://registries/azureml/components/mlflow_model_local_validation/versions/0.0.14
@@ -378,20 +331,6 @@ jobs:
     outputs:
       mlflow_model_folder: ${{parent.outputs.mlflow_model_folder}}
 
-  mlflow_model_local_validation_mmd:
-    component: azureml:mlflow_model_local_validation:0.0.16
-    # component: azureml://registries/azureml/components/mlflow_model_local_validation/versions/0.0.14
-    compute: ${{parent.inputs.compute}}
-    resources:
-      instance_type: '${{parent.inputs.instance_type}}'
-    inputs:
-      model_path: ${{parent.jobs.convert_model_to_mlflow_mmd.outputs.mlflow_model_folder}}
-      test_data_path: ${{parent.inputs.local_validation_test_data}}
-      column_rename_map: ${{parent.inputs.local_validation_column_rename_map}}
-      task_name: ${{parent.inputs.task_name}}
-    outputs:
-      mlflow_model_folder: ${{parent.outputs.mlflow_model_folder_mmd}}
-
   register_model:
     component: azureml:register_model:0.0.19
     # component: azureml://registries/azureml/components/register_model/versions/0.0.17
@@ -412,24 +351,5 @@ jobs:
     outputs:
       registration_details_folder: ${{parent.outputs.model_registration_details}}
 
-  register_model_mmd:
-    component: azureml:register_model:0.0.19
-    # component: azureml://registries/azureml/components/register_model/versions/0.0.17
-    compute: ${{parent.inputs.compute}}
-    resources:
-      instance_type: '${{parent.inputs.instance_type}}'
-    identity:
-      type: user_identity
-    inputs:
-      model_name: ${{parent.inputs.custom_model_name}}
-      model_version: ${{parent.inputs.model_version}}
-      model_description: ${{parent.inputs.model_description}}
-      registry_name: ${{parent.inputs.registry_name}}
-      model_metadata: ${{parent.inputs.model_metadata}}
-      model_type: mlflow_model
-      model_path: ${{parent.jobs.mlflow_model_local_validation_mmd.outputs.mlflow_model_folder}}
-    outputs:
-      registration_details_folder: ${{parent.outputs.model_registration_details_mmd}}
-
 tags:
     Preview: ""
diff --git a/assets/training/model_management/src/run_model_preprocess.py b/assets/training/model_management/src/run_model_preprocess.py
index 6dfebedfb2..c6f9b7011e 100644
--- a/assets/training/model_management/src/run_model_preprocess.py
+++ b/assets/training/model_management/src/run_model_preprocess.py
@@ -74,6 +74,7 @@ def _get_parser():
         help="Model download details",
     )
     parser.add_argument("--model-path", type=Path, required=True, help="Model input path")
+    parser.add_argument("--model-path-mmd", type=Path, required=True, help="MMD Model input path")
     parser.add_argument("--license-file-path", type=Path, required=False, help="License file path")
     parser.add_argument(
         "--mlflow-model-output-dir",
@@ -107,7 +108,7 @@ def run():
     inference_base_image = args.inference_base_image
 
     model_download_metadata_path = args.model_download_metadata
-    model_path = args.model_path
+    model_path = args.model_path_mmd if model_framework == 'MMLab' else args.model_path
     mlflow_model_output_dir = args.mlflow_model_output_dir
     license_file_path = args.license_file_path
     TRUST_CODE_KEY = "trust_remote_code=True"

From 71ad8a33b9a23ec4c217775d171eb4ff51002d35 Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Thu, 14 Nov 2024 11:18:28 +0530
Subject: [PATCH 03/25] Fixes

---
 .../components/convert_model_to_mlflow/spec.yaml            | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
index 84dae892f0..b66304e848 100644
--- a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
+++ b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -19,7 +19,7 @@ command: |
   pip_pkg_str="${pip_pkgs[*]}"
   if [[ -n "$pip_pkg_str" ]]; then echo "Installing $pip_pkg_str"; pip install $pip_pkg_str; echo "pip installation completed. For any installation error please check above logs"; fi;
   echo "Running model conversion ... "
-  python -u run_model_preprocess.py $[[--model-id ${{inputs.model_id}}]] $[[--task-name ${{inputs.task_name}}]] $[[--model-download-metadata ${{inputs.model_download_metadata}}]] $[[--license-file-path ${{inputs.license_file_path}}]] $[[--hf-config-args "${{inputs.hf_config_args}}"]] $[[--hf-tokenizer-args "${{inputs.hf_tokenizer_args}}"]] $[[--hf-model-args "${{inputs.hf_model_args}}"]] $[[--hf-pipeline-args "${{inputs.hf_pipeline_args}}"]] $[[--hf-config-class ${{inputs.hf_config_class}}]] $[[--hf-model-class ${{inputs.hf_model_class}}]] $[[--hf-tokenizer-class ${{inputs.hf_tokenizer_class}}]] $[[--hf-use-experimental-features ${{inputs.hf_use_experimental_features}}]] $[[--extra-pip-requirements "${{inputs.extra_pip_requirements}}"]] $[[--inference-base-image "${{inputs.inference_base_image}}"]] --vllm-enabled ${{inputs.vllm_enabled}} --model-framework ${{inputs.model_framework}} --model-path ${{inputs.model_path}} --mlflow-model-output-dir ${{outputs.mlflow_model_folder}} --model-flavor ${{inputs.model_flavor}}
+  python -u run_model_preprocess.py $[[--model-id ${{inputs.model_id}}]] $[[--task-name ${{inputs.task_name}}]] $[[--model-download-metadata ${{inputs.model_download_metadata}}]] $[[--license-file-path ${{inputs.license_file_path}}]] $[[--hf-config-args "${{inputs.hf_config_args}}"]] $[[--hf-tokenizer-args "${{inputs.hf_tokenizer_args}}"]] $[[--hf-model-args "${{inputs.hf_model_args}}"]] $[[--hf-pipeline-args "${{inputs.hf_pipeline_args}}"]] $[[--hf-config-class ${{inputs.hf_config_class}}]] $[[--hf-model-class ${{inputs.hf_model_class}}]] $[[--hf-tokenizer-class ${{inputs.hf_tokenizer_class}}]] $[[--hf-use-experimental-features ${{inputs.hf_use_experimental_features}}]] $[[--extra-pip-requirements "${{inputs.extra_pip_requirements}}"]] $[[--inference-base-image "${{inputs.inference_base_image}}"]] --vllm-enabled ${{inputs.vllm_enabled}} --model-framework ${{inputs.model_framework}} $[[--model-path ${{inputs.model_path}}]] $[[--model-path-mmd ${{inputs.model_path_mmd}}]] --mlflow-model-output-dir ${{outputs.mlflow_model_folder}} --model-flavor ${{inputs.model_flavor}}
   echo "Completed model conversion ... "
 
 inputs:
@@ -156,13 +156,13 @@ inputs:
     type: uri_folder
     description: Path to the model.
     mode: ro_mount
-    optional: false
+    optional: true
 
   model_path_mmd:
     type: uri_folder
     description: Path to the MMD model.
     mode: ro_mount
-    optional: false
+    optional: true
 
   license_file_path:
     type: uri_file

From 1256ae79ad4acc14977b8a69612a7fbd8cb433be Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Thu, 14 Nov 2024 11:44:22 +0530
Subject: [PATCH 04/25] Fixes

---
 .../components/convert_model_to_mlflow/spec.yaml                | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
index b66304e848..db8f0d7af2 100644
--- a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
+++ b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -19,7 +19,7 @@ command: |
   pip_pkg_str="${pip_pkgs[*]}"
   if [[ -n "$pip_pkg_str" ]]; then echo "Installing $pip_pkg_str"; pip install $pip_pkg_str; echo "pip installation completed. For any installation error please check above logs"; fi;
   echo "Running model conversion ... "
-  python -u run_model_preprocess.py $[[--model-id ${{inputs.model_id}}]] $[[--task-name ${{inputs.task_name}}]] $[[--model-download-metadata ${{inputs.model_download_metadata}}]] $[[--license-file-path ${{inputs.license_file_path}}]] $[[--hf-config-args "${{inputs.hf_config_args}}"]] $[[--hf-tokenizer-args "${{inputs.hf_tokenizer_args}}"]] $[[--hf-model-args "${{inputs.hf_model_args}}"]] $[[--hf-pipeline-args "${{inputs.hf_pipeline_args}}"]] $[[--hf-config-class ${{inputs.hf_config_class}}]] $[[--hf-model-class ${{inputs.hf_model_class}}]] $[[--hf-tokenizer-class ${{inputs.hf_tokenizer_class}}]] $[[--hf-use-experimental-features ${{inputs.hf_use_experimental_features}}]] $[[--extra-pip-requirements "${{inputs.extra_pip_requirements}}"]] $[[--inference-base-image "${{inputs.inference_base_image}}"]] --vllm-enabled ${{inputs.vllm_enabled}} --model-framework ${{inputs.model_framework}} $[[--model-path ${{inputs.model_path}}]] $[[--model-path-mmd ${{inputs.model_path_mmd}}]] --mlflow-model-output-dir ${{outputs.mlflow_model_folder}} --model-flavor ${{inputs.model_flavor}}
+  python -u run_model_preprocess.py $[[--model-id ${{inputs.model_id}}]] $[[--task-name ${{inputs.task_name}}]] $[[--model-download-metadata ${{inputs.model_download_metadata}}]] $[[--license-file-path ${{inputs.license_file_path}}]] $[[--hf-config-args "${{inputs.hf_config_args}}"]] $[[--hf-tokenizer-args "${{inputs.hf_tokenizer_args}}"]] $[[--hf-model-args "${{inputs.hf_model_args}}"]] $[[--hf-pipeline-args "${{inputs.hf_pipeline_args}}"]] $[[--hf-config-class ${{inputs.hf_config_class}}]] $[[--hf-model-class ${{inputs.hf_model_class}}]] $[[--hf-tokenizer-class ${{inputs.hf_tokenizer_class}}]] $[[--hf-use-experimental-features ${{inputs.hf_use_experimental_features}}]] $[[--extra-pip-requirements "${{inputs.extra_pip_requirements}}"]] $[[--inference-base-image "${{inputs.inference_base_image}}"]] --vllm-enabled ${{inputs.vllm_enabled}} --model-framework ${{inputs.model_framework}} $[[--model-path "${{inputs.model_path}}"]] $[[--model-path-mmd "${{inputs.model_path_mmd}}"]] --mlflow-model-output-dir ${{outputs.mlflow_model_folder}} --model-flavor ${{inputs.model_flavor}}
   echo "Completed model conversion ... "
 
 inputs:

From a13de9f3129a365f2e59c2fec9f22ff20674c4b8 Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Thu, 14 Nov 2024 13:11:25 +0530
Subject: [PATCH 05/25] Fixes

---
 assets/training/model_management/src/run_model_preprocess.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/assets/training/model_management/src/run_model_preprocess.py b/assets/training/model_management/src/run_model_preprocess.py
index c6f9b7011e..3c8efaafc1 100644
--- a/assets/training/model_management/src/run_model_preprocess.py
+++ b/assets/training/model_management/src/run_model_preprocess.py
@@ -73,8 +73,8 @@ def _get_parser():
         required=False,
         help="Model download details",
     )
-    parser.add_argument("--model-path", type=Path, required=True, help="Model input path")
-    parser.add_argument("--model-path-mmd", type=Path, required=True, help="MMD Model input path")
+    parser.add_argument("--model-path", type=Path, required=False, help="Model input path")
+    parser.add_argument("--model-path-mmd", type=Path, required=False, help="MMD Model input path")
     parser.add_argument("--license-file-path", type=Path, required=False, help="License file path")
     parser.add_argument(
         "--mlflow-model-output-dir",

From 1e45634c16c908f3e05cefb5a29e9d6a6b76de70 Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Wed, 20 Nov 2024 12:02:35 +0530
Subject: [PATCH 06/25] Fixes

---
 .../model_management/components/import_model/spec.yaml       | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index b416f2822e..c4244504b2 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -219,7 +219,6 @@ outputs:
 jobs:
   validation_trigger_import:
     component: azureml:validation_trigger_import:0.0.13
-    # component: azureml://registries/azureml/components/validation_trigger_import/versions/0.0.11
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -254,7 +253,6 @@ jobs:
 
   download_model:
     component: azureml:download_model:0.0.29
-    # component: azureml://registries/azureml/components/download_model/versions/0.0.28
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -289,7 +287,6 @@ jobs:
 
   convert_model_to_mlflow:
     component: azureml:convert_model_to_mlflow:0.0.34
-    # component: azureml://registries/azureml-preview-test1/components/convert_model_to_mlflow/versions/0.0.33-groopa-test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -319,7 +316,6 @@ jobs:
 
   mlflow_model_local_validation:
     component: azureml:mlflow_model_local_validation:0.0.16
-    # component: azureml://registries/azureml/components/mlflow_model_local_validation/versions/0.0.14
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -333,7 +329,6 @@ jobs:
 
   register_model:
     component: azureml:register_model:0.0.19
-    # component: azureml://registries/azureml/components/register_model/versions/0.0.17
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'

From 74de8796d5359733f282074990b45110d8bd2dbe Mon Sep 17 00:00:00 2001
From: Roopa G <groopa@microsoft.com>
Date: Mon, 2 Dec 2024 14:03:22 +0530
Subject: [PATCH 07/25] If else rough

---
 .../components/condition_block/asset.yaml     |  3 +
 .../components/condition_block/spec.yaml      | 34 +++++++++++
 .../src/run_condition_block.py                | 58 +++++++++++++++++++
 3 files changed, 95 insertions(+)
 create mode 100644 assets/training/model_management/components/condition_block/asset.yaml
 create mode 100644 assets/training/model_management/components/condition_block/spec.yaml
 create mode 100644 assets/training/model_management/src/run_condition_block.py

diff --git a/assets/training/model_management/components/condition_block/asset.yaml b/assets/training/model_management/components/condition_block/asset.yaml
new file mode 100644
index 0000000000..ea30e0bd5b
--- /dev/null
+++ b/assets/training/model_management/components/condition_block/asset.yaml
@@ -0,0 +1,3 @@
+type: component
+spec: spec.yaml
+categories: ["Models"]
diff --git a/assets/training/model_management/components/condition_block/spec.yaml b/assets/training/model_management/components/condition_block/spec.yaml
new file mode 100644
index 0000000000..02a4786d13
--- /dev/null
+++ b/assets/training/model_management/components/condition_block/spec.yaml
@@ -0,0 +1,34 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
+
+name: condition_block
+version: 0.0.1
+type: command
+
+is_deterministic: True
+
+display_name: Condition Block
+description: Component takes a condition and executes the command based on the condition.
+
+environment: azureml://registries/azureml/environments/model-management/versions/34
+
+code: ../../src/
+command: |
+  set -ex
+  IFS=',' read -ra input_args <<< "${{inputs.args}}"
+  echo "Running condition block ..."
+  python -u run_condition_block.py --condition ${{inputs.condition}} --input-args "${{inputs.args}}"
+  echo "Completed condition block ... "
+
+inputs:
+  condition:
+    type: string
+    description: Input condition.
+    optional: false
+
+  extra_pip_requirements:
+    type: string
+    description: |
+      Input arguments.
+      Arguments expressed as below. Do not use quotes for passing.
+      eg: a==1.0, b==1
+    optional: false
\ No newline at end of file
diff --git a/assets/training/model_management/src/run_condition_block.py b/assets/training/model_management/src/run_condition_block.py
new file mode 100644
index 0000000000..9366112381
--- /dev/null
+++ b/assets/training/model_management/src/run_condition_block.py
@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Run Model preprocessor module."""
+
+import argparse
+import os
+import json
+import shutil
+from azureml.model.mgmt.config import AppName, ModelFramework
+from azureml.model.mgmt.processors.transformers.config import HF_CONF
+from azureml.model.mgmt.processors.preprocess import run_preprocess, check_for_py_files
+from azureml.model.mgmt.processors.transformers.config import SupportedTasks as TransformersSupportedTasks
+from azureml.model.mgmt.processors.pyfunc.config import SupportedTasks as PyFuncSupportedTasks
+from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions, UnsupportedTaskType
+from azureml._common.exceptions import AzureMLException
+from azureml._common._error_definition.azureml_error import AzureMLError
+from azureml.model.mgmt.utils.logging_utils import custom_dimensions, get_logger
+from pathlib import Path
+from tempfile import TemporaryDirectory
+import json
+
+
+logger = get_logger(__name__)
+custom_dimensions.app_name = AppName.CONVERT_MODEL_TO_MLFLOW
+
+
+def _get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--condition", type=str, required=True, help="Condition")
+    parser.add_argument(
+        "--input-args",
+        type=str,
+        required=True,
+        help="Input args",
+    )
+    return parser
+
+
+@swallow_all_exceptions(logger)
+def run():
+    """Run preprocess."""
+    parser = _get_parser()
+    args, _ = parser.parse_known_args()
+
+    input_args = args.input_args
+    condition = args.condition
+    input_args = json.loads(input_args)
+    result = None
+    try:
+        result = eval(condition, input_args)
+    except Exception as e:
+        logger.error(f"Error evaluating condition: {e}")
+        result = False
+    
+
+if __name__ == "__main__":
+    run()

From 970d78ef29a781cef12e05f4d4defcf4f5340d88 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Tue, 3 Dec 2024 22:30:36 +0530
Subject: [PATCH 08/25] test conditional block changes

---
 .../components/condition_block/spec.yaml      | 23 ++++---
 .../components/import_model/spec.yaml         | 68 ++++++++++++++++---
 .../validation_trigger_import/spec.yaml       | 25 +++++--
 .../src/run_condition_block.py                | 48 +++++--------
 .../src/run_model_validate.py                 | 27 ++++++--
 5 files changed, 125 insertions(+), 66 deletions(-)

diff --git a/assets/training/model_management/components/condition_block/spec.yaml b/assets/training/model_management/components/condition_block/spec.yaml
index 02a4786d13..6e2752affe 100644
--- a/assets/training/model_management/components/condition_block/spec.yaml
+++ b/assets/training/model_management/components/condition_block/spec.yaml
@@ -1,7 +1,7 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: condition_block
-version: 0.0.1
+version: 0.0.1.8.test
 type: command
 
 is_deterministic: True
@@ -12,23 +12,24 @@ description: Component takes a condition and executes the command based on the c
 environment: azureml://registries/azureml/environments/model-management/versions/34
 
 code: ../../src/
-command: |
-  set -ex
-  IFS=',' read -ra input_args <<< "${{inputs.args}}"
-  echo "Running condition block ..."
-  python -u run_condition_block.py --condition ${{inputs.condition}} --input-args "${{inputs.args}}"
-  echo "Completed condition block ... "
+command: >
+  python run_condition_block.py
+  --condition ${{inputs.condition}}
+  --input-args ${{inputs.input_args}}
+  --result ${{outputs.result}}
 
 inputs:
   condition:
     type: string
     description: Input condition.
-    optional: false
-
-  extra_pip_requirements:
+  input_args:
     type: string
     description: |
       Input arguments.
       Arguments expressed as below. Do not use quotes for passing.
       eg: a==1.0, b==1
-    optional: false
\ No newline at end of file
+outputs:
+  result:
+    type: uri_file
+    description: Result of the condition evaluation.
+    mode: rw_mount
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index c4244504b2..c497dc7d74 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: 0.0.40
+version: mmd.test.41
 
 # Pipeline inputs
 inputs:
@@ -218,10 +218,13 @@ outputs:
 
 jobs:
   validation_trigger_import:
-    component: azureml:validation_trigger_import:0.0.13
+    type: command
+    component: azureml:validation_trigger_import:0.0.1.19.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
+    identity:
+      type: user_identity
     inputs:
       compute: ${{parent.inputs.compute}}
       task_name: ${{parent.inputs.task_name}}
@@ -231,6 +234,7 @@ jobs:
       model_id: ${{parent.inputs.model_id}}
       model_version: ${{parent.inputs.model_version}}
       model_flavor: ${{parent.inputs.model_flavor}}
+      model_framework: ${{parent.inputs.model_framework}}
       model_description: ${{parent.inputs.model_description}}
       model_metadata: ${{parent.inputs.model_metadata}}
       registry_name: ${{parent.inputs.registry_name}}
@@ -247,12 +251,54 @@ jobs:
       hf_tokenizer_class: ${{parent.inputs.hf_tokenizer_class}}
       hf_use_experimental_features: ${{parent.inputs.hf_use_experimental_features}}
       update_existing_model: ${{parent.inputs.update_existing_model}}
-    outputs:
-      validation_info:
-        type: uri_file
+    # outputs:
+    #   # validation_info:
+    #   #   type: uri_file
+    #   output:
+    #     type: boolean
+
+  # is_vision_model:
+  #   component: azureml:condition_block:0.0.1.12.test
+  #   compute: ${{parent.inputs.compute}}
+  #   resources:
+  #     instance_type: '${{parent.inputs.instance_type}}'
+  #   identity:
+  #     type: user_identity
+  #   inputs:
+  #     condition: "model_framework == 'MMLab'"
+  #     input_args: '{\"model_framework\": \"${{parent.inputs.model_framework}}\"}'
+  #   outputs:
+  #     result:
+  #       type: uri_file
+
+  # model_download_selector:
+  #   component: azureml://registries/azureml-preview-test1/components/output_selector/versions/0.0.1
+  #   inputs:
+  #     condition: ${{parent.jobs.is_vision_model.outputs.result}}
+  #     input_a: ${{parent.inputs.model_framework}}
+  #     input_b: 'MMLab'
+  #   outputs:
+  #     output: ${{parent.outputs.is_mmd_model_framework}}
+
+  # read_result:
+  #   type: command
+  #   command: |
+  #     python -c "import json; print(open('${{inputs.result_file}}').read().strip().lower())" 
+  #   environment: azureml:AzureML-sklearn-1.0-ubuntu20.04-py38-cpu:1
+  #   inputs:
+  #     result_file: ${{parent.jobs.is_vision_model.outputs.result}}
+  #   outputs:
+  #     result:
+        # type: uri_file
+
+  is_mlflow_model:
+    type: if_else
+    condition: ${{parent.jobs.validation_trigger_import.outputs.output}}
+    true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
+    false_block: ${{parent.jobs.download_model}}
 
   download_model:
-    component: azureml:download_model:0.0.29
+    component: azureml:download_model:0.0.30.0
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -261,7 +307,7 @@ jobs:
     inputs:
       model_source: ${{parent.inputs.model_source}}
       model_id: ${{parent.inputs.model_id}}
-      validation_info: ${{parent.jobs.validation_trigger_import.outputs.validation_info}}
+      # validation_info: ${{parent.jobs.validation_trigger_import.outputs.validation_info}}
       update_existing_model: ${{parent.inputs.update_existing_model}}
       token: ${{parent.inputs.token}}
     outputs:
@@ -271,7 +317,7 @@ jobs:
         type: uri_folder
 
   mmdetection_image_objectdetection_instancesegmentation_model_import:
-    component: azureml:mmdetection_image_objectdetection_instancesegmentation_model_import/versions/0.0.19
+    component: azureml:mmdetection_image_objectdetection_instancesegmentation_model_import:0.0.19
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -286,7 +332,7 @@ jobs:
         type: uri_file
 
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.34
+    component: azureml:convert_model_to_mlflow:0.0.33.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -315,7 +361,7 @@ jobs:
         type: mlflow_model
 
   mlflow_model_local_validation:
-    component: azureml:mlflow_model_local_validation:0.0.16
+    component: azureml:mlflow_model_local_validation:0.0.16.0
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -328,7 +374,7 @@ jobs:
       mlflow_model_folder: ${{parent.outputs.mlflow_model_folder}}
 
   register_model:
-    component: azureml:register_model:0.0.19
+    component: azureml:register_model:0.0.19.0
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
diff --git a/assets/training/model_management/components/validation_trigger_import/spec.yaml b/assets/training/model_management/components/validation_trigger_import/spec.yaml
index 94b87f080b..9a92adbdf7 100644
--- a/assets/training/model_management/components/validation_trigger_import/spec.yaml
+++ b/assets/training/model_management/components/validation_trigger_import/spec.yaml
@@ -4,7 +4,7 @@ type: command
 name: validation_trigger_import
 display_name: Validation Trigger Import model
 description: Component for enabling validation of import pipeline.
-version: 0.0.13
+version: 0.0.1.19.test
 
 # Pipeline inputs
 inputs:
@@ -47,6 +47,17 @@ inputs:
     optional: false
     description: Flavor of MLFlow to which model the model is converted to.
 
+  model_framework:
+    type: string
+    enum:
+      - Huggingface
+      - MMLab
+      - llava
+      - AutoML
+    default: Huggingface
+    optional: false
+    description: Framework from which model is imported from.
+
   ## Inputs for the MlFLow conversion
   license_file_path:
     type: uri_file
@@ -174,15 +185,17 @@ inputs:
     optional: true
 
 outputs:
-    validation_info:
-        type: uri_file
-        description: Validation status of the import pipeline inputs.
+  output:
+    type: boolean
+    is_control: true
 
 is_deterministic: True
 
-environment: azureml://registries/azureml/environments/python-sdk-v2/versions/23
+environment: azureml:python-sdk-v2-mmd-7:1
+
 code: ../../src
-command: python run_model_validate.py --validation-info ${{outputs.validation_info}}
+command: >  
+  mldesigner execute --source run_model_validate.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs output='${{outputs.output}}'
 
 tags:
     Preview: ""
diff --git a/assets/training/model_management/src/run_condition_block.py b/assets/training/model_management/src/run_condition_block.py
index 9366112381..c455035ecd 100644
--- a/assets/training/model_management/src/run_condition_block.py
+++ b/assets/training/model_management/src/run_condition_block.py
@@ -1,58 +1,44 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
-"""Run Model preprocessor module."""
-
 import argparse
-import os
-import json
-import shutil
-from azureml.model.mgmt.config import AppName, ModelFramework
-from azureml.model.mgmt.processors.transformers.config import HF_CONF
-from azureml.model.mgmt.processors.preprocess import run_preprocess, check_for_py_files
-from azureml.model.mgmt.processors.transformers.config import SupportedTasks as TransformersSupportedTasks
-from azureml.model.mgmt.processors.pyfunc.config import SupportedTasks as PyFuncSupportedTasks
-from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions, UnsupportedTaskType
-from azureml._common.exceptions import AzureMLException
-from azureml._common._error_definition.azureml_error import AzureMLError
-from azureml.model.mgmt.utils.logging_utils import custom_dimensions, get_logger
-from pathlib import Path
-from tempfile import TemporaryDirectory
 import json
+import logging
 
+# Assuming swallow_all_exceptions and get_logger are defined elsewhere
+from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
+from azureml.model.mgmt.utils.logging_utils import get_logger
 
 logger = get_logger(__name__)
-custom_dimensions.app_name = AppName.CONVERT_MODEL_TO_MLFLOW
-
 
 def _get_parser():
     parser = argparse.ArgumentParser()
     parser.add_argument("--condition", type=str, required=True, help="Condition")
-    parser.add_argument(
-        "--input-args",
-        type=str,
-        required=True,
-        help="Input args",
-    )
+    parser.add_argument("--input-args", type=str, required=True, help="Input args")
+    parser.add_argument("--result", type=str, required=True, help="Output result file path")
     return parser
 
-
 @swallow_all_exceptions(logger)
 def run():
     """Run preprocess."""
     parser = _get_parser()
     args, _ = parser.parse_known_args()
 
-    input_args = args.input_args
     condition = args.condition
+    input_args = args.input_args
+    result_path = args.result
     input_args = json.loads(input_args)
-    result = None
+    logger.info(f"Run preprocess with input args: {input_args}, and condition: {condition}")
+
     try:
+        # Use eval to evaluate the condition with a context dictionary
         result = eval(condition, input_args)
     except Exception as e:
         logger.error(f"Error evaluating condition: {e}")
         result = False
-    
+
+    logger.info(f"Condition result: {result}")
+
+    # Write the result to the output file
+    with open(result_path, 'w') as f:
+        f.write(str(result))
 
 if __name__ == "__main__":
     run()
diff --git a/assets/training/model_management/src/run_model_validate.py b/assets/training/model_management/src/run_model_validate.py
index af7b4f6556..833dde64bf 100644
--- a/assets/training/model_management/src/run_model_validate.py
+++ b/assets/training/model_management/src/run_model_validate.py
@@ -2,13 +2,26 @@
 # Licensed under the MIT License.
 
 """Validate Import Pipeline Parameters."""
-import argparse
+# import argparse
 
-parser = argparse.ArgumentParser()
-parser.add_argument("--validation-info", required=True, help="Model source ")
+from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
+from azureml.model.mgmt.utils.logging_utils import get_logger
+from mldesigner import Input, Output, command_component
 
-args = parser.parse_args()
+logger = get_logger(name=__name__)
 
-print("Validation info: ", args.validation_info)
-with open(args.validation_info, "w") as f:
-    f.write("Validation info")
+@command_component
+@swallow_all_exceptions(logger)
+def validate(
+        model_framework: Input(type="string", optional=True, default=None)
+        ) -> Output(type="boolean", is_control=True):
+    """Entry function of model validation script."""
+    
+    print(f"Model framework: {model_framework}")
+
+    model_framework = model_framework
+    result = model_framework == "MMLab"
+    
+    print(f"Model framework: {model_framework}, result: {result}")
+    logger.info(f"Model framework: {model_framework}, result: {result}")
+    return result
\ No newline at end of file

From 1e76215f8c7262a282d92009783f9837077fdb92 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Tue, 3 Dec 2024 22:34:11 +0530
Subject: [PATCH 09/25] test conditional block changes

---
 .../components/import_model/spec.yaml         | 34 -------------------
 1 file changed, 34 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index c497dc7d74..1b522c33eb 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -257,40 +257,6 @@ jobs:
     #   output:
     #     type: boolean
 
-  # is_vision_model:
-  #   component: azureml:condition_block:0.0.1.12.test
-  #   compute: ${{parent.inputs.compute}}
-  #   resources:
-  #     instance_type: '${{parent.inputs.instance_type}}'
-  #   identity:
-  #     type: user_identity
-  #   inputs:
-  #     condition: "model_framework == 'MMLab'"
-  #     input_args: '{\"model_framework\": \"${{parent.inputs.model_framework}}\"}'
-  #   outputs:
-  #     result:
-  #       type: uri_file
-
-  # model_download_selector:
-  #   component: azureml://registries/azureml-preview-test1/components/output_selector/versions/0.0.1
-  #   inputs:
-  #     condition: ${{parent.jobs.is_vision_model.outputs.result}}
-  #     input_a: ${{parent.inputs.model_framework}}
-  #     input_b: 'MMLab'
-  #   outputs:
-  #     output: ${{parent.outputs.is_mmd_model_framework}}
-
-  # read_result:
-  #   type: command
-  #   command: |
-  #     python -c "import json; print(open('${{inputs.result_file}}').read().strip().lower())" 
-  #   environment: azureml:AzureML-sklearn-1.0-ubuntu20.04-py38-cpu:1
-  #   inputs:
-  #     result_file: ${{parent.jobs.is_vision_model.outputs.result}}
-  #   outputs:
-  #     result:
-        # type: uri_file
-
   is_mlflow_model:
     type: if_else
     condition: ${{parent.jobs.validation_trigger_import.outputs.output}}

From 2241eaf48c94efd735929961484663a0f0f932c7 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 4 Dec 2024 15:11:50 +0530
Subject: [PATCH 10/25] test conditional block changes

---
 .../components/import_model/spec.yaml          | 18 ++++++++++--------
 .../validation_trigger_import/spec.yaml        | 13 +++++++++----
 .../model_management/src/run_model_validate.py | 10 +++++-----
 3 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 1b522c33eb..7d057532dc 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.41
+version: mmd.test.48
 
 # Pipeline inputs
 inputs:
@@ -212,6 +212,10 @@ outputs:
     description: Output path for the converted MLflow model
     type: mlflow_model
 
+  result:
+    description: Vision model registration details
+    type: boolean
+
   model_registration_details:
     description: Output folder with a file which captures transformations applied above and registration details in JSON file
     type: uri_folder
@@ -219,7 +223,7 @@ outputs:
 jobs:
   validation_trigger_import:
     type: command
-    component: azureml:validation_trigger_import:0.0.1.19.test
+    component: azureml:validation_trigger_import:0.0.1.25.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -251,15 +255,13 @@ jobs:
       hf_tokenizer_class: ${{parent.inputs.hf_tokenizer_class}}
       hf_use_experimental_features: ${{parent.inputs.hf_use_experimental_features}}
       update_existing_model: ${{parent.inputs.update_existing_model}}
-    # outputs:
-    #   # validation_info:
-    #   #   type: uri_file
-    #   output:
-    #     type: boolean
+    outputs:
+      # validation_result: ${{parent.outputs.validation_result}}
+      result: ${{parent.outputs.result}}
 
   is_mlflow_model:
     type: if_else
-    condition: ${{parent.jobs.validation_trigger_import.outputs.output}}
+    condition: ${{parent.jobs.validation_trigger_import.outputs.result}}
     true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
     false_block: ${{parent.jobs.download_model}}
 
diff --git a/assets/training/model_management/components/validation_trigger_import/spec.yaml b/assets/training/model_management/components/validation_trigger_import/spec.yaml
index 9a92adbdf7..055023a142 100644
--- a/assets/training/model_management/components/validation_trigger_import/spec.yaml
+++ b/assets/training/model_management/components/validation_trigger_import/spec.yaml
@@ -4,7 +4,7 @@ type: command
 name: validation_trigger_import
 display_name: Validation Trigger Import model
 description: Component for enabling validation of import pipeline.
-version: 0.0.1.19.test
+version: 0.0.1.25.test
 
 # Pipeline inputs
 inputs:
@@ -184,18 +184,23 @@ inputs:
     description: If set to true, will update the already existing model in the registry. If set to false, will create a new model in the provided registry.
     optional: true
 
+# Pipeline outputs
 outputs:
-  output:
+  # validation_result:
+  #   type: uri_file
+  #   description: Validation result of the model
+  result:
     type: boolean
     is_control: true
 
 is_deterministic: True
 
-environment: azureml:python-sdk-v2-mmd-7:1
+# environment: azureml:python-sdk-v2-mmd-8:1
+environment: azureml://registries/azureml/environments/model-management/versions/34
 
 code: ../../src
 command: >  
-  mldesigner execute --source run_model_validate.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs output='${{outputs.output}}'
+  mldesigner execute --source run_model_validate.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs result='${{outputs.result}}'
 
 tags:
     Preview: ""
diff --git a/assets/training/model_management/src/run_model_validate.py b/assets/training/model_management/src/run_model_validate.py
index 833dde64bf..f8edbc3729 100644
--- a/assets/training/model_management/src/run_model_validate.py
+++ b/assets/training/model_management/src/run_model_validate.py
@@ -13,15 +13,15 @@
 @command_component
 @swallow_all_exceptions(logger)
 def validate(
-        model_framework: Input(type="string", optional=True, default=None)
+        model_framework: Input(type="string", optional=False, default='HuggingFace')
         ) -> Output(type="boolean", is_control=True):
     """Entry function of model validation script."""
-    
-    print(f"Model framework: {model_framework}")
 
-    model_framework = model_framework
+    print(f"Model framework: {model_framework}")
     result = model_framework == "MMLab"
     
     print(f"Model framework: {model_framework}, result: {result}")
     logger.info(f"Model framework: {model_framework}, result: {result}")
-    return result
\ No newline at end of file
+    
+    # Return the result as a boolean control output
+    return result

From a3459aaebd51e83d9e84214a354e857942e195fd Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 4 Dec 2024 22:33:10 +0530
Subject: [PATCH 11/25] test conditional block changes

---
 .../components/condition_block/spec.yaml      | 73 +++++++++++++------
 .../components/import_model/spec.yaml         | 29 +++++---
 .../validation_trigger_import/spec.yaml       | 32 ++------
 .../src/run_condition_block.py                | 59 ++++++---------
 .../src/run_model_validate.py                 | 27 ++-----
 5 files changed, 105 insertions(+), 115 deletions(-)

diff --git a/assets/training/model_management/components/condition_block/spec.yaml b/assets/training/model_management/components/condition_block/spec.yaml
index 6e2752affe..27b9c1e847 100644
--- a/assets/training/model_management/components/condition_block/spec.yaml
+++ b/assets/training/model_management/components/condition_block/spec.yaml
@@ -1,35 +1,62 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: condition_block
-version: 0.0.1.8.test
-type: command
-
-is_deterministic: True
-
 display_name: Condition Block
 description: Component takes a condition and executes the command based on the condition.
+version: 0.0.1.16.test
+type: command
 
-environment: azureml://registries/azureml/environments/model-management/versions/34
-
-code: ../../src/
-command: >
-  python run_condition_block.py
-  --condition ${{inputs.condition}}
-  --input-args ${{inputs.input_args}}
-  --result ${{outputs.result}}
+is_deterministic: True
 
 inputs:
-  condition:
+  # pipeline specific compute
+  compute:
+    type: string
+    optional: true
+    default: serverless
+    description: Common compute for model download, MLflow conversion and registration. eg. provide 'FT-Cluster' if
+      your compute is named 'FT-Cluster'. Special characters like \ and ' are invalid in the parameter value.
+      If compute name is provided, instance_type field will be ignored and the respective cluster will be used
+
+  instance_type:
     type: string
-    description: Input condition.
-  input_args:
+    optional: true
+    default: Standard_DS12_v2
+    description: Instance type to be used for the component in case of serverless compute, eg. STANDARD_NC6s_v3. 
+      The parameter compute must be set to 'serverless' for instance_type to be used
+
+  model_framework:
     type: string
-    description: |
-      Input arguments.
-      Arguments expressed as below. Do not use quotes for passing.
-      eg: a==1.0, b==1
+    enum:
+      - Huggingface
+      - MMLab
+      - llava
+      - AutoML
+    default: Huggingface
+    optional: false
+    description: Framework from which model is imported from.
+
+  validation_info:
+    type: uri_file
+    description: Path to the validation info file
+    optional: true
+
+# Pipeline outputs
 outputs:
+  # validation_result:
+  #   type: uri_file
+  #   description: Validation result of the model
   result:
-    type: uri_file
-    description: Result of the condition evaluation.
-    mode: rw_mount
+    type: boolean
+    is_control: true
+
+environment: azureml:model-management-mmd-1:1
+
+code: ../../src
+command: >  
+  mldesigner execute --source run_condition_block.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs result='${{outputs.result}}'
+
+
+tags:
+    Preview: ""
+    Internal: ""
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 7d057532dc..18a3863dae 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.48
+version: mmd.test.50
 
 # Pipeline inputs
 inputs:
@@ -222,13 +222,10 @@ outputs:
 
 jobs:
   validation_trigger_import:
-    type: command
-    component: azureml:validation_trigger_import:0.0.1.25.test
+    component: azureml:validation_trigger_import:0.0.13.5.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
-    identity:
-      type: user_identity
     inputs:
       compute: ${{parent.inputs.compute}}
       task_name: ${{parent.inputs.task_name}}
@@ -238,7 +235,6 @@ jobs:
       model_id: ${{parent.inputs.model_id}}
       model_version: ${{parent.inputs.model_version}}
       model_flavor: ${{parent.inputs.model_flavor}}
-      model_framework: ${{parent.inputs.model_framework}}
       model_description: ${{parent.inputs.model_description}}
       model_metadata: ${{parent.inputs.model_metadata}}
       registry_name: ${{parent.inputs.registry_name}}
@@ -256,12 +252,27 @@ jobs:
       hf_use_experimental_features: ${{parent.inputs.hf_use_experimental_features}}
       update_existing_model: ${{parent.inputs.update_existing_model}}
     outputs:
-      # validation_result: ${{parent.outputs.validation_result}}
-      result: ${{parent.outputs.result}}
+      validation_info:
+        type: uri_file
+
+  is_vision_model:
+    type: command
+    component: azureml:condition_block:0.0.1.16.test
+    compute: ${{parent.inputs.compute}}
+    resources:
+      instance_type: '${{parent.inputs.instance_type}}'
+    identity:
+      type: user_identity
+    inputs:
+      validation_info: ${{parent.jobs.validation_trigger_import.outputs.validation_info}}
+      model_framework: ${{parent.inputs.model_framework}}
+    # outputs:
+    #   result:
+    #     type: boolean
 
   is_mlflow_model:
     type: if_else
-    condition: ${{parent.jobs.validation_trigger_import.outputs.result}}
+    condition: ${{parent.jobs.is_vision_model.outputs.result}}
     true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
     false_block: ${{parent.jobs.download_model}}
 
diff --git a/assets/training/model_management/components/validation_trigger_import/spec.yaml b/assets/training/model_management/components/validation_trigger_import/spec.yaml
index 055023a142..cc2b231081 100644
--- a/assets/training/model_management/components/validation_trigger_import/spec.yaml
+++ b/assets/training/model_management/components/validation_trigger_import/spec.yaml
@@ -4,7 +4,7 @@ type: command
 name: validation_trigger_import
 display_name: Validation Trigger Import model
 description: Component for enabling validation of import pipeline.
-version: 0.0.1.25.test
+version: 0.0.13.5.test
 
 # Pipeline inputs
 inputs:
@@ -47,17 +47,6 @@ inputs:
     optional: false
     description: Flavor of MLFlow to which model the model is converted to.
 
-  model_framework:
-    type: string
-    enum:
-      - Huggingface
-      - MMLab
-      - llava
-      - AutoML
-    default: Huggingface
-    optional: false
-    description: Framework from which model is imported from.
-
   ## Inputs for the MlFLow conversion
   license_file_path:
     type: uri_file
@@ -184,24 +173,17 @@ inputs:
     description: If set to true, will update the already existing model in the registry. If set to false, will create a new model in the provided registry.
     optional: true
 
-# Pipeline outputs
 outputs:
-  # validation_result:
-  #   type: uri_file
-  #   description: Validation result of the model
-  result:
-    type: boolean
-    is_control: true
+    validation_info:
+        type: uri_file
+        description: Validation status of the import pipeline inputs.
 
 is_deterministic: True
 
-# environment: azureml:python-sdk-v2-mmd-8:1
-environment: azureml://registries/azureml/environments/model-management/versions/34
-
+environment: azureml://registries/azureml/environments/python-sdk-v2/versions/23
 code: ../../src
-command: >  
-  mldesigner execute --source run_model_validate.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs result='${{outputs.result}}'
+command: python run_model_validate.py --validation-info ${{outputs.validation_info}}
 
 tags:
     Preview: ""
-    Internal: ""
+    Internal: ""
\ No newline at end of file
diff --git a/assets/training/model_management/src/run_condition_block.py b/assets/training/model_management/src/run_condition_block.py
index c455035ecd..f8edbc3729 100644
--- a/assets/training/model_management/src/run_condition_block.py
+++ b/assets/training/model_management/src/run_condition_block.py
@@ -1,44 +1,27 @@
-import argparse
-import json
-import logging
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Validate Import Pipeline Parameters."""
+# import argparse
 
-# Assuming swallow_all_exceptions and get_logger are defined elsewhere
 from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
 from azureml.model.mgmt.utils.logging_utils import get_logger
+from mldesigner import Input, Output, command_component
 
-logger = get_logger(__name__)
-
-def _get_parser():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--condition", type=str, required=True, help="Condition")
-    parser.add_argument("--input-args", type=str, required=True, help="Input args")
-    parser.add_argument("--result", type=str, required=True, help="Output result file path")
-    return parser
+logger = get_logger(name=__name__)
 
+@command_component
 @swallow_all_exceptions(logger)
-def run():
-    """Run preprocess."""
-    parser = _get_parser()
-    args, _ = parser.parse_known_args()
-
-    condition = args.condition
-    input_args = args.input_args
-    result_path = args.result
-    input_args = json.loads(input_args)
-    logger.info(f"Run preprocess with input args: {input_args}, and condition: {condition}")
-
-    try:
-        # Use eval to evaluate the condition with a context dictionary
-        result = eval(condition, input_args)
-    except Exception as e:
-        logger.error(f"Error evaluating condition: {e}")
-        result = False
-
-    logger.info(f"Condition result: {result}")
-
-    # Write the result to the output file
-    with open(result_path, 'w') as f:
-        f.write(str(result))
-
-if __name__ == "__main__":
-    run()
+def validate(
+        model_framework: Input(type="string", optional=False, default='HuggingFace')
+        ) -> Output(type="boolean", is_control=True):
+    """Entry function of model validation script."""
+
+    print(f"Model framework: {model_framework}")
+    result = model_framework == "MMLab"
+    
+    print(f"Model framework: {model_framework}, result: {result}")
+    logger.info(f"Model framework: {model_framework}, result: {result}")
+    
+    # Return the result as a boolean control output
+    return result
diff --git a/assets/training/model_management/src/run_model_validate.py b/assets/training/model_management/src/run_model_validate.py
index f8edbc3729..6312d8a22e 100644
--- a/assets/training/model_management/src/run_model_validate.py
+++ b/assets/training/model_management/src/run_model_validate.py
@@ -2,26 +2,13 @@
 # Licensed under the MIT License.
 
 """Validate Import Pipeline Parameters."""
-# import argparse
+import argparse
 
-from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
-from azureml.model.mgmt.utils.logging_utils import get_logger
-from mldesigner import Input, Output, command_component
+parser = argparse.ArgumentParser()
+parser.add_argument("--validation-info", required=True, help="Model source ")
 
-logger = get_logger(name=__name__)
+args = parser.parse_args()
 
-@command_component
-@swallow_all_exceptions(logger)
-def validate(
-        model_framework: Input(type="string", optional=False, default='HuggingFace')
-        ) -> Output(type="boolean", is_control=True):
-    """Entry function of model validation script."""
-
-    print(f"Model framework: {model_framework}")
-    result = model_framework == "MMLab"
-    
-    print(f"Model framework: {model_framework}, result: {result}")
-    logger.info(f"Model framework: {model_framework}, result: {result}")
-    
-    # Return the result as a boolean control output
-    return result
+print("Validation info: ", args.validation_info)
+with open(args.validation_info, "w") as f:
+    f.write("Validation info")
\ No newline at end of file

From a88e05a83bde93f259c29d7e467c70aad455114a Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Thu, 5 Dec 2024 14:48:15 +0530
Subject: [PATCH 12/25] test conditional block changes

---
 .../components/condition_block/spec.yaml            | 13 +++++--------
 .../components/import_model/spec.yaml               |  8 ++++----
 ...ion_block.py => run_model_framework_selector.py} |  4 ++--
 3 files changed, 11 insertions(+), 14 deletions(-)
 rename assets/training/model_management/src/{run_condition_block.py => run_model_framework_selector.py} (85%)

diff --git a/assets/training/model_management/components/condition_block/spec.yaml b/assets/training/model_management/components/condition_block/spec.yaml
index 27b9c1e847..13540d0190 100644
--- a/assets/training/model_management/components/condition_block/spec.yaml
+++ b/assets/training/model_management/components/condition_block/spec.yaml
@@ -1,9 +1,9 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
-name: condition_block
-display_name: Condition Block
+name: model_framework_selector
+display_name: Model Framework Selector
 description: Component takes a condition and executes the command based on the condition.
-version: 0.0.1.16.test
+version: 0.0.1.20.test
 type: command
 
 is_deterministic: True
@@ -43,10 +43,7 @@ inputs:
 
 # Pipeline outputs
 outputs:
-  # validation_result:
-  #   type: uri_file
-  #   description: Validation result of the model
-  result:
+  output:
     type: boolean
     is_control: true
 
@@ -54,7 +51,7 @@ environment: azureml:model-management-mmd-1:1
 
 code: ../../src
 command: >  
-  mldesigner execute --source run_condition_block.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs result='${{outputs.result}}'
+  mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs output="${{outputs.output}}"
 
 
 tags:
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 18a3863dae..22a54e68a4 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.50
+version: mmd.test.53
 
 # Pipeline inputs
 inputs:
@@ -255,9 +255,9 @@ jobs:
       validation_info:
         type: uri_file
 
-  is_vision_model:
+  framework_selector:
     type: command
-    component: azureml:condition_block:0.0.1.16.test
+    component: azureml:model_framework_selector:0.0.1.20.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -272,7 +272,7 @@ jobs:
 
   is_mlflow_model:
     type: if_else
-    condition: ${{parent.jobs.is_vision_model.outputs.result}}
+    condition: ${{parent.jobs.framework_selector.outputs.output}}
     true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
     false_block: ${{parent.jobs.download_model}}
 
diff --git a/assets/training/model_management/src/run_condition_block.py b/assets/training/model_management/src/run_model_framework_selector.py
similarity index 85%
rename from assets/training/model_management/src/run_condition_block.py
rename to assets/training/model_management/src/run_model_framework_selector.py
index f8edbc3729..a6cf368a8c 100644
--- a/assets/training/model_management/src/run_condition_block.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -14,8 +14,8 @@
 @swallow_all_exceptions(logger)
 def validate(
         model_framework: Input(type="string", optional=False, default='HuggingFace')
-        ) -> Output(type="boolean", is_control=True):
-    """Entry function of model validation script."""
+        ) -> Output(type="boolean", is_control=True): # noqa: F821
+    """Return true if model_framework is MMLab, False otherwise."""
 
     print(f"Model framework: {model_framework}")
     result = model_framework == "MMLab"

From 1cc8bb541ccf5d09d1601ddf3f5ee4745c5d5749 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Mon, 9 Dec 2024 11:41:00 +0530
Subject: [PATCH 13/25] test conditional block changes

---
 .../components/condition_block/spec.yaml      | 59 -------------------
 .../components/import_model/spec.yaml         | 15 ++---
 .../asset.yaml                                |  0
 .../model_framework_selector/spec.yaml        | 44 ++++++++++++++
 .../validation_trigger_import/spec.yaml       |  4 +-
 .../src/run_model_framework_selector.py       | 29 ++++-----
 6 files changed, 67 insertions(+), 84 deletions(-)
 delete mode 100644 assets/training/model_management/components/condition_block/spec.yaml
 rename assets/training/model_management/components/{condition_block => model_framework_selector}/asset.yaml (100%)
 create mode 100644 assets/training/model_management/components/model_framework_selector/spec.yaml

diff --git a/assets/training/model_management/components/condition_block/spec.yaml b/assets/training/model_management/components/condition_block/spec.yaml
deleted file mode 100644
index 13540d0190..0000000000
--- a/assets/training/model_management/components/condition_block/spec.yaml
+++ /dev/null
@@ -1,59 +0,0 @@
-$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
-
-name: model_framework_selector
-display_name: Model Framework Selector
-description: Component takes a condition and executes the command based on the condition.
-version: 0.0.1.20.test
-type: command
-
-is_deterministic: True
-
-inputs:
-  # pipeline specific compute
-  compute:
-    type: string
-    optional: true
-    default: serverless
-    description: Common compute for model download, MLflow conversion and registration. eg. provide 'FT-Cluster' if
-      your compute is named 'FT-Cluster'. Special characters like \ and ' are invalid in the parameter value.
-      If compute name is provided, instance_type field will be ignored and the respective cluster will be used
-
-  instance_type:
-    type: string
-    optional: true
-    default: Standard_DS12_v2
-    description: Instance type to be used for the component in case of serverless compute, eg. STANDARD_NC6s_v3. 
-      The parameter compute must be set to 'serverless' for instance_type to be used
-
-  model_framework:
-    type: string
-    enum:
-      - Huggingface
-      - MMLab
-      - llava
-      - AutoML
-    default: Huggingface
-    optional: false
-    description: Framework from which model is imported from.
-
-  validation_info:
-    type: uri_file
-    description: Path to the validation info file
-    optional: true
-
-# Pipeline outputs
-outputs:
-  output:
-    type: boolean
-    is_control: true
-
-environment: azureml:model-management-mmd-1:1
-
-code: ../../src
-command: >  
-  mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework="${{inputs.model_framework}}" --outputs output="${{outputs.output}}"
-
-
-tags:
-    Preview: ""
-    Internal: ""
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 22a54e68a4..f6e48f1f7a 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.53
+version: mmd.test.64
 
 # Pipeline inputs
 inputs:
@@ -255,24 +255,21 @@ jobs:
       validation_info:
         type: uri_file
 
-  framework_selector:
+  model_framework_selector:
     type: command
-    component: azureml:model_framework_selector:0.0.1.20.test
+    component: azureml:model_framework_selector:0.0.1.30.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
     identity:
       type: user_identity
     inputs:
-      validation_info: ${{parent.jobs.validation_trigger_import.outputs.validation_info}}
-      model_framework: ${{parent.inputs.model_framework}}
-    # outputs:
-    #   result:
-    #     type: boolean
+      validation_info: '${{parent.jobs.validation_trigger_import.outputs.validation_info}}'
+      model_framework: '${{parent.inputs.model_framework}}'
 
   is_mlflow_model:
     type: if_else
-    condition: ${{parent.jobs.framework_selector.outputs.output}}
+    condition: ${{parent.jobs.model_framework_selector.outputs.decision}}
     true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
     false_block: ${{parent.jobs.download_model}}
 
diff --git a/assets/training/model_management/components/condition_block/asset.yaml b/assets/training/model_management/components/model_framework_selector/asset.yaml
similarity index 100%
rename from assets/training/model_management/components/condition_block/asset.yaml
rename to assets/training/model_management/components/model_framework_selector/asset.yaml
diff --git a/assets/training/model_management/components/model_framework_selector/spec.yaml b/assets/training/model_management/components/model_framework_selector/spec.yaml
new file mode 100644
index 0000000000..9480249510
--- /dev/null
+++ b/assets/training/model_management/components/model_framework_selector/spec.yaml
@@ -0,0 +1,44 @@
+$schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
+
+name: model_framework_selector
+display_name: Model Framework Selector
+description: Component takes a condition and executes the command based on the condition.
+version: 0.0.1.30.test
+type: command
+
+is_deterministic: True
+
+inputs:
+
+  model_framework:
+    type: string
+    enum:
+      - Huggingface
+      - MMLab
+      - llava
+      - AutoML
+    default: Huggingface
+    optional: false
+    description: Framework from which model is imported from.
+
+  validation_info:
+    type: uri_file
+    description: Path to the validation info file
+    optional: true
+
+# Pipeline outputs
+outputs:
+  decision:
+    type: boolean
+    is_control: true
+
+# environment: azureml:model-management-mmd-1:1
+environment: azureml://registries/azureml/environments/model-evaluation/versions/37
+
+code: ../../src
+command: mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework='${{inputs.model_framework}}' --outputs output='${{outputs.decision}}'
+
+
+tags:
+    Preview: ""
+    Internal: ""
diff --git a/assets/training/model_management/components/validation_trigger_import/spec.yaml b/assets/training/model_management/components/validation_trigger_import/spec.yaml
index cc2b231081..55e65074e8 100644
--- a/assets/training/model_management/components/validation_trigger_import/spec.yaml
+++ b/assets/training/model_management/components/validation_trigger_import/spec.yaml
@@ -4,7 +4,7 @@ type: command
 name: validation_trigger_import
 display_name: Validation Trigger Import model
 description: Component for enabling validation of import pipeline.
-version: 0.0.13.5.test
+version: 0.0.13.6.test
 
 # Pipeline inputs
 inputs:
@@ -180,7 +180,7 @@ outputs:
 
 is_deterministic: True
 
-environment: azureml://registries/azureml/environments/python-sdk-v2/versions/23
+environment: azureml://registries/azureml/environments/python-sdk-v2/versions/24
 code: ../../src
 command: python run_model_validate.py --validation-info ${{outputs.validation_info}}
 
diff --git a/assets/training/model_management/src/run_model_framework_selector.py b/assets/training/model_management/src/run_model_framework_selector.py
index a6cf368a8c..191957609f 100644
--- a/assets/training/model_management/src/run_model_framework_selector.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -1,27 +1,28 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-"""Validate Import Pipeline Parameters."""
-# import argparse
+"""Select Model Framework Component."""
 
-from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
+from azure.ai.ml import Input
+from mldesigner import Output, command_component
 from azureml.model.mgmt.utils.logging_utils import get_logger
-from mldesigner import Input, Output, command_component
+from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
 
-logger = get_logger(name=__name__)
+logger = get_logger(__name__)
 
 @command_component
 @swallow_all_exceptions(logger)
 def validate(
-        model_framework: Input(type="string", optional=False, default='HuggingFace')
-        ) -> Output(type="boolean", is_control=True): # noqa: F821
-    """Return true if model_framework is MMLab, False otherwise."""
+        model_framework: Input(type="string", optional=False)  # noqa: F821
+) -> Output(type="boolean", is_control=True):  # noqa: F821
+    """Entry function of model validation script."""
+
+    if model_framework == "MMLab":
+        result = True
+    else:
+        result = False
 
-    print(f"Model framework: {model_framework}")
-    result = model_framework == "MMLab"
-    
-    print(f"Model framework: {model_framework}, result: {result}")
     logger.info(f"Model framework: {model_framework}, result: {result}")
-    
-    # Return the result as a boolean control output
+    print(f"Model framework: {model_framework}, result: {result}")
+
     return result

From 246d53251d615d5a415a562d1d9aa66a93b29d0e Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Mon, 9 Dec 2024 17:52:25 +0530
Subject: [PATCH 14/25] test conditional block changes

---
 .../model_management/components/import_model/spec.yaml | 10 +++++-----
 .../components/model_framework_selector/spec.yaml      |  5 ++++-
 .../src/run_model_framework_selector.py                |  5 +----
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index f6e48f1f7a..1f2bb6c5ec 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.64
+version: mmd.test.67
 
 # Pipeline inputs
 inputs:
@@ -212,9 +212,9 @@ outputs:
     description: Output path for the converted MLflow model
     type: mlflow_model
 
-  result:
-    description: Vision model registration details
-    type: boolean
+  # result:
+  #   description: Vision model registration details
+  #   type: boolean
 
   model_registration_details:
     description: Output folder with a file which captures transformations applied above and registration details in JSON file
@@ -257,7 +257,7 @@ jobs:
 
   model_framework_selector:
     type: command
-    component: azureml:model_framework_selector:0.0.1.30.test
+    component: azureml:model_framework_selector:0.0.1.32.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
diff --git a/assets/training/model_management/components/model_framework_selector/spec.yaml b/assets/training/model_management/components/model_framework_selector/spec.yaml
index 9480249510..603b3f1083 100644
--- a/assets/training/model_management/components/model_framework_selector/spec.yaml
+++ b/assets/training/model_management/components/model_framework_selector/spec.yaml
@@ -3,7 +3,7 @@ $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.jso
 name: model_framework_selector
 display_name: Model Framework Selector
 description: Component takes a condition and executes the command based on the condition.
-version: 0.0.1.30.test
+version: 0.0.1.32.test
 type: command
 
 is_deterministic: True
@@ -30,12 +30,15 @@ inputs:
 outputs:
   decision:
     type: boolean
+    mode: rw_mount
     is_control: true
 
 # environment: azureml:model-management-mmd-1:1
 environment: azureml://registries/azureml/environments/model-evaluation/versions/37
 
 code: ../../src
+# command: python run_model_framework_selector.py --model-framework ${{inputs.model_framework}} --decision ${{outputs.decision}}
+
 command: mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework='${{inputs.model_framework}}' --outputs output='${{outputs.decision}}'
 
 
diff --git a/assets/training/model_management/src/run_model_framework_selector.py b/assets/training/model_management/src/run_model_framework_selector.py
index 191957609f..b05d7516b3 100644
--- a/assets/training/model_management/src/run_model_framework_selector.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -12,9 +12,7 @@
 
 @command_component
 @swallow_all_exceptions(logger)
-def validate(
-        model_framework: Input(type="string", optional=False)  # noqa: F821
-) -> Output(type="boolean", is_control=True):  # noqa: F821
+def validate(model_framework: Input(type="string", optional=False))  -> bool:
     """Entry function of model validation script."""
 
     if model_framework == "MMLab":
@@ -23,6 +21,5 @@ def validate(
         result = False
 
     logger.info(f"Model framework: {model_framework}, result: {result}")
-    print(f"Model framework: {model_framework}, result: {result}")
 
     return result

From a5f46287298272db25fb99797526f7c81bf76558 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Mon, 9 Dec 2024 18:35:36 +0530
Subject: [PATCH 15/25] test conditional block changes

---
 .../model_management/components/import_model/spec.yaml    | 8 ++++++--
 .../components/model_framework_selector/spec.yaml         | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 1f2bb6c5ec..5d83e5297f 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.67
+version: mmd.test.69
 
 # Pipeline inputs
 inputs:
@@ -257,7 +257,7 @@ jobs:
 
   model_framework_selector:
     type: command
-    component: azureml:model_framework_selector:0.0.1.32.test
+    component: azureml:model_framework_selector:0.0.1.33.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -266,6 +266,10 @@ jobs:
     inputs:
       validation_info: '${{parent.jobs.validation_trigger_import.outputs.validation_info}}'
       model_framework: '${{parent.inputs.model_framework}}'
+    outputs:
+      decision:
+        type: uri_file
+        path: azureml://datastores/singularityblobstore/paths/azureml/${{name}}/${{output_name}}.jsonl
 
   is_mlflow_model:
     type: if_else
diff --git a/assets/training/model_management/components/model_framework_selector/spec.yaml b/assets/training/model_management/components/model_framework_selector/spec.yaml
index 603b3f1083..25d8bbf853 100644
--- a/assets/training/model_management/components/model_framework_selector/spec.yaml
+++ b/assets/training/model_management/components/model_framework_selector/spec.yaml
@@ -3,7 +3,7 @@ $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.jso
 name: model_framework_selector
 display_name: Model Framework Selector
 description: Component takes a condition and executes the command based on the condition.
-version: 0.0.1.32.test
+version: 0.0.1.33.test
 type: command
 
 is_deterministic: True

From 346c228dfda7bdf23a88969ea198e5360d6dad9f Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 11 Dec 2024 12:18:49 +0530
Subject: [PATCH 16/25] Removed unused code

---
 .../components/model_framework_selector/spec.yaml          | 7 +------
 .../model_management/src/run_model_framework_selector.py   | 7 ++++---
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/assets/training/model_management/components/model_framework_selector/spec.yaml b/assets/training/model_management/components/model_framework_selector/spec.yaml
index 25d8bbf853..63f7cb2a58 100644
--- a/assets/training/model_management/components/model_framework_selector/spec.yaml
+++ b/assets/training/model_management/components/model_framework_selector/spec.yaml
@@ -33,15 +33,10 @@ outputs:
     mode: rw_mount
     is_control: true
 
-# environment: azureml:model-management-mmd-1:1
 environment: azureml://registries/azureml/environments/model-evaluation/versions/37
-
-code: ../../src
-# command: python run_model_framework_selector.py --model-framework ${{inputs.model_framework}} --decision ${{outputs.decision}}
-
+code: ../../srcgit
 command: mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework='${{inputs.model_framework}}' --outputs output='${{outputs.decision}}'
 
-
 tags:
     Preview: ""
     Internal: ""
diff --git a/assets/training/model_management/src/run_model_framework_selector.py b/assets/training/model_management/src/run_model_framework_selector.py
index b05d7516b3..720e72926c 100644
--- a/assets/training/model_management/src/run_model_framework_selector.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -3,8 +3,7 @@
 
 """Select Model Framework Component."""
 
-from azure.ai.ml import Input
-from mldesigner import Output, command_component
+from mldesigner import Input, Output, command_component
 from azureml.model.mgmt.utils.logging_utils import get_logger
 from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
 
@@ -12,7 +11,9 @@
 
 @command_component
 @swallow_all_exceptions(logger)
-def validate(model_framework: Input(type="string", optional=False))  -> bool:
+def validate(
+        model_framework: Input(type="string", optional=False)  # noqa: F821
+) -> Output(type="boolean", is_control=True):  # noqa: F821
     """Entry function of model validation script."""
 
     if model_framework == "MMLab":

From cfde392ddb87daac640c2d0b36617bca172343c1 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 11 Dec 2024 14:06:23 +0530
Subject: [PATCH 17/25] Added output selector

---
 .../components/import_model/spec.yaml         | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 5d83e5297f..da1a438db4 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -212,9 +212,9 @@ outputs:
     description: Output path for the converted MLflow model
     type: mlflow_model
 
-  # result:
-  #   description: Vision model registration details
-  #   type: boolean
+  model_path:
+    description: Path to the model.
+    type: uri_folder
 
   model_registration_details:
     description: Output folder with a file which captures transformations applied above and registration details in JSON file
@@ -311,6 +311,17 @@ jobs:
       output_dir:
         type: uri_file
 
+  output_selector:
+    type: command
+    component: azureml://registries/azureml-preview-test1/components/output_selector/versions/0.0.1
+    compute: ${{parent.inputs.compute}}
+    inputs:
+      model_output_t: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
+      condition: ${{parent.jobs.framework_selector.outputs.output}}
+      model_output_f: ${{parent.jobs.download_model.outputs.model_output}}
+    outputs:
+      model_path: ${{parent.outputs.model_path}}
+
   convert_model_to_mlflow:
     component: azureml:convert_model_to_mlflow:0.0.33.test
     compute: ${{parent.inputs.compute}}
@@ -326,7 +337,7 @@ jobs:
       model_framework: ${{parent.inputs.model_framework}}
       model_download_metadata: ${{parent.jobs.download_model.outputs.model_download_metadata}}
       model_path_mmd:  ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
-      model_path: ${{parent.jobs.download_model.outputs.model_output}}
+      model_path: ${{parent.outputs.model_path}}
       hf_config_args: ${{parent.inputs.hf_config_args}}
       hf_tokenizer_args:  ${{parent.inputs.hf_tokenizer_args}}
       hf_model_args: ${{parent.inputs.hf_model_args}}

From 6c6fb01dad1608bc3a0480bedb9ea6d05bee82e4 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 11 Dec 2024 14:22:21 +0530
Subject: [PATCH 18/25] Added output selector

---
 .../components/import_model/spec.yaml            | 16 ++++++++--------
 .../model_management/src/run_model_preprocess.py |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index da1a438db4..e1182f7f5f 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.69
+version: mmd.test.75.test
 
 # Pipeline inputs
 inputs:
@@ -311,19 +311,19 @@ jobs:
       output_dir:
         type: uri_file
 
-  output_selector:
+  model_output_selector:
     type: command
     component: azureml://registries/azureml-preview-test1/components/output_selector/versions/0.0.1
     compute: ${{parent.inputs.compute}}
     inputs:
-      model_output_t: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
-      condition: ${{parent.jobs.framework_selector.outputs.output}}
-      model_output_f: ${{parent.jobs.download_model.outputs.model_output}}
+      input_a: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
+      condition: ${{parent.jobs.model_framework_selector.outputs.decision}}
+      input_b: ${{parent.jobs.download_model.outputs.model_output}}
     outputs:
-      model_path: ${{parent.outputs.model_path}}
+      output: ${{parent.outputs.model_path}}
 
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.33.test
+    component: azureml:convert_model_to_mlflow:0.0.34.test
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -337,7 +337,7 @@ jobs:
       model_framework: ${{parent.inputs.model_framework}}
       model_download_metadata: ${{parent.jobs.download_model.outputs.model_download_metadata}}
       model_path_mmd:  ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
-      model_path: ${{parent.outputs.model_path}}
+      model_path: ${{parent.jobs.model_output_selector.outputs.output}}
       hf_config_args: ${{parent.inputs.hf_config_args}}
       hf_tokenizer_args:  ${{parent.inputs.hf_tokenizer_args}}
       hf_model_args: ${{parent.inputs.hf_model_args}}
diff --git a/assets/training/model_management/src/run_model_preprocess.py b/assets/training/model_management/src/run_model_preprocess.py
index 3c8efaafc1..18557dd9c9 100644
--- a/assets/training/model_management/src/run_model_preprocess.py
+++ b/assets/training/model_management/src/run_model_preprocess.py
@@ -108,7 +108,7 @@ def run():
     inference_base_image = args.inference_base_image
 
     model_download_metadata_path = args.model_download_metadata
-    model_path = args.model_path_mmd if model_framework == 'MMLab' else args.model_path
+    model_path = args.model_path
     mlflow_model_output_dir = args.mlflow_model_output_dir
     license_file_path = args.license_file_path
     TRUST_CODE_KEY = "trust_remote_code=True"

From e05036bbc6fd95d79772fe8ddfbe1cbb00cac6a9 Mon Sep 17 00:00:00 2001
From: Deepansh Goyal <188957625+deepanshMS@users.noreply.github.com>
Date: Wed, 18 Dec 2024 14:53:40 +0530
Subject: [PATCH 19/25] remove un-needed components

---
 .../convert_model_to_mlflow/spec.yaml         |  2 +-
 .../components/import_model/spec.yaml         | 64 +++++++------------
 .../model_framework_selector/spec.yaml        | 14 ++--
 .../validation_trigger_import/spec.yaml       |  7 +-
 .../src/run_model_preprocess.py               |  4 ++
 .../src/run_model_validate.py                 |  2 +-
 6 files changed, 41 insertions(+), 52 deletions(-)

diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
index db8f0d7af2..f2cff6b9ed 100644
--- a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
+++ b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -1,7 +1,7 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: convert_model_to_mlflow
-version: 0.0.34
+version: 0.0.34.0.test-dgoyal
 type: command
 
 is_deterministic: True
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index e1182f7f5f..4b51514622 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.75.test
+version: mmd.test.77.test-dgoyal
 
 # Pipeline inputs
 inputs:
@@ -212,10 +212,6 @@ outputs:
     description: Output path for the converted MLflow model
     type: mlflow_model
 
-  model_path:
-    description: Path to the model.
-    type: uri_folder
-
   model_registration_details:
     description: Output folder with a file which captures transformations applied above and registration details in JSON file
     type: uri_folder
@@ -257,7 +253,7 @@ jobs:
 
   model_framework_selector:
     type: command
-    component: azureml:model_framework_selector:0.0.1.33.test
+    component: azureml:model_framework_selector:0.0.1.38.test-dgoyal
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -267,16 +263,30 @@ jobs:
       validation_info: '${{parent.jobs.validation_trigger_import.outputs.validation_info}}'
       model_framework: '${{parent.inputs.model_framework}}'
     outputs:
-      decision:
+      is_mmd_framework:
         type: uri_file
-        path: azureml://datastores/singularityblobstore/paths/azureml/${{name}}/${{output_name}}.jsonl
 
-  is_mlflow_model:
+  is_mmd_model:
     type: if_else
-    condition: ${{parent.jobs.model_framework_selector.outputs.decision}}
-    true_block: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import}}
+    condition: ${{parent.jobs.model_framework_selector.outputs.is_mmd_framework}}
+    true_block: ${{parent.jobs.download_mmd_model}}
     false_block: ${{parent.jobs.download_model}}
 
+  download_mmd_model:
+    component: azureml:mmdetection_image_objectdetection_instancesegmentation_model_import:0.0.19
+    compute: ${{parent.inputs.compute}}
+    resources:
+      instance_type: '${{parent.inputs.instance_type}}'
+    identity:
+      type: user_identity
+    inputs:
+      model_family: 'MmDetectionImage'
+      model_name: ${{parent.inputs.model_id}}
+      download_from_source: False
+    outputs:
+      output_dir:
+        type: uri_file
+
   download_model:
     component: azureml:download_model:0.0.30.0
     compute: ${{parent.inputs.compute}}
@@ -296,34 +306,8 @@ jobs:
       model_output:
         type: uri_folder
 
-  mmdetection_image_objectdetection_instancesegmentation_model_import:
-    component: azureml:mmdetection_image_objectdetection_instancesegmentation_model_import:0.0.19
-    compute: ${{parent.inputs.compute}}
-    resources:
-      instance_type: '${{parent.inputs.instance_type}}'
-    identity:
-      type: user_identity
-    inputs:
-      model_family: 'MmDetectionImage'
-      model_name: ${{parent.inputs.model_id}}
-      download_from_source: False
-    outputs:
-      output_dir:
-        type: uri_file
-
-  model_output_selector:
-    type: command
-    component: azureml://registries/azureml-preview-test1/components/output_selector/versions/0.0.1
-    compute: ${{parent.inputs.compute}}
-    inputs:
-      input_a: ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
-      condition: ${{parent.jobs.model_framework_selector.outputs.decision}}
-      input_b: ${{parent.jobs.download_model.outputs.model_output}}
-    outputs:
-      output: ${{parent.outputs.model_path}}
-
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.34.test
+    component: azureml:convert_model_to_mlflow:0.0.34.0.test-dgoyal
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -336,8 +320,8 @@ jobs:
       license_file_path: ${{parent.inputs.license_file_path}}
       model_framework: ${{parent.inputs.model_framework}}
       model_download_metadata: ${{parent.jobs.download_model.outputs.model_download_metadata}}
-      model_path_mmd:  ${{parent.jobs.mmdetection_image_objectdetection_instancesegmentation_model_import.outputs.output_dir}}
-      model_path: ${{parent.jobs.model_output_selector.outputs.output}}
+      model_path: ${{parent.jobs.download_model.outputs.model_output}}
+      model_path_mmd:  ${{parent.jobs.download_mmd_model.outputs.output_dir}}
       hf_config_args: ${{parent.inputs.hf_config_args}}
       hf_tokenizer_args:  ${{parent.inputs.hf_tokenizer_args}}
       hf_model_args: ${{parent.inputs.hf_model_args}}
diff --git a/assets/training/model_management/components/model_framework_selector/spec.yaml b/assets/training/model_management/components/model_framework_selector/spec.yaml
index 63f7cb2a58..ee8f53c57d 100644
--- a/assets/training/model_management/components/model_framework_selector/spec.yaml
+++ b/assets/training/model_management/components/model_framework_selector/spec.yaml
@@ -1,10 +1,10 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
+type: command
 name: model_framework_selector
 display_name: Model Framework Selector
-description: Component takes a condition and executes the command based on the condition.
-version: 0.0.1.33.test
-type: command
+description: Checks the framework for model
+version: 0.0.1.38.test-dgoyal
 
 is_deterministic: True
 
@@ -24,18 +24,18 @@ inputs:
   validation_info:
     type: uri_file
     description: Path to the validation info file
-    optional: true
+    optional: false
 
 # Pipeline outputs
 outputs:
-  decision:
+  is_mmd_framework:
     type: boolean
     mode: rw_mount
     is_control: true
 
 environment: azureml://registries/azureml/environments/model-evaluation/versions/37
-code: ../../srcgit
-command: mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework='${{inputs.model_framework}}' --outputs output='${{outputs.decision}}'
+code: ../../src
+command: mldesigner execute --source run_model_framework_selector.py --name validate --inputs model_framework='${{inputs.model_framework}}' --outputs output='${{outputs.is_mmd_framework}}'
 
 tags:
     Preview: ""
diff --git a/assets/training/model_management/components/validation_trigger_import/spec.yaml b/assets/training/model_management/components/validation_trigger_import/spec.yaml
index 55e65074e8..afa560c6e0 100644
--- a/assets/training/model_management/components/validation_trigger_import/spec.yaml
+++ b/assets/training/model_management/components/validation_trigger_import/spec.yaml
@@ -4,7 +4,7 @@ type: command
 name: validation_trigger_import
 display_name: Validation Trigger Import model
 description: Component for enabling validation of import pipeline.
-version: 0.0.13.6.test
+version: 0.0.14
 
 # Pipeline inputs
 inputs:
@@ -67,6 +67,7 @@ inputs:
     - image-classification
     - text-to-image
     - chat-completion
+    - image-feature-extraction
     optional: true
     type: string
 
@@ -180,10 +181,10 @@ outputs:
 
 is_deterministic: True
 
-environment: azureml://registries/azureml/environments/python-sdk-v2/versions/24
+environment: azureml://registries/azureml/environments/python-sdk-v2/versions/23
 code: ../../src
 command: python run_model_validate.py --validation-info ${{outputs.validation_info}}
 
 tags:
     Preview: ""
-    Internal: ""
\ No newline at end of file
+    Internal: ""
diff --git a/assets/training/model_management/src/run_model_preprocess.py b/assets/training/model_management/src/run_model_preprocess.py
index 18557dd9c9..9d3b38f94f 100644
--- a/assets/training/model_management/src/run_model_preprocess.py
+++ b/assets/training/model_management/src/run_model_preprocess.py
@@ -108,7 +108,11 @@ def run():
     inference_base_image = args.inference_base_image
 
     model_download_metadata_path = args.model_download_metadata
+    
     model_path = args.model_path
+    if model_framework == "MMLab":
+        model_path = args.model_path_mmd
+
     mlflow_model_output_dir = args.mlflow_model_output_dir
     license_file_path = args.license_file_path
     TRUST_CODE_KEY = "trust_remote_code=True"
diff --git a/assets/training/model_management/src/run_model_validate.py b/assets/training/model_management/src/run_model_validate.py
index 6312d8a22e..af7b4f6556 100644
--- a/assets/training/model_management/src/run_model_validate.py
+++ b/assets/training/model_management/src/run_model_validate.py
@@ -11,4 +11,4 @@
 
 print("Validation info: ", args.validation_info)
 with open(args.validation_info, "w") as f:
-    f.write("Validation info")
\ No newline at end of file
+    f.write("Validation info")

From 2d012601dd5d7b039ee0d727c2b6ef40d9b2f248 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 18 Dec 2024 15:09:56 +0530
Subject: [PATCH 20/25] updated component versions

---
 .../components/convert_model_to_mlflow/spec.yaml   |  2 +-
 .../components/import_model/spec.yaml              | 14 +++++++-------
 .../components/model_framework_selector/spec.yaml  |  2 +-
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
index f2cff6b9ed..db8f0d7af2 100644
--- a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
+++ b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -1,7 +1,7 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: convert_model_to_mlflow
-version: 0.0.34.0.test-dgoyal
+version: 0.0.34
 type: command
 
 is_deterministic: True
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 4b51514622..2dcbd103f1 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: mmd.test.77.test-dgoyal
+version: 0.0.40
 
 # Pipeline inputs
 inputs:
@@ -218,7 +218,7 @@ outputs:
 
 jobs:
   validation_trigger_import:
-    component: azureml:validation_trigger_import:0.0.13.5.test
+    component: azureml:validation_trigger_import:0.0.14
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -253,7 +253,7 @@ jobs:
 
   model_framework_selector:
     type: command
-    component: azureml:model_framework_selector:0.0.1.38.test-dgoyal
+    component: azureml:model_framework_selector:0.0.1
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -288,7 +288,7 @@ jobs:
         type: uri_file
 
   download_model:
-    component: azureml:download_model:0.0.30.0
+    component: azureml:download_model:0.0.29
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -307,7 +307,7 @@ jobs:
         type: uri_folder
 
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.34.0.test-dgoyal
+    component: azureml:convert_model_to_mlflow:0.0.34
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -336,7 +336,7 @@ jobs:
         type: mlflow_model
 
   mlflow_model_local_validation:
-    component: azureml:mlflow_model_local_validation:0.0.16.0
+    component: azureml:mlflow_model_local_validation:0.0.16
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
@@ -349,7 +349,7 @@ jobs:
       mlflow_model_folder: ${{parent.outputs.mlflow_model_folder}}
 
   register_model:
-    component: azureml:register_model:0.0.19.0
+    component: azureml:register_model:0.0.19
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
diff --git a/assets/training/model_management/components/model_framework_selector/spec.yaml b/assets/training/model_management/components/model_framework_selector/spec.yaml
index ee8f53c57d..74b2ab6ae6 100644
--- a/assets/training/model_management/components/model_framework_selector/spec.yaml
+++ b/assets/training/model_management/components/model_framework_selector/spec.yaml
@@ -4,7 +4,7 @@ type: command
 name: model_framework_selector
 display_name: Model Framework Selector
 description: Checks the framework for model
-version: 0.0.1.38.test-dgoyal
+version: 0.0.1
 
 is_deterministic: True
 

From 031a3103bf0539383b379649d6b5c20c062ff633 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Wed, 18 Dec 2024 15:19:45 +0530
Subject: [PATCH 21/25] removed unused code

---
 .../training/model_management/components/import_model/spec.yaml  | 1 -
 1 file changed, 1 deletion(-)

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 2dcbd103f1..9eb2969307 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -297,7 +297,6 @@ jobs:
     inputs:
       model_source: ${{parent.inputs.model_source}}
       model_id: ${{parent.inputs.model_id}}
-      # validation_info: ${{parent.jobs.validation_trigger_import.outputs.validation_info}}
       update_existing_model: ${{parent.inputs.update_existing_model}}
       token: ${{parent.inputs.token}}
     outputs:

From fa226b4b72faf1836bd29670b582e4a77d9762e8 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Tue, 24 Dec 2024 09:11:00 +0530
Subject: [PATCH 22/25] updated the download component for mmlab

---
 .../components/import_model/spec.yaml         |  2 +-
 .../mgmt/processors/pyfunc/convertors.py      |  8 ++++----
 .../processors/pyfunc/vision/conda_is.yaml    | 19 +++++++++++++++++++
 .../processors/pyfunc/vision/conda_od.yaml    | 17 +++++++++++++++++
 4 files changed, 41 insertions(+), 5 deletions(-)
 create mode 100644 assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_is.yaml
 create mode 100644 assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_od.yaml

diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index 9eb2969307..a875ee0451 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -282,7 +282,7 @@ jobs:
     inputs:
       model_family: 'MmDetectionImage'
       model_name: ${{parent.inputs.model_id}}
-      download_from_source: False
+      download_from_source: true
     outputs:
       output_dir:
         type: uri_file
diff --git a/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/convertors.py b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/convertors.py
index e20fc66285..557e954708 100644
--- a/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/convertors.py
+++ b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/convertors.py
@@ -191,11 +191,11 @@ def save_as_mlflow(self):
         mlflow_model_wrapper = ImagesDetectionMLflowModelWrapper(task_type=self._task)
         artifacts_dict = self._prepare_artifacts_dict()
         if self._task == MMLabDetectionTasks.MM_OBJECT_DETECTION.value:
-            pip_requirements = os.path.join(self.MODEL_DIR, "mmdet-od-requirements.txt")
+            conda_env_file = os.path.join(self.MODEL_DIR, "conda_od.yaml")
         elif self._task == MMLabDetectionTasks.MM_INSTANCE_SEGMENTATION.value:
-            pip_requirements = os.path.join(self.MODEL_DIR, "mmdet-is-requirements.txt")
+            conda_env_file = os.path.join(self.MODEL_DIR, "conda_is.yaml")
         else:
-            pip_requirements = None
+            conda_env_file = None
         code_path = [
             os.path.join(self.MODEL_DIR, "detection_predict.py"),
             os.path.join(self.MODEL_DIR, "config.py"),
@@ -204,8 +204,8 @@ def save_as_mlflow(self):
         super()._save(
             mlflow_model_wrapper=mlflow_model_wrapper,
             artifacts_dict=artifacts_dict,
-            pip_requirements=pip_requirements,
             code_path=code_path,
+            conda_env=conda_env_file,
         )
 
     def _prepare_artifacts_dict(self) -> Dict:
diff --git a/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_is.yaml b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_is.yaml
new file mode 100644
index 0000000000..96a48f8c67
--- /dev/null
+++ b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_is.yaml
@@ -0,0 +1,19 @@
+channels:
+- conda-forge
+dependencies:
+- python=3.9.19
+- pip<=24.0
+- pip:
+  - mlflow==2.12.1
+  - cloudpickle==2.2.1
+  - datasets==2.15.0
+  - openmim==0.3.9
+  - torch==2.0.1
+  - torchvision==0.15.2
+  - transformers==4.38.2
+  - accelerate==0.27.2
+  - albumentations==1.3.0
+  - scikit-image==0.19.3
+  - simplification==0.7.10
+  - fairscale==0.4.13
+name: mlflow-env
diff --git a/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_od.yaml b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_od.yaml
new file mode 100644
index 0000000000..738e6fd5f8
--- /dev/null
+++ b/assets/training/model_management/src/azureml/model/mgmt/processors/pyfunc/vision/conda_od.yaml
@@ -0,0 +1,17 @@
+channels:
+- conda-forge
+dependencies:
+- python=3.9.19
+- pip<=24.0
+- pip:
+  - mlflow==2.12.1
+  - cloudpickle==2.2.1
+  - datasets==2.15.0
+  - openmim==0.3.9
+  - torch==2.0.1
+  - torchvision==0.15.2
+  - transformers==4.38.2
+  - accelerate==0.27.2
+  - albumentations==1.3.0
+  - fairscale==0.4.13
+name: mlflow-env

From 15f5bad050761988645f3e940ad57caf094f8c00 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Tue, 24 Dec 2024 10:53:25 +0530
Subject: [PATCH 23/25] updated the import model component

---
 .../components/convert_model_to_mlflow/spec.yaml              | 2 +-
 .../model_management/components/import_model/spec.yaml        | 4 ++--
 .../model_management/src/run_model_framework_selector.py      | 4 ++--
 assets/training/model_management/src/run_model_preprocess.py  | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
index fc8e168813..4a3d383277 100644
--- a/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
+++ b/assets/training/model_management/components/convert_model_to_mlflow/spec.yaml
@@ -1,7 +1,7 @@
 $schema: https://azuremlschemas.azureedge.net/latest/commandComponent.schema.json
 
 name: convert_model_to_mlflow
-version: 0.0.35
+version: 0.0.36
 type: command
 
 is_deterministic: True
diff --git a/assets/training/model_management/components/import_model/spec.yaml b/assets/training/model_management/components/import_model/spec.yaml
index d6b39a6ef2..eda6b8aaab 100644
--- a/assets/training/model_management/components/import_model/spec.yaml
+++ b/assets/training/model_management/components/import_model/spec.yaml
@@ -4,7 +4,7 @@ type: pipeline
 name: import_model
 display_name: Import model
 description: Import a model into a workspace or a registry
-version: 0.0.41
+version: 0.0.42
 
 # Pipeline inputs
 inputs:
@@ -307,7 +307,7 @@ jobs:
         type: uri_folder
 
   convert_model_to_mlflow:
-    component: azureml:convert_model_to_mlflow:0.0.35
+    component: azureml:convert_model_to_mlflow:0.0.36
     compute: ${{parent.inputs.compute}}
     resources:
       instance_type: '${{parent.inputs.instance_type}}'
diff --git a/assets/training/model_management/src/run_model_framework_selector.py b/assets/training/model_management/src/run_model_framework_selector.py
index 720e72926c..e692ca4297 100644
--- a/assets/training/model_management/src/run_model_framework_selector.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -2,7 +2,7 @@
 # Licensed under the MIT License.
 
 """Select Model Framework Component."""
-
+from azureml.model.mgmt.config import ModelFramework
 from mldesigner import Input, Output, command_component
 from azureml.model.mgmt.utils.logging_utils import get_logger
 from azureml.model.mgmt.utils.exceptions import swallow_all_exceptions
@@ -16,7 +16,7 @@ def validate(
 ) -> Output(type="boolean", is_control=True):  # noqa: F821
     """Entry function of model validation script."""
 
-    if model_framework == "MMLab":
+    if model_framework == ModelFramework.MMLAB.value:
         result = True
     else:
         result = False
diff --git a/assets/training/model_management/src/run_model_preprocess.py b/assets/training/model_management/src/run_model_preprocess.py
index 9d3b38f94f..f812e9fe2a 100644
--- a/assets/training/model_management/src/run_model_preprocess.py
+++ b/assets/training/model_management/src/run_model_preprocess.py
@@ -110,7 +110,7 @@ def run():
     model_download_metadata_path = args.model_download_metadata
     
     model_path = args.model_path
-    if model_framework == "MMLab":
+    if model_framework == ModelFramework.MMLAB.value:
         model_path = args.model_path_mmd
 
     mlflow_model_output_dir = args.mlflow_model_output_dir

From 067952cb8c2538ca59a2496cbdc53adfffbdcbfc Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Tue, 24 Dec 2024 12:08:25 +0530
Subject: [PATCH 24/25] code refactoring

---
 .../model_management/src/run_model_framework_selector.py        | 1 +
 assets/training/model_management/src/run_model_preprocess.py    | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/assets/training/model_management/src/run_model_framework_selector.py b/assets/training/model_management/src/run_model_framework_selector.py
index e692ca4297..0e6a95ecc8 100644
--- a/assets/training/model_management/src/run_model_framework_selector.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -9,6 +9,7 @@
 
 logger = get_logger(__name__)
 
+
 @command_component
 @swallow_all_exceptions(logger)
 def validate(
diff --git a/assets/training/model_management/src/run_model_preprocess.py b/assets/training/model_management/src/run_model_preprocess.py
index f812e9fe2a..5a6ab6e32d 100644
--- a/assets/training/model_management/src/run_model_preprocess.py
+++ b/assets/training/model_management/src/run_model_preprocess.py
@@ -108,7 +108,7 @@ def run():
     inference_base_image = args.inference_base_image
 
     model_download_metadata_path = args.model_download_metadata
-    
+
     model_path = args.model_path
     if model_framework == ModelFramework.MMLAB.value:
         model_path = args.model_path_mmd

From 79817c8eb42e196530e068e3a9042bd317014538 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhisheku@microsoft.com>
Date: Tue, 24 Dec 2024 12:44:34 +0530
Subject: [PATCH 25/25] Fixed syntax issue

---
 .../model_management/src/run_model_framework_selector.py         | 1 -
 1 file changed, 1 deletion(-)

diff --git a/assets/training/model_management/src/run_model_framework_selector.py b/assets/training/model_management/src/run_model_framework_selector.py
index 0e6a95ecc8..ee6ff82240 100644
--- a/assets/training/model_management/src/run_model_framework_selector.py
+++ b/assets/training/model_management/src/run_model_framework_selector.py
@@ -16,7 +16,6 @@ def validate(
         model_framework: Input(type="string", optional=False)  # noqa: F821
 ) -> Output(type="boolean", is_control=True):  # noqa: F821
     """Entry function of model validation script."""
-
     if model_framework == ModelFramework.MMLAB.value:
         result = True
     else: