Update demo.py with the released checkpoint

demo.py

@@ -5,7 +5,7 @@ import copy
 import time
 
 ORIGINAL_CHECKPOINT_ID = "CompVis/stable-diffusion-v1-4"
-COMPRESSED_UNET_PATH = "checkpoints/BK-SDM-Small_iter50000"
+COMPRESSED_UNET_ID = "nota-ai/bk-sdm-small"
 
 DEVICE='cuda'
 # DEVICE='cpu'
@@ -18,7 +18,7 @@ class SdmCompressionDemo:
         self.pipe_original = StableDiffusionPipeline.from_pretrained(ORIGINAL_CHECKPOINT_ID,
                                                                      torch_dtype=self.torch_dtype)        
         self.pipe_compressed = copy.deepcopy(self.pipe_original)
-        self.pipe_compressed.unet = UNet2DConditionModel.from_pretrained(COMPRESSED_UNET_PATH, 
+        self.pipe_compressed.unet = UNet2DConditionModel.from_pretrained(COMPRESSED_UNET_ID, 
                                                                          subfolder="unet",
                                                                          torch_dtype=self.torch_dtype)
         if 'cuda' in self.device: