merge

app.py

@@ -123,20 +123,11 @@ pipe.unet.fuse_qkv_projections()
 pipe.to(device=DEVICE)
 #pipe.unet = torch.compile(pipe.unet)
 #pipe.vae = torch.compile(pipe.vae)
-# TODO cannot compile on Spaces or we time out; don't run leave_imb stuff either
-#im_embs = torch.zeros(1, 1, 1, 1280, device=DEVICE, dtype=dtype)
-#output = pipe(prompt='a person', guidance_scale=0, added_cond_kwargs={}, ip_adapter_image_embeds=[im_embs], num_inference_steps=STEPS)
-#leave_im_emb, _ = pipe.encode_image(
-#                output.frames[0][len(output.frames[0])//2], DEVICE, 1, output_hidden_state
-#)
-#assert len(output.frames[0]) == 16
-#leave_im_emb.detach().to('cpu')
 
 @spaces.GPU(duration=10)
 def generate_gpu(in_im_embs):
     print('start gen')
     in_im_embs = in_im_embs.to('cuda').unsqueeze(0).unsqueeze(0)
-    #im_embs = torch.cat((torch.zeros(1, 1280, device=DEVICE, dtype=dtype), in_im_embs), 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')
@@ -148,7 +139,6 @@ def generate_gpu(in_im_embs):
     return output, im_emb
 
 def generate(in_im_embs):
-    
     output, im_emb = generate_gpu(in_im_embs)
     nsfw = maybe_nsfw(output.frames[0][len(output.frames[0])//2])
     
@@ -358,8 +348,7 @@ def choose(img, choice, calibrate_prompts, user_id, request: gr.Request):
         choice = 0
     
     row_mask = [p.split('/')[-1] in img for p in prevs_df['paths'].to_list()]
-    
-    
+    # if it's still in the dataframe, add the choice
     if len(prevs_df.loc[row_mask, 'user:rating']) > 0:
         prevs_df.loc[row_mask, 'user:rating'][0][user_id] = choice
         prevs_df.loc[row_mask, 'latest_user_to_rate'] = [user_id]

safety_checker_improved.py

@@ -9,11 +9,12 @@ import sys
 sys.path.append('/home/ryn_mote/Misc/generative_recommender/gradio_video/automl/efficientnetv2/')
 import tensorflow as tf
 from tensorflow.keras import mixed_precision
-#physical_devices = tf.config.list_physical_devices('GPU')
 
-#tf.config.experimental.set_memory_growth(
-#    physical_devices[0], True
-#)
+physical_devices = tf.config.list_physical_devices('GPU')
+if len(physical_devices) > 0:
+    tf.config.experimental.set_memory_growth(
+        physical_devices[0], True
+    )
 
 model = tf.keras.models.load_model('nsfweffnetv2-b02-3epochs.h5',custom_objects={"KerasLayer":hub.KerasLayer})
 # "The image classifier had been trained on 682550 images from the 5 classes "Drawing" (39026), "Hentai" (28134), "Neutral" (369507), "Porn" (207969) & "Sexy" (37914).