Add benchmark method

hydra_node/models.py

@@ -258,6 +258,33 @@ class StableDiffusionModel(nn.Module):
             self.prior = VectorAdjustPrior.load_model(config.prior_path).to(self.device)
         self.copied_ema = False
 
+    def benchmark(self):
+        settings = self.get_default_config
+        settings.prompt = "bru"
+        settings.steps = 50
+        settings.sampler = "k_euler_ancestral"
+        settings.uc = ""
+        res_list = [(512,512), (512,704), (512,768), (576,768), (768,768), (1024,704), (1024,768), (1024,1024), (1280,1280)]
+        results = []
+        print("benchmarking")
+        for res in res_list:
+            for w, h in [res, (res[1], res[0])]:
+                for n in [1, 2, 4]:
+                    settings.width = w
+                    settings.height = h
+                    settings.n_samples = n
+                    self.sample(settings)
+                    now = time.perf_counter()
+                    self.sample(settings)
+                    self.sample(settings)
+                    dur = time.perf_counter() - now
+                    results.append((w, h, n, f"{dur/2.0:.4f}s"))
+                break
+            break
+        print("results")
+        for w,h,n,dur in results:
+            print(f"{w}x{h} with {n} images ({settings.steps} steps {settings.sampler}) took: {dur}")
+
     @property
     def get_default_config(self):
         dict_config = {