feat: Add testing_models (#378)

Haidra-Org · Feb 26, 2024 · 72abd5f · 72abd5f
1 parent e8075d7
commit 72abd5f
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 41 deletions.
diff --git a/horde/classes/stable/worker.py b/horde/classes/stable/worker.py
@@ -174,7 +174,7 @@ def parse_models(self, unchecked_models):
                 if self.user.get_unique_alias() != user_alias:
                     raise e.BadRequest(f"This model can only be hosted by {user_alias}")
                 models.add(model)
-            elif model in model_reference.stable_diffusion_names or self.user.customizer:
+            elif model in model_reference.stable_diffusion_names or self.user.customizer or model in model_reference.testing_models:
                 models.add(model)
             else:
                 logger.debug(f"Rejecting unknown model '{model}' from {self.name} ({self.id})")

diff --git a/horde/model_reference.py b/horde/model_reference.py
@@ -12,6 +12,13 @@ class ModelReference(PrimaryTimedFunction):
     text_model_names = set()
     nsfw_models = set()
     controlnet_models = set()
+    testing_models = {
+        "Juggernaut XL",
+        "Animagine XL",
+        "DreamShaper XL",
+        "Stable Cascade 1.0",
+        "Anime Illust Diffusion XL",
+    }
 
     def call_function(self):
         """Retrieves to nataili and text model reference and stores in it a var"""

diff --git a/horde/sandbox.py b/horde/sandbox.py
@@ -16,44 +16,44 @@
 
 
 def test():
-    # with HORDE.app_context():
-    #     logger.debug(stats.get_model_avg("Deliberate"))
-    #     logger.debug(stats.get_model_avg("stable_diffusion"))
-    # logger.debug(database.count_totals())
-
-    # pp = pprint.PrettyPrinter(depth=3)
-    # pp.pprint(patrons.get_monthly_kudos(42742))
-    # pp.pprint(patrons.get_ids())
-
-    # Test sus discord webhook
-    # send_pause_notification("Hello World")
-    # with HORDE.app_context():
-    #     worker = database.find_worker_by_name("Db0_Test_Worker", worker_class=ImageWorker)
-    #     logger.debug(worker.get_bridge_kudos_multiplier())
-    # worker.report_suspicion(amount = 1, reason = Suspicions.UNREASONABLY_FAST, formats = [9999])
-    # threads.store_patreon_members()
-
-    # Test problem userdiscord webhook
-    # with HORDE.app_context():
-    #     worker = database.find_worker_by_name("Db0_Test_Worker", worker_class=ImageWorker)
-    #     worker.report_suspicion(amount = 1, reason = Suspicions.UNREASONABLY_FAST, formats = [9999])
-
-    # Cache testing
-    # with HORDE.app_context():
-    # logger.info(database.retrieve_totals(True))
-    # logger.info(database.retrieve_totals())
-    # database.get_available_models()
-
-    # IP timeout testing
-    # CounterMeasures.set_block_timeout("2001:db8::/64",1)
-    # logger.debug(CounterMeasures.get_block_timeouts())
-    # logger.debug(CounterMeasures.retrieve_timeout('2001:db8:0000:0000:0000:0000:0000:0001'))
-    # logger.debug(CounterMeasures.extract_ipv6_subnet('2001:db8:0000:0000:0000:0000:0000:0001'))
-
-    # Worker rewards test
-    # with HORDE.app_context():
-    #     w = database.find_worker_by_id("43f5f639-134f-4687-a130-b4bf13821d8c")
-    #     logger.debug(w.max_context_length)
-    #     logger.debug(w.calculate_uptime_reward())
-
+    with HORDE.app_context():
+        #     logger.debug(stats.get_model_avg("Deliberate"))
+        #     logger.debug(stats.get_model_avg("stable_diffusion"))
+        # logger.debug(database.count_totals())
+
+        # pp = pprint.PrettyPrinter(depth=3)
+        # pp.pprint(patrons.get_monthly_kudos(42742))
+        # pp.pprint(patrons.get_ids())
+
+        # Test sus discord webhook
+        # send_pause_notification("Hello World")
+        # with HORDE.app_context():
+        #     worker = database.find_worker_by_name("Db0_Test_Worker", worker_class=ImageWorker)
+        #     logger.debug(worker.get_bridge_kudos_multiplier())
+        # worker.report_suspicion(amount = 1, reason = Suspicions.UNREASONABLY_FAST, formats = [9999])
+        # threads.store_patreon_members()
+
+        # Test problem userdiscord webhook
+        # with HORDE.app_context():
+        #     worker = database.find_worker_by_name("Db0_Test_Worker", worker_class=ImageWorker)
+        #     worker.report_suspicion(amount = 1, reason = Suspicions.UNREASONABLY_FAST, formats = [9999])
+
+        # Cache testing
+        # with HORDE.app_context():
+        # logger.info(database.retrieve_totals(True))
+        # logger.info(database.retrieve_totals())
+        # database.get_available_models()
+
+        # IP timeout testing
+        # CounterMeasures.set_block_timeout("2001:db8::/64",1)
+        # logger.debug(CounterMeasures.get_block_timeouts())
+        # logger.debug(CounterMeasures.retrieve_timeout('2001:db8:0000:0000:0000:0000:0000:0001'))
+        # logger.debug(CounterMeasures.extract_ipv6_subnet('2001:db8:0000:0000:0000:0000:0000:0001'))
+
+        # Worker rewards test
+        # with HORDE.app_context():
+        #     w = database.find_worker_by_id("43f5f639-134f-4687-a130-b4bf13821d8c")
+        #     logger.debug(w.max_context_length)
+        #     logger.debug(w.calculate_uptime_reward())
+        pass
     sys.exit()