match multimodlar

app.py

@@ -15,14 +15,12 @@ import matplotlib.pyplot as plt
 import matplotlib
 import logging
 
-from sklearn.linear_model import Ridge
 
 import os
 import imageio
 import gradio as gr
 import numpy as np
 from sklearn.svm import SVC
-from sklearn.inspection import permutation_importance
 from sklearn import preprocessing
 import pandas as pd
 from apscheduler.schedulers.background import BackgroundScheduler
@@ -39,14 +37,13 @@ torch.set_grad_enabled(False)
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 
-prevs_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating', 'latest_user_to_rate', 'from_user_id'])
+prevs_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating', 'latest_user_to_rate', 'from_user_id', 'text', 'gemb'])
 
 import spaces
 start_time = time.time()
 
 ####################### Setup Model
-from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, LCMScheduler, AutoencoderTiny, UNet2DConditionModel, AutoencoderKL, utils
-utils.logging.disable_progress_bar
+from diffusers import AnimateDiffPipeline, MotionAdapter, EulerDiscreteScheduler, LCMScheduler, AutoencoderTiny, UNet2DConditionModel, AutoencoderKL
 from transformers import CLIPTextModel
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
@@ -54,6 +51,7 @@ from PIL import Image
 from transformers import CLIPVisionModelWithProjection
 import uuid
 import av
+import torchvision
 
 def write_video(file_name, images, fps=17):
     container = av.open(file_name, mode="w")
@@ -92,6 +90,9 @@ device_map='cuda')
 # vae = ConsistencyDecoderVAE.from_pretrained("openai/consistency-decoder", torch_dtype=dtype)
 # vae = compile_unet(vae, config=config)
 
+#finetune_path = '''/home/ryn_mote/Misc/finetune-sd1.5/dreambooth-model best'''''
+#unet = UNet2DConditionModel.from_pretrained(finetune_path+'/unet/').to(dtype)
+#text_encoder = CLIPTextModel.from_pretrained(finetune_path+'/text_encoder/').to(dtype)
 
 
 unet = UNet2DConditionModel.from_pretrained('rynmurdock/Sea_Claws', subfolder='unet',).to(dtype).to('cpu')
@@ -99,7 +100,8 @@ text_encoder = CLIPTextModel.from_pretrained('rynmurdock/Sea_Claws', subfolder='
 device_map='cpu').to(dtype)
 
 adapter = MotionAdapter.from_pretrained("wangfuyun/AnimateLCM")
-pipe = AnimateDiffPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", motion_adapter=adapter, image_encoder=image_encoder, torch_dtype=dtype, unet=unet, text_encoder=text_encoder)
+pipe = AnimateDiffPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", motion_adapter=adapter, image_encoder=image_encoder, torch_dtype=dtype,     
+                                            unet=unet, text_encoder=text_encoder)
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config, beta_schedule="linear")
 pipe.load_lora_weights("wangfuyun/AnimateLCM", weight_name="AnimateLCM_sd15_t2v_lora.safetensors", adapter_name="lcm-lora",)
 pipe.set_adapters(["lcm-lora"], [.9])
@@ -114,7 +116,7 @@ pipe.fuse_lora()
 
 pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15_vit-G.bin", map_location='cpu')
 # This IP adapter improves outputs substantially.
-pipe.set_ip_adapter_scale(.9)
+pipe.set_ip_adapter_scale(.6)
 pipe.unet.fuse_qkv_projections()
 #pipe.enable_free_init(method="gaussian", use_fast_sampling=True)
 
@@ -122,21 +124,71 @@ pipe.to(device=DEVICE)
 #pipe.unet = torch.compile(pipe.unet)
 #pipe.vae = torch.compile(pipe.vae)
 
