Update app.py

app.py

@@ -45,7 +45,8 @@ import spaces
 start_time = time.time()
 
 ####################### Setup Model
-from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, LCMScheduler, AutoencoderTiny, UNet2DConditionModel, AutoencoderKL
+from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, LCMScheduler, AutoencoderTiny, UNet2DConditionModel, AutoencoderKL, utils
+utils.logging.disable_progress_bar
 from transformers import CLIPTextModel
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
@@ -55,7 +56,6 @@ import uuid
 import av
 
 def write_video(file_name, images, fps=17):
-    print('Saving')
     container = av.open(file_name, mode="w")
 
     stream = container.add_stream("h264", rate=fps)
@@ -76,7 +76,6 @@ def write_video(file_name, images, fps=17):
         container.mux(packet)
     # Close the file
     container.close()
-    print('Saved')
 
 def imio_write_video(file_name, images, fps=15):
     writer = imageio.get_writer(file_name, fps=fps)
@@ -128,14 +127,11 @@ pipe.to(device=DEVICE)
 
 @spaces.GPU()
 def generate_gpu(in_im_embs):
-    print('start gen')
     in_im_embs = in_im_embs.to('cuda').unsqueeze(0).unsqueeze(0)
     output = pipe(prompt='', guidance_scale=0, added_cond_kwargs={}, ip_adapter_image_embeds=[in_im_embs], num_inference_steps=STEPS)
