Update README.md

README.md

@@ -1,3 +1,109 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+---
+---
+license: apache-2.0
+---
+
+### Load PixCell-1024 model
+
+```python
+import torch
+
+from diffusers import DiffusionPipeline
+from diffusers import AutoencoderKL
+
+device = torch.device('cuda')
+
+# We do not host the weights of the SD3 VAE -- load it from StabilityAI
+sd3_vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-3.5-large", subfolder="vae")
+
+pipeline = DiffusionPipeline.from_pretrained(
+    "StonyBrook-CVLab/pixcell-1024-diffusers",
+    vae=sd3_vae,
+    custom_pipeline="StonyBrook-CVLab/pixcell-pipeline",
+    trust_remote_code=True,
+)
+
+pipeline.to(device);
+```
+
+### Load [[UNI-2h]](https://huggingface.co/MahmoodLab/UNI2-h) for conditioning
+```python
+import timm
+from timm.data import resolve_data_config
+from timm.data.transforms_factory import create_transform
+
+timm_kwargs = {
+            'img_size': 224, 
+            'patch_size': 14, 
+            'depth': 24,
+            'num_heads': 24,
+            'init_values': 1e-5, 
+            'embed_dim': 1536,
+            'mlp_ratio': 2.66667*2,
+            'num_classes': 0, 
+            'no_embed_class': True,
+            'mlp_layer': timm.layers.SwiGLUPacked, 
+            'act_layer': torch.nn.SiLU, 
+            'reg_tokens': 8, 
+            'dynamic_img_size': True
+        }
+uni_model = timm.create_model("hf-hub:MahmoodLab/UNI2-h", pretrained=True, **timm_kwargs)
+transform = create_transform(**resolve_data_config(uni_model.pretrained_cfg, model=uni_model))
+uni_model.eval()
+uni_model.to(device);
+```
+
+### Unconditional generation
+```python
+uncond = pipeline.get_unconditional_embedding(1)
+samples = pipeline(uni_embeds=uncond, negative_uni_embeds=None, guidance_scale=1.0)
+```
+
+### Conditional generation
+```python
+# Load image
+import numpy as np
+import einops
+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-1024-diffusers", filename="test_image.jpg")
+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)
+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)
+```