updates

app.py

@@ -8,7 +8,9 @@ import glob
 import config
 from model import get_model_and_tokenizer
 
-model, model.prior_pipe.image_encoder = get_model_and_tokenizer(config.model_path, 
+torch.set_float32_matmul_precision('high')
+
+model, model.prior_pipe.image_encoder = get_model_and_tokenizer(config.model_path,
                                                                 'cuda', torch.bfloat16)
 
 # TODO unify/merge origin and this
@@ -16,6 +18,7 @@ model, model.prior_pipe.image_encoder = get_model_and_tokenizer(config.model_pat
 
 device = "cuda"
 
+k = config.k
 
 import spaces
 import matplotlib.pyplot as plt
@@ -51,14 +54,14 @@ def generate_gpu(in_im_embs, prompt='the scene'):
     with torch.no_grad():
         in_im_embs = in_im_embs.to('cuda')
 
-        negative_image_embeds = in_im_embs[0] # model.prior_pipe.get_zero_embed()
+        negative_image_embeds = in_im_embs[0]# if random.random() < .3 else model.prior_pipe.get_zero_embed()
         positive_image_embeds = in_im_embs[1]
 
         images = model.kandinsky_pipe(
             num_inference_steps=50,
             image_embeds=positive_image_embeds,
             negative_image_embeds=negative_image_embeds,
-            guidance_scale=15,
+            guidance_scale=8,
         ).images[0]
         cond = (
                     model.prior_pipe.image_processor(images, return_tensors="pt")
@@ -91,11 +94,10 @@ def generate(in_im_embs, ):
 @spaces.GPU()
 def sample_embs(prompt_embeds):
     latent = torch.randn(prompt_embeds.shape[0], 1, prompt_embeds.shape[-1])
-    if prompt_embeds.shape[1] < 8: # TODO grab as `k` arg from config
-            prompt_embeds = torch.nn.functional.pad(prompt_embeds, [0, 0, 0, 8-prompt_embeds.shape[1]])
-    assert prompt_embeds.shape[1] == 8, f"The model is set to take `k`` cond image embeds but is shape {prompt_embeds.shape}"
+    if prompt_embeds.shape[1] < k:
+            prompt_embeds = torch.nn.functional.pad(prompt_embeds, [0, 0, 0, k-prompt_embeds.shape[1]])
+    assert prompt_embeds.shape[1] == k, f"The model is set to take `k`` cond image embeds but is shape {prompt_embeds.shape}"
     image_embeds = model(latent.to('cuda'), prompt_embeds.to('cuda')).predicted_image_embedding
-
     return image_embeds
 
 @spaces.GPU()
@@ -113,6 +115,8 @@ def get_user_emb(embs, ys):
     else:
         negative_embs = random.sample(negs, min(4, len(negs))) + negs[-4:]
         negatives = torch.stack(negative_embs, 1)
+        # if random.random() < .5:
+        #     negatives = torch.zeros_like(negatives)
 
     image_embeds = torch.stack([sample_embs(negatives), sample_embs(positives)])
 
@@ -175,6 +179,7 @@ def background_next_image():
                     prevs_df = pd.concat((prevs_df.iloc[:6], prevs_df.iloc[7:]))
     
 def pluck_img(user_id):
+    # TODO pluck images based on similarity but also based on diversity by cluster every few times.
     rated_rows = prevs_df[[i[1]['user:rating'].get(user_id, None) is not None for i in prevs_df.iterrows()]]
     ems = rated_rows['embeddings'].to_list()
     ys = [i[user_id][0] for i in rated_rows['user:rating'].to_list()]

config.py

@@ -12,5 +12,5 @@ batch_size = 16
 number_k_clip_embed = 16 # divide by this to determine bundling together of sequences -> CLIP
 num_workers = 32
 seed = 107
-
+k = 8
 # TODO config option to swap to diffusion?

last_epoch_ckpt/diffusion_pytorch_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d65a902c101345526b244420a5e6f495a947909db28015840afa9bacd557936b
+oid sha256:ae34b5c319b9c804e1e82c93f78821b880553d2ac60ff628003175334ee9066d
 size 136790920

prior/pipeline_kandinsky_prior.py

@@ -498,14 +498,13 @@ class KandinskyPriorPipeline(DiffusionPipeline):
         if negative_prompt is None:
             # zero_embeds = self.get_zero_embed(latents.shape[0], device=latents.device)
 
-            # using the same hidden states or different hidden states?
-            
-            hidden_states = torch.randn(
-                (batch_size, prompt_embeds.shape[-1]),
-                device=prompt_embeds.device,
-                dtype=prompt_embeds.dtype,
-                generator=generator,
-            )
+            # use the same hidden states or different hidden states?
+            # hidden_states = torch.randn(
+            #     (batch_size, prompt_embeds.shape[-1]),
+            #     device=prompt_embeds.device,
+            #     dtype=prompt_embeds.dtype,
+            #     generator=generator,
+            # )
 
             latents = self.prior(
                 hidden_states,
@@ -541,7 +540,7 @@ class KandinskyPriorPipeline(DiffusionPipeline):
 
         if not return_dict:
             return (image_embeddings, zero_embeds)
-
+        
         return KandinskyPriorPipelineOutput(
             image_embeds=image_embeddings, negative_image_embeds=zero_embeds
         )