-    print('image is made')
     im_emb, _ = pipe.encode_image(
                 output.frames[0][len(output.frames[0])//2], 'cuda', 1, output_hidden_state
             )
-    print('im_emb is made')
     im_emb = im_emb.detach().to('cpu').to(torch.float32)
     return output, im_emb
 
@@ -168,7 +164,6 @@ def get_user_emb(embs, ys):
         embs.append(.01*torch.randn(1280))
         ys.append(0)
         ys.append(1)
-        print('Fixing only one feedback class available.\n')
     
     indices = list(range(len(embs)))
     # sample only as many negatives as there are positives
@@ -177,14 +172,12 @@ def get_user_emb(embs, ys):
     #lower = min(len(pos_indices), len(neg_indices))
     #neg_indices = random.sample(neg_indices, lower)
     #pos_indices = random.sample(pos_indices, lower)
-    print(len(neg_indices), len(pos_indices))
     
     
     # we may have just encountered a rare multi-threading diffusers issue (https://github.com/huggingface/diffusers/issues/5749);
     # this ends up adding a rating but losing an embedding, it seems.
     # let's take off a rating if so to continue without indexing errors.
     if len(ys) > len(embs):
-        print('ys are longer than embs; popping latest rating')
         ys.pop(-1)
     
     feature_embs = np.array(torch.stack([embs[i].squeeze().to('cpu') for i in indices]).to('cpu'))
@@ -192,12 +185,10 @@ def get_user_emb(embs, ys):
     #feature_embs = scaler.transform(feature_embs)
     chosen_y = np.array([ys[i] for i in indices])
     
-    print('Gathering coefficients')
     #lin_class = Ridge(fit_intercept=False).fit(feature_embs, chosen_y)
     lin_class = SVC(max_iter=20, kernel='linear', C=.1, class_weight='balanced').fit(feature_embs, chosen_y)
     coef_ = torch.tensor(lin_class.coef_, dtype=torch.double).detach().to('cpu')
     coef_ = coef_ / coef_.abs().max() * 3
-    print('Gathered')
 
     w = 1# if len(embs) % 2 == 0 else 0
     im_emb = w * coef_.to(dtype=dtype)
@@ -205,7 +196,6 @@ def get_user_emb(embs, ys):
 
 
 def pluck_img(user_id, user_emb):
-    print(user_id, 'user_id')
     not_rated_rows = prevs_df[[i[1]['user:rating'].get(user_id, 'gone') == 'gone' for i in prevs_df.iterrows()]]
     while len(not_rated_rows) == 0:
         not_rated_rows = prevs_df[[i[1]['user:rating'].get(user_id, 'gone') == 'gone' for i in prevs_df.iterrows()]]
@@ -231,7 +221,6 @@ def background_next_image():
         #    not_rated_rows = prevs_df[[i[1]['user:rating'] == {' ': ' '} for i in prevs_df.iterrows()]]
             rated_rows = prevs_df[[i[1]['user:rating'] != {' ': ' '} for i in prevs_df.iterrows()]]
             time.sleep(.01)
-            print('all users have 4 or less rows rated')
 
         user_id_list = set(rated_rows['latest_user_to_rate'].to_list())
         for uid in user_id_list:
@@ -253,15 +242,12 @@ def background_next_image():
                 continue
             
             if len(rated_rows) < 4:
-                print(f'latest user {uid} has < 4 rows') # or > 7 unrated rows')
                 continue
             
-            print(uid)
             embs, ys = pluck_embs_ys(uid)
             
             user_emb = get_user_emb(embs, ys)
             img, embs = generate(user_emb)
-            print(img)
             if img:
                 tmp_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating', 'latest_user_to_rate'])
                 tmp_df['paths'] = [img]
@@ -276,16 +262,10 @@ def background_next_image():
                     cands['sum_bad_ratings'] = [sum([int(t==0) for t in i.values()]) for i in cands['user:rating']]
                     worst_row = cands.loc[cands['sum_bad_ratings']==cands['sum_bad_ratings'].max()].iloc[0]
                     worst_path = worst_row['paths']
-                    print('Removing worst row:', worst_row, 'from prevs_df of len', len(prevs_df))
                     if os.path.isfile(worst_path):
                         os.remove(worst_path)
-                    else:
-                        # If it fails, inform the user.
-                        print("Error: %s file not found" % worst_path)
-
                     # only keep x images & embeddings & ips, then remove the most often disliked besides calibrating
                     prevs_df = prevs_df[prevs_df['paths'] != worst_path]
-                    print('prevs_df is now length:', len(prevs_df))
 
 def pluck_embs_ys(user_id):
     rated_rows = prevs_df[[i[1]['user:rating'].get(user_id, None) != None for i in prevs_df.iterrows()]]
@@ -298,21 +278,17 @@ def pluck_embs_ys(user_id):
     
     embs = rated_rows['embeddings'].to_list()
     ys = [i[user_id] for i in rated_rows['user:rating'].to_list()]
-    print('embs', 'ys', embs, ys)
     return embs, ys
 
 def next_image(calibrate_prompts, user_id):
-    print(prevs_df)
     
     with torch.no_grad():
         if len(calibrate_prompts) > 0:
-            print('######### Calibrating with sample media #########')
             cal_video = calibrate_prompts.pop(0)
             image = prevs_df[prevs_df['paths'] == cal_video]['paths'].to_list()[0]
             
             return image, calibrate_prompts
         else:
-            print('######### Roaming #########')
             embs, ys = pluck_embs_ys(user_id)
             user_emb = get_user_emb(embs, ys)
             image = pluck_img(user_id, user_emb)
@@ -355,7 +331,6 @@ def choose(img, choice, calibrate_prompts, user_id, request: gr.Request):
     # if we detected NSFW, leave that area of latent space regardless of how they rated chosen.
     # TODO skip allowing rating & just continue
     if img == None:
-        print('NSFW -- choice is disliked')
         choice = 0
     
     row_mask = [p.split('/')[-1] in img for p in prevs_df['paths'].to_list()]
@@ -425,7 +400,6 @@ with gr.Blocks(css=css, head=js_head) as demo:
 Explore the latent space without text prompts based on your preferences. Learn more on [the write-up](https://rynmurdock.github.io/posts/2024/3/generative_recomenders/).
     ''', elem_id="description")
     user_id = gr.State()
-    print('USER_ID: ',user_id)
     # calibration videos -- this is a misnomer now :D
     calibrate_prompts = gr.State([
     './first.mp4',
@@ -487,7 +461,7 @@ log = logging.getLogger('log_here')
 log.setLevel(logging.ERROR)
 
 scheduler = BackgroundScheduler()
-scheduler.add_job(func=background_next_image, trigger="interval", seconds=.1)
+scheduler.add_job(func=background_next_image, trigger="interval", seconds=.3)
 scheduler.start()
 
 #thread = threading.Thread(target=background_next_image,)