Use DiffusersModelVariant

Author: xiaoyu-work

olive/common/hf/io_config/input_generators.py

@@ -19,6 +19,8 @@
 if TYPE_CHECKING:
     from transformers import PretrainedConfig
 
+    from olive.constants import DiffusersModelVariant
+
 
 class DummyInputGenerator(ABC):
     """Generate dummy inputs for ONNX export."""
@@ -70,7 +72,7 @@ class DiffusersDummyInputGenerator(DummyInputGenerator):
     Reads input specifications from diffusers.yaml.
     """
 
-    def __init__(self, component_name: str, config: PretrainedConfig, pipeline: str | None = None):
+    def __init__(self, component_name: str, config: PretrainedConfig, pipeline: DiffusersModelVariant | None = None):
         self.component_name = component_name
         self.config = config
 
@@ -145,7 +147,7 @@ def generate(self, input_name: str):
 def generate_diffusers_dummy_inputs(
     component_name: str,
     config: PretrainedConfig,
-    pipeline: str | None = None,
+    pipeline: DiffusersModelVariant | None = None,
 ) -> dict[str, Any]:
     """Create all dummy inputs for a diffusers component.
 

olive/common/hf/io_config/io_resolver.py

@@ -7,7 +7,10 @@
 import functools
 import logging
 from importlib.resources import files
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from olive.constants import DiffusersModelVariant
 
 import yaml
 
@@ -54,7 +57,9 @@ def get_task_template(task: str) -> dict[str, Any] | None:
     return tasks.get(task)
 
 
-def get_diffusers_component_config(component_name: str, pipeline: str | None = None) -> dict[str, Any] | None:
+def get_diffusers_component_config(
+    component_name: str, pipeline: DiffusersModelVariant | None = None
+) -> dict[str, Any] | None:
     """Get diffusers component configuration.
 
     Args:

olive/common/hf/io_config/task_config.py

@@ -17,6 +17,8 @@
 if TYPE_CHECKING:
     from transformers import PretrainedConfig, PreTrainedModel
 
+    from olive.constants import DiffusersModelVariant
+
 logger = logging.getLogger(__name__)
 
 
@@ -365,7 +367,7 @@ def generate_dummy_inputs(
 def get_diffusers_io_config(
     component_name: str,
     config: PretrainedConfig,
-    pipeline: str | None = None,
+    pipeline: DiffusersModelVariant | None = None,
     **kwargs,
 ) -> dict[str, Any]:
     """Get IO configuration for a diffusers component.

olive/passes/onnx/conversion.py

@@ -673,14 +673,14 @@ def _convert_diffusers_model(
             dummy_inputs = generate_diffusers_dummy_inputs(
                 component_name=component_name,
                 config=component_config,
-                pipeline=str(pipeline_type),
+                pipeline=pipeline_type,
             )
 
             # Get IO config using new task-driven API
             io_config = get_diffusers_io_config(
                 component_name=component_name,
                 config=component_config,
-                pipeline=str(pipeline_type),
+                pipeline=pipeline_type,
             )
 
             # Create output directory for this component