Update app.py

app.py

@@ -125,8 +125,8 @@ pipe.to(device=DEVICE)
 
 from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
 
-pali = PaliGemmaForConditionalGeneration.from_pretrained('google/paligemma-3b-pt-224', torch_dtype=dtype).eval().to('cuda')
-processor = AutoProcessor.from_pretrained('google/paligemma-3b-pt-224')
+pali = PaliGemmaForConditionalGeneration.from_pretrained('google/paligemma-3b-mix-224', torch_dtype=dtype).eval().to('cuda')
+processor = AutoProcessor.from_pretrained('google/paligemma-3b-mix-224')
 
 #pali = torch.compile(pali)
 
@@ -162,7 +162,7 @@ def generate_pali(n_embs):
             decoded = processor.decode(generation[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
             descs += f'Description: {decoded}\n'
         else:
-            prompt = f'en {descs} Describe a new image that is similar.'
+            prompt = f'en {descs} Describe a new image that is similar. Description:'
             model_inputs = processor(text=prompt, images=torch.zeros(1, 3, 224, 224), return_tensors="pt")
             input_len = model_inputs["input_ids"].shape[-1]
             input_embeds = to_wanted_embs(emb,