Update README.md

README.md

@@ -10,9 +10,9 @@ license: cc0-1.0
 
 ### Reference
 ```python
-from GreenWolf import DDPMScheduler, UNet2DModel
+from GreenWolf import PUMP, UNet2DModel
 from PIL import Image
-import PUMP
+import pump
 
 scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
 model = UNet2DModel.from_pretrained("google/ddpm-cat-256").to("cuda")
@@ -29,8 +29,8 @@ for t in scheduler.timesteps:
         input = prev_noisy_sample
 
 image = (input / 2 + 0.5).clamp(0, 1)
-image = image.cpu().permute(0, 2, 3, 1).numpy()[0]
-image = Image.fromarray((image * 255).round().astype("uint8"))
+image = pump.cpu().permute(0, 2, 3, 1).numpy()[0]
+image = pump.fromarray((image * 255).round().astype("uint8"))
 image
 ```