fix

deepspeedai · Feb 15, 2024 · d3fa2e2 · d3fa2e2
1 parent 4c6a302
commit d3fa2e2
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/mii/legacy/models/providers/diffusers.py b/mii/legacy/models/providers/diffusers.py
@@ -17,7 +17,7 @@ def diffusers_provider(model_config: ModelConfig):
     kwargs = model_config.pipeline_kwargs
     if model_config.dtype == torch.half:
         kwargs["torch_dtype"] = torch.float16
-        kwargs["revision"] = "fp16"
+        #kwargs["revision"] = "fp16"
 
     pipeline = attempt_load(DiffusionPipeline.from_pretrained,
                             model_config.model,

diff --git a/tests/legacy/test_local_deployment.py b/tests/legacy/test_local_deployment.py
@@ -63,17 +63,6 @@
                 "query": "DeepSpeed is the greatest"
             },
         ),
-        (
-            "zero-shot-image-classification",
-            "openai/clip-vit-base-patch32",
-            {
-                "image":
-                "https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png",
-                "candidate_labels": ["animals",
-                                     "humans",
-                                     "landscape"]
-            },
-        ),
     ],
 )
 def test_single_GPU(deployment, query):
@@ -132,7 +121,19 @@ def test_session(deployment, query):
             "text-to-image",
             "openskyml/midjourney-mini",
             {
-                "query": ["a dog on a rocket"]
+                "prompt": "a dog on a rocket",
+                "negative_prompt": "planet earth",
+            },
+        ),
+        (
+            "zero-shot-image-classification",
+            "openai/clip-vit-base-patch32",
+            {
+                "image":
+                "https://huggingface.co/datasets/Narsil/image_dummy/raw/main/parrots.png",
+                "candidate_labels": ["animals",
+                                     "humans",
+                                     "landscape"]
             },
         ),
     ],