fix diffusers test

Signed-off-by: Raphael Glon <oOraph@users.noreply.github.com>
huggingface · Nov 30, 2023 · b8263dd · b8263dd
1 parent afe2260
commit b8263dd
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/tests/test_dockers.py b/tests/test_dockers.py
@@ -313,12 +313,16 @@ def test_diffusers(self):
             "diffusers",
             "text-to-image",
             "hf-internal-testing/tiny-stable-diffusion-pipe",
+            {"num_inference_steps": 1},
+            custom_environment={"UVICORN_TIMEOUT": "1200"},
         )
         self.framework_docker_test(
             "diffusers",
             "image-to-image",
             "hf-internal-testing/tiny-controlnet",
+            {"num_inference_steps": 1},
             timeout=600,
+            custom_environment={"UVICORN_TIMEOUT": "1200"},
         )
         self.framework_invalid_test("diffusers")
 
@@ -479,6 +483,7 @@ def framework_docker_test(
             Any
         ] = None,  # if given, check inference with this specific input
         timeout=60,
+        custom_environment: Optional[dict] = None,
     ):
         tag = self.create_docker(framework)
         run_docker_command = [
@@ -490,6 +495,12 @@ def framework_docker_test(
             f"TASK={task}",
             "-e",
             f"MODEL_ID={model_id}",
+        ]
+
+        if custom_environment:
+            for k, v in custom_environment.items():
+                run_docker_command += ["-e", f"{k}={v}"]
+        run_docker_command += [
             "-v",
             "/tmp:/data",
             "-t",