-@spaces.GPU()
-def generate_gpu(in_im_embs):
-    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)
-    im_emb, _ = pipe.encode_image(
-                output.frames[0][len(output.frames[0])//2], 'cuda', 1, output_hidden_state
-            )
-    im_emb = im_emb.detach().to('cpu').to(torch.float32)
-    return output, im_emb
 
+#############################################################
+
+from transformers import AutoProcessor, PaliGemmaForConditionalGeneration, BitsAndBytesConfig
+
+quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+pali = PaliGemmaForConditionalGeneration.from_pretrained('google/paligemma-3b-pt-224', torch_dtype=dtype, quantization_config=quantization_config).eval()
+processor = AutoProcessor.from_pretrained('google/paligemma-3b-pt-224')
 
-def generate(in_im_embs):
-    output, im_emb = generate_gpu(in_im_embs)
-    nsfw = maybe_nsfw(output.frames[0][len(output.frames[0])//2])
+
+
+def to_wanted_embs(image_outputs, input_ids, attention_mask, cache_position=None):
+    inputs_embeds = pali.get_input_embeddings()(input_ids)
+    selected_image_feature = image_outputs.to(dtype).to(device)
+    image_features = pali.multi_modal_projector(selected_image_feature)
+
+    if cache_position is None:
+        cache_position = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device)
+    inputs_embeds, attention_mask, labels, position_ids = pali._merge_input_ids_with_image_features(
+        image_features, inputs_embeds, input_ids, attention_mask, None, None, cache_position
+    )
+    return inputs_embeds
     
+
+
+def generate_pali(user_emb):
+    prompt = 'caption en'
+    model_inputs = processor(text=prompt, images=torch.zeros(1, 3, 224, 224), return_tensors="pt")
+    # we need to get im_embs taken in here.
+    input_len = model_inputs["input_ids"].shape[-1]
+    input_embeds = to_wanted_embs(user_emb.squeeze()[None, None, :].repeat(1, 256, 1), 
+                        model_inputs["input_ids"].to(device), 
+                        model_inputs["attention_mask"].to(device))
+
+    generation = pali.generate(max_new_tokens=100, do_sample=True, top_p=.94, temperature=1.2, inputs_embeds=input_embeds)
+    decoded = processor.decode(generation[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    return decoded
+
+
+
+
+#############################################################
+
+
+
+@spaces.GPU()
+def generate_gpu(in_im_embs, prompt='the scene'):
+    with torch.no_grad():
+        in_im_embs = in_im_embs.to('cuda').unsqueeze(0).unsqueeze(0)
+        output = pipe(prompt=prompt, guidance_scale=1, added_cond_kwargs={}, ip_adapter_image_embeds=[in_im_embs], num_inference_steps=STEPS)
+        im_emb, _ = pipe.encode_image(
+                    output.frames[0][len(output.frames[0])//2], 'cuda', 1, output_hidden_state
+                )
+        im_emb = im_emb.detach().to('cpu').to(torch.float32)
+        im = torchvision.transforms.ToTensor()(output.frames[0][len(output.frames[0])//2]).unsqueeze(0)
+        im = torch.nn.functional.interpolate(im, (224, 224))
+        im = (im - .5) * 2
+        gemb = pali.vision_tower(im.to(device).to(dtype)).last_hidden_state.detach().to('cpu').to(torch.float32).mean(1)
+    return output, im_emb, gemb
+
+
+def generate(in_im_embs, prompt='the scene'):
+    output, im_emb, gemb = generate_gpu(in_im_embs, prompt)
+    nsfw =maybe_nsfw(output.frames[0][len(output.frames[0])//2])
+    print(prompt)
     name = str(uuid.uuid4()).replace("-", "")
     path = f"/tmp/{name}.mp4"
     
@@ -149,19 +201,19 @@ def generate(in_im_embs):
     output.frames[0] = output.frames[0] + list(reversed(output.frames[0]))
 
     write_video(path, output.frames[0])
-    return path, im_emb
+    return path, im_emb, gemb
 
 
 #######################
 
 def get_user_emb(embs, ys):
     # handle case where every instance of calibration videos is 'Neither' or 'Like' or 'Dislike'
+    
     if len(list(ys)) <= 7:
-        aways = [.01*torch.randn(1280) for i in range(3)]
+        aways = [.01*torch.randn_like(embs[0]) for i in range(3)]
         embs += aways
         awal = [0 for i in range(3)]
         ys += awal
-        print('Fixing only one feedback class available.\n')
     
     indices = list(range(len(embs)))
     # sample only as many negatives as there are positives
@@ -176,21 +228,20 @@ def get_user_emb(embs, ys):
     # 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 = torch.stack([embs[i].squeeze().to('cpu') for i in indices]).to('cpu')
     #scaler = preprocessing.StandardScaler().fit(feature_embs)
     #feature_embs = scaler.transform(feature_embs)
-    
+    chosen_y = np.array([ys[i] for i in indices])
     
     if feature_embs.norm() != 0:
         feature_embs = feature_embs / feature_embs.norm()
     
-    chosen_y = np.array([ys[i] for i in indices])
-    
     #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')
+    lin_class = SVC(max_iter=20, kernel='linear', C=.1, class_weight='balanced').fit(feature_embs.squeeze(), chosen_y)
+    coef_ = torch.tensor(lin_class.coef_, dtype=torch.float32).detach().to('cpu')
     coef_ = coef_ / coef_.abs().max() * 3
 
     w = 1# if len(embs) % 2 == 0 else 0
@@ -212,7 +263,8 @@ def pluck_img(user_id, user_emb):
             best_sim = sim
             best_row = i[1]
     img = best_row['paths']
-    return img
+    text = best_row.get('text', '')
+    return img, text
 
 
 def background_next_image():
@@ -236,39 +288,48 @@ def background_next_image():
             unrated_from_user = not_rated_rows[[i[1]['from_user_id'] == uid for i in not_rated_rows.iterrows()]]
             rated_from_user = rated_rows[[i[1]['from_user_id'] == uid for i in rated_rows.iterrows()]]
 
-            # we pop previous ratings if there are > 10
-            if len(rated_from_user) >= 10:
+            # we pop previous ratings if there are > n
+            if len(rated_from_user) >= 15:
                 oldest = rated_from_user.iloc[0]['paths']
                 prevs_df = prevs_df[prevs_df['paths'] != oldest]
-            # we don't compute more after 10 are in the queue for them
+            # we don't compute more after n are in the queue for them
             if len(unrated_from_user) >= 10:
                 continue
             
-            if len(rated_rows) < 4:
+            if len(rated_rows) < 5:
                 continue
             
-            embs, ys = pluck_embs_ys(uid)
+            embs, ys, gembs = pluck_embs_ys(uid)
             
             user_emb = get_user_emb(embs, ys)
-            img, embs = generate(user_emb)
+                        
+            if len(gembs) > 4:
+                user_gem = get_user_emb(gembs, ys) / 4 # TODO scale this correctly; matplotlib, etc.
+                text = generate_pali(user_gem)
+            else:
+                text = generate_pali(torch.zeros(1, 1152))
+            img, embs, new_gem = generate(user_emb, text)
+            
             if img:
-                tmp_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating', 'latest_user_to_rate'])
+                tmp_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating', 'latest_user_to_rate', 'text', 'gemb'])
                 tmp_df['paths'] = [img]
                 tmp_df['embeddings'] = [embs]
                 tmp_df['user:rating'] = [{' ': ' '}]
                 tmp_df['from_user_id'] = [uid]
+                tmp_df['text'] = [text]
+                tmp_df['gemb'] = [new_gem]
                 prevs_df = pd.concat((prevs_df, tmp_df))
-                
                 # we can free up storage by deleting the image
-                if len(prevs_df) > 30:
-                    cands = prevs_df.iloc[6:]
-                    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']
-                    if os.path.isfile(worst_path):
-                        os.remove(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]
+                if len(prevs_df) > 500:
+                    oldest_path = prevs_df.iloc[6]['paths']
+                    if os.path.isfile(oldest_path):
+                        os.remove(oldest_path)
+                    else:
+                        # If it fails, inform the user.
+                        print("Error: %s file not found" % oldest_path)
+                    # only keep 50 images & embeddings & ips, then remove oldest besides calibrating
+                    prevs_df = pd.concat((prevs_df.iloc[:6], prevs_df.iloc[7:]))
+    
 
 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()]]
@@ -281,21 +342,21 @@ 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()]
-    return embs, ys
+    gembs = rated_rows['gemb'].to_list()
+    return embs, ys, gembs
 
 def next_image(calibrate_prompts, user_id):
-    
     with torch.no_grad():
         if len(calibrate_prompts) > 0:
             cal_video = calibrate_prompts.pop(0)
             image = prevs_df[prevs_df['paths'] == cal_video]['paths'].to_list()[0]
             
-            return image, calibrate_prompts
+            return image, calibrate_prompts, ''
         else:
-            embs, ys = pluck_embs_ys(user_id)
+            embs, ys, gembs = pluck_embs_ys(user_id)
             user_emb = get_user_emb(embs, ys)
-            image = pluck_img(user_id, user_emb)
-            return image, calibrate_prompts
+            image, text = pluck_img(user_id, user_emb)
+            return image, calibrate_prompts, text
 
 
 
@@ -307,7 +368,7 @@ def next_image(calibrate_prompts, user_id):
 
 def start(_, calibrate_prompts, user_id, request: gr.Request):
     user_id = int(str(time.time())[-7:].replace('.', ''))
-    image, calibrate_prompts = next_image(calibrate_prompts, user_id)
+    image, calibrate_prompts, text = next_image(calibrate_prompts, user_id)
     return [
             gr.Button(value='Like (L)', interactive=True), 
             gr.Button(value='Neither (Space)', interactive=True, visible=False), 
@@ -326,14 +387,15 @@ def choose(img, choice, calibrate_prompts, user_id, request: gr.Request):
     if choice == 'Like (L)':
         choice = 1
     elif choice == 'Neither (Space)':
-        img, calibrate_prompts = next_image(calibrate_prompts, user_id)
-        return img, calibrate_prompts
+        img, calibrate_prompts, text = next_image(calibrate_prompts, user_id)
+        return img, calibrate_prompts, text
     else:
         choice = 0
     
     # 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()]
@@ -341,8 +403,8 @@ def choose(img, choice, calibrate_prompts, user_id, request: gr.Request):
     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]
-    img, calibrate_prompts = next_image(calibrate_prompts, user_id)
-    return img, calibrate_prompts
+    img, calibrate_prompts, text = next_image(calibrate_prompts, user_id)
+    return img, calibrate_prompts, text
 
 css = '''.gradio-container{max-width: 700px !important}
 #description{text-align: center}
@@ -426,6 +488,8 @@ Explore the latent space without text prompts based on your preferences. Learn m
         elem_id="video_output"
        )
         img.play(l, js='''document.querySelector('[data-testid="Lightning-player"]').loop = true''')
+    with gr.Row():
+        text = gr.Textbox(interactive=False, visible=True, label='Text')
     with gr.Row(equal_height=True):
         b3 = gr.Button(value='Dislike (A)', interactive=False, elem_id="dislike")
         b2 = gr.Button(value='Neither (Space)', interactive=False, elem_id="neither", visible=False)
@@ -433,17 +497,17 @@ Explore the latent space without text prompts based on your preferences. Learn m
         b1.click(
         choose, 
         [img, b1, calibrate_prompts, user_id],
-        [img, calibrate_prompts],
+        [img, calibrate_prompts, text],
         )
         b2.click(
         choose, 
         [img, b2, calibrate_prompts, user_id],
-        [img, calibrate_prompts],
+        [img, calibrate_prompts, text],
         )
         b3.click(
         choose, 
         [img, b3, calibrate_prompts, user_id],
-        [img, calibrate_prompts],
+        [img, calibrate_prompts, text],
         )
     with gr.Row():
         b4 = gr.Button(value='Start')
@@ -464,20 +528,28 @@ log = logging.getLogger('log_here')
 log.setLevel(logging.ERROR)
 
 scheduler = BackgroundScheduler()
-scheduler.add_job(func=background_next_image, trigger="interval", seconds=.3)
+scheduler.add_job(func=background_next_image, trigger="interval", seconds=.5)
 scheduler.start()
 
 #thread = threading.Thread(target=background_next_image,)
 #thread.start()
 
+# TODO shouldn't call this before gradio launch, yeah?
 @spaces.GPU()
 def encode_space(x):
     im_emb, _ = pipe.encode_image(
                 image, DEVICE, 1, output_hidden_state
             )
-    return im_emb.detach().to('cpu').to(torch.float32)
+            
+            
+    im = torchvision.transforms.ToTensor()(x).unsqueeze(0)
+    im = torch.nn.functional.interpolate(im, (224, 224))
+    im = (im - .5) * 2
+    gemb = pali.vision_tower(im.to(device).to(dtype)).last_hidden_state.detach().to('cpu').to(torch.float32).mean(1)
+            
+    return im_emb.detach().to('cpu').to(torch.float32), gemb
 
-# prep our calibration prompts
+# prep our calibration videos
 for im in [
     './first.mp4',
     './second.mp4',
@@ -485,16 +557,23 @@ for im in [
     './fourth.mp4',
     './fifth.mp4',
     './sixth.mp4',
+    './seventh.mp4',
+    './eigth.mp4',
+    './ninth.mp4',
+    './tenth.mp4',
     ]:
-    tmp_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating'])
+    tmp_df = pd.DataFrame(columns=['paths', 'embeddings', 'ips', 'user:rating', 'text', 'gemb'])
     tmp_df['paths'] = [im]
     image = list(imageio.imiter(im))
     image = image[len(image)//2]
-    im_emb = encode_space(image)
+    im_emb, gemb = encode_space(image)
 
     tmp_df['embeddings'] = [im_emb.detach().to('cpu')]
+    tmp_df['gemb'] = [gemb.detach().to('cpu')]
     tmp_df['user:rating'] = [{' ': ' '}]
     prevs_df = pd.concat((prevs_df, tmp_df))
 
 
-demo.launch(share=True)
+demo.launch(share=True, server_port=8443)
+
+