mechanics

app.py

@@ -236,6 +236,18 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
             return image, embs, img_embs, ys, calibrate_prompts
         else:
             print('######### Roaming #########')
+            
+            pos_indices = [i for i in range(len(embs)) if ys[i] == 1]
+            neg_indices = [i for i in range(len(embs)) if ys[i] == 0]
+
+            if len(neg_indices) > 40:
+                neg_indices = neg_indices[1:]
+                # popping first negative rating due to > 25
+            
+            indices = pos_indices + neg_indices
+            embs = [embs[i] for i in indices]
+            img_embs = [img_embs[i] for i in indices]
+            ys = [ys[i] for i in indices]
 
             im_s = get_coeff(embs, ys)
             rng_prompt = random.choice(prompt_list)
@@ -251,11 +263,7 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
             img_emb = w * learn_emb.to(dtype=torch.float16)
             image, img_emb = predict(prompt, im_emb=img_emb)
             img_embs.append(img_emb)
-
-            if len(embs) > 100:
-                embs.pop(0)
-                img_embs.pop(0)
-                ys.pop(0)
+            
             return image, embs, img_embs, ys, calibrate_prompts
 
 
@@ -292,7 +300,6 @@ def choose(img, choice, embs, img_embs, ys, calibrate_prompts):
     else:
         choice = 0
     
-    print(img, 'img')
     if img is None:
         print('NSFW -- choice is disliked')
         choice = 0