Make benchmark lighter

hydra_node/models.py

@@ -268,12 +268,14 @@ class StableDiffusionModel(nn.Module):
         settings.sampler = "k_euler_ancestral"
         settings.uc = ""
         settings.seed = 42
-        res_list = [(512,512), (512,704), (512,768), (576,768), (768,768), (1024,704), (1024,768), (1024,1024), (1280,1280)]
+        res_list = [(512,512), (512,704), (512,768), (576,768), (768,768), (1024,704), (1024,768), (1024,1024)]
         results = []
         print("benchmarking")
         for res in res_list:
-            for n in [1, 2, 4]:
+            for n in [1, 4]:
                 for w, h in [res]:
+                    if w * h > 1024 * 704 and n > 1:
+                        continue
                     settings.width = w
                     settings.height = h
                     settings.n_samples = n
@@ -281,7 +283,7 @@ class StableDiffusionModel(nn.Module):
                     self.sample(settings)
                     dur = time.perf_counter() - now
                     results.append((w, h, n, dur/2.0))
-                    print(f"{w}x{h} with {n} images ({settings.steps} steps {settings.sampler}) took: {dur:.4f}s ({1./(dur/settings.steps):.4f}it/s, {(dur/settings.steps)*28:.4f}s for 28 steps)")
+                    print(f"{w}x{h} with {n} images ({settings.steps} steps {settings.sampler}) took: {dur:.4f}s ({1./(dur/settings.steps):.4f}it/s, {(dur/settings.steps)*28:.4f}s for 28 steps)", flush=True)
         print("results")
         for w,h,n,dur in results:
             print(f"{w}x{h} with {n} images ({settings.steps} steps {settings.sampler}) took: {dur:.4f}s ({1./(dur/settings.steps):.4f}it/s, {(dur/settings.steps)*28:.4f}s for 28 steps)")