You can deploy this model and start earning money today!

StonyBrook-CVLab/
PixCell-256-Cell-ControlNet

Name: PixCell-256-Cell-ControlNet
Rating: 2.5 (38 reviews)

PixCellControlNetPipeline

Views

No views yet

PixCell: A generative foundation model for digital histopathology images

[📄 arXiv][🔬 PixCell-1024] [🔬 PixCell-256] [🔬 Pixcell-256-Cell-ControlNet] [💾 Synthetic SBU-1M]

Load PixCell-256-Cell-ControlNet 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-256-Cell-ControlNet",
    vae=sd3_vae,
    custom_pipeline="StonyBrook-CVLab/PixCell-pipeline-ControlNet",
    trust_remote_code=True,
)

pipeline.to(device);

Load [UNI-2h] 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)
uni_transforms = create_transform(**resolve_data_config(uni_model.pretrained_cfg, model=uni_model))
uni_model.eval()
uni_model.to(device);

Mask-conditioned generation

python

# Load image
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

# This is an example image/mask pair we provide
image_path = hf_hub_download(repo_id="StonyBrook-CVLab/PixCell-256-Cell-ControlNet", filename="test_image.png")
mask_path = hf_hub_download(repo_id="StonyBrook-CVLab/PixCell-256-Cell-ControlNet", filename="test_mask.png")
image = Image.open(image_path).convert("RGB")
mask = np.asarray(Image.open(mask_path).convert("RGB"))

# Extract UNI embedding from the image
uni_inp = uni_transforms(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 using the given mask
samples = pipeline(uni_embeds=uni_emb, controlnet_input=mask, negative_uni_embeds=uncond, guidance_scale=2.5, num_images_per_prompt=1).images

License & Usage

License: Creative Commons Attribution-NonCommercial-NoDerivatives 4.0 International (CC BY-NC-ND 4.0)

Notice: This model is a derivative work conditioned on embeddings from the [UNI-2h] foundation model. As such, it is subject to the original terms of the UNI2 license.

Academic & Research Use Only: You may use these weights for non-commercial research purposes.
No Commercial Use: You may not use this model for any commercial purpose, including product development or commercial services.