From f32463936d3b8205df7b66dbd9c3f9a2fd69668a Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 18 Sep 2023 23:23:25 -0400
Subject: [PATCH] Unhardcode sampler and scheduler list in test.

---
 tests/inference/test_inference.py | 20 ++++++--------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/tests/inference/test_inference.py b/tests/inference/test_inference.py
index a96f9455..141cc5c7 100644
--- a/tests/inference/test_inference.py
+++ b/tests/inference/test_inference.py
@@ -16,10 +16,8 @@ import uuid
 import urllib.request
 import urllib.parse
 
-# Currently causes an error when running pytest with built-in pytest args
-# TODO: modify cli_args.py to not parse args on import
-# We will hard-code sampler and scheduler lists for now
-# from comfy.samplers import KSampler
+
+from comfy.samplers import KSampler
 
 """
 These tests generate and save images through a range of parameters
@@ -140,16 +138,10 @@ comfy_graph_ids = [DEFAULT_COMFY_GRAPH_ID]
 prompt_list = [
     'a painting of a cat',
 ]
-#TODO use sampler and scheduler list from comfy.samplers.KSampler
-# sampler_list = KSampler.SAMPLERS
-# scheduler_list = KSampler.SCHEDULERS
-# Hard coded sampler and scheduler lists for now
-SCHEDULERS = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
-SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
-                "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
-                "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"]
-sampler_list = SAMPLERS
-scheduler_list = SCHEDULERS
+
+sampler_list = KSampler.SAMPLERS
+scheduler_list = KSampler.SCHEDULERS
+
 @pytest.mark.inference
 @pytest.mark.parametrize("sampler", sampler_list)
 @pytest.mark.parametrize("scheduler", scheduler_list)