update readme

README.md

@@ -68,39 +68,25 @@ from huggingface_hub import hf_hub_download
 
 # This is an example image we provide
 path = hf_hub_download(repo_id="StonyBrook-CVLab/PixCell-1024", filename="test_image.png")
-image = Image.open(path)
-
-# Extract UNI from random patches in the image
-n_patches = 1
-patches = []
-uni_emb = []
-for k in range(n_patches):
-    # Extract random crop
-    sz = pipeline.transformer.config.sample_size * pipeline.vae_scale_factor
-
-    x1 = np.random.randint(0, image.size[0] - sz+1)
-    y1 = np.random.randint(0, image.size[1] - sz+1)
-    image_patch = image.crop((x1, y1, x1+sz, y1+sz))
-    patches.append(image_patch)
-    print("Extracted patch:", patches[-1].size)
-
-    # Rearrange 1024x1024 image into 16 256x256 patches
-    uni_patches = np.array(image_patch)
-    uni_patches = einops.rearrange(uni_patches, '(d1 h) (d2 w) c -> (d1 d2) h w c', d1=4, d2=4)
-
-    # Extract UNIs for each patch
-    uni_images = torch.cat(
-        [transform(Image.fromarray(x)).unsqueeze(0) for x in uni_patches],
-        dim=0)
-    with torch.inference_mode():
-        feature_emb = uni_model(uni_images.to(device))
-    uni_emb.append(feature_emb)
-
-uni_emb = torch.stack(uni_emb, dim=0)
+image = Image.open(path).convert("RGB")
+
+
+# Rearrange 1024x1024 image into 16 256x256 patches
+uni_patches = np.array(image)
+uni_patches = einops.rearrange(uni_patches, '(d1 h) (d2 w) c -> (d1 d2) h w c', d1=4, d2=4)
+uni_input = torch.stack([transform(Image.fromarray(item)) for item in uni_patches])
+
+# Extract UNI embeddings
+with torch.inference_mode():
+    uni_emb = uni_model(uni_input.to(device))
+
+# reshape UNI to (bs, 16, D)
+uni_emb = uni_emb.unsqueeze(0)
 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=1.5, num_images_per_prompt=1)
+samples = pipeline(uni_embeds=uni_emb, negative_uni_embeds=uncond, guidance_scale=1.5).images
+
 ```