support testing without diffusers

huggingface · Sep 11, 2024 · 475efdf · 475efdf
1 parent 56d06d4
commit 475efdf
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 38 deletions.
diff --git a/optimum/onnxruntime/__init__.py b/optimum/onnxruntime/__init__.py
@@ -83,6 +83,10 @@
         "ORTPipelineForInpainting",
         "ORTPipelineForText2Image",
         "ORTDiffusionPipeline",
+        "ORTModelTextEncoder",
+        "ORTModelUnet",
+        "ORTModelVaeDecoder",
+        "ORTModelVaeEncoder",
     ]
 else:
     _import_structure["modeling_diffusion"] = [
@@ -96,6 +100,10 @@
         "ORTPipelineForInpainting",
         "ORTPipelineForText2Image",
         "ORTDiffusionPipeline",
+        "ORTModelTextEncoder",
+        "ORTModelUnet",
+        "ORTModelVaeDecoder",
+        "ORTModelVaeEncoder",
     ]
 
 
@@ -147,6 +155,10 @@
         from ..utils.dummy_diffusers_objects import (
             ORTDiffusionPipeline,
             ORTLatentConsistencyModelPipeline,
+            ORTModelTextEncoder,
+            ORTModelUnet,
+            ORTModelVaeDecoder,
+            ORTModelVaeEncoder,
             ORTPipelineForImage2Image,
             ORTPipelineForInpainting,
             ORTPipelineForText2Image,
@@ -160,6 +172,10 @@
         from .modeling_diffusion import (
             ORTDiffusionPipeline,
             ORTLatentConsistencyModelPipeline,
+            ORTModelTextEncoder,
+            ORTModelUnet,
+            ORTModelVaeDecoder,
+            ORTModelVaeEncoder,
             ORTPipelineForImage2Image,
             ORTPipelineForInpainting,
             ORTPipelineForText2Image,

diff --git a/optimum/utils/dummy_diffusers_objects.py b/optimum/utils/dummy_diffusers_objects.py
@@ -123,3 +123,47 @@ def __init__(self, *args, **kwargs):
     @classmethod
     def from_pretrained(cls, *args, **kwargs):
         requires_backends(cls, ["diffusers"])
+
+
+class ORTModelTextEncoder(metaclass=DummyObject):
+    _backends = ["diffusers"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["diffusers"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["diffusers"])
+
+
+class ORTModelVaeDecoder(metaclass=DummyObject):
+    _backends = ["diffusers"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["diffusers"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["diffusers"])
+
+
+class ORTModelVaeEncoder(metaclass=DummyObject):
+    _backends = ["diffusers"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["diffusers"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["diffusers"])
+
+
+class ORTModelUnet(metaclass=DummyObject):
+    _backends = ["diffusers"]
+
+    def __init__(self, *args, **kwargs):
+        requires_backends(self, ["diffusers"])
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        requires_backends(cls, ["diffusers"])