Update README.md

README.md

@@ -88,6 +88,52 @@ APTP
     ├── ...
     └── arch7
 ```
+## Simple Inference Example
+
+Make sure follow the installation instructions in the [Github Repository](https://github.com/rezashkv/diffusion_pruning) to install pdm from source.
+
+```python
+from diffusers import StableDiffusionPipeline, PNDMScheduler
+from pdm.models import HyperStructure, StructureVectorQuantizer, UNet2DConditionModelPruned
+from pdm.utils.data_utils import get_mpnet_embeddings
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+prompt_encoder_model_name_or_path = "sentence-transformers/all-mpnet-base-v2"
+prompt_encoder_tokenizer = AutoTokenizer.from_pretrained(prompt_encoder_model_name_or_path)
+prompt_encoder = AutoModel.from_pretrained(prompt_encoder_model_name_or_path)
+
+aptp_model_name_or_path = f"rezashkv/APTP"
+aptp_variant = "APTP-Base-CC3M"
+hyper_net = HyperStructure.from_pretrained(aptp_model_name_or_path, subfolder=f"{aptp_variant}/hypernet")
+quantizer = StructureVectorQuantizer.from_pretrained(aptp_model_name_or_path, subfolder=f"{aptp_variant}/quantizer")
+
+prompts = ["a woman on a white background looks down and away from the camera the a forlorn look on her face"]
+prompt_embedding = get_mpnet_embeddings(prompts, prompt_encoder, prompt_encoder_tokenizer)
+
+arch_embedding = hyper_net(prompt_embedding)
+expert_id = quantizer.get_cosine_sim_min_encoding_indices(arch_embedding)[0].item()
+
+sd_model_name_or_path = "stabilityai/stable-diffusion-2-1"
+
+unet = UNet2DConditionModelPruned.from_pretrained(aptp_model_name_or_path,
+                                                  subfolder=f"{aptp_variant}/arch{expert_id}/checkpoint-30000/unet")
+
+noise_scheduler = PNDMScheduler.from_pretrained(sd_model_name_or_path, subfolder="scheduler")
+pipeline = StableDiffusionPipeline.from_pretrained(sd_model_name_or_path, unet=unet, scheduler=noise_scheduler)
+
+pipeline.to('cuda')
+generator = torch.Generator(device='cuda').manual_seed(43)
+
+image = pipeline(
+    prompt=prompts[0],
+    guidance_scale=7.5,
+    generator=generator,
+    output_type='pil',
+).images[0]
+
+image.save("image.png")
+```
 
 
 ## Uses