img to rgb

README.md

@@ -67,17 +67,19 @@ from huggingface_hub import hf_hub_download
 
 # This is an example image we provide
 path = hf_hub_download(repo_id="StonyBrook-CVLab/PixCell-256", filename="test_image.png")
-image = Image.open(path)
+image = Image.open(path).convert("RGB")
 
 # Extract UNI embedding from the image
 uni_inp = transform(image).unsqueeze(dim=0)
 with torch.inference_mode():
     uni_emb = uni_model(uni_inp.to(device))
 
+# reshape UNI to (bs, 1, D)
+uni_emb = uni_emb.unsqueeze(1)
 print("Extracted UNI:", uni_emb.shape)
 
 # Get unconditional embedding for classifier-free guidance
 uncond = pipeline.get_unconditional_embedding(uni_emb.shape[0])
 # Generate new samples
-samples = pipeline(uni_embeds=uni_emb, negative_uni_embeds=uncond, guidance_scale=3., num_images_per_prompt=1)
+samples = pipeline(uni_embeds=uni_emb, negative_uni_embeds=uncond, guidance_scale=3., num_images_per_prompt=1).images
 ```