Upload folder using huggingface_hub

.gitattributes

@@ -34,3 +34,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 test_image.jpg filter=lfs diff=lfs merge=lfs -text
+test_image.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -66,30 +66,14 @@ from PIL import Image
 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.jpg")
+path = hf_hub_download(repo_id="StonyBrook-CVLab/PixCell-256", filename="test_image.png")
 image = Image.open(path)
 
-# Extract UNI from random patches in the image
-n_patches = 8
-patches = []
-uni_emb = []
-for k in range(n_patches):
-    # Extract random crop
-    sz = pipeline.transformer.config.sample_size * pipeline.vae_scale_factor
+# 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)).unsqueeze(dim=0)
 
-    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)
-
-    # For 256x256 directly pass through UNI
-    uni_image = transform(image_patch).unsqueeze(dim=0)
-    with torch.inference_mode():
-        feature_emb = uni_model(uni_image.to(device))
-    uni_emb.append(feature_emb)
-
-uni_emb = torch.stack(uni_emb, dim=0)
 print("Extracted UNI:", uni_emb.shape)
 
 # Get unconditional embedding for classifier-free guidance