Update README.md

README.md

@@ -31,13 +31,15 @@ checkpoint = 'takuoko/tiny_sd_xl_pokemon_blip'
 prompt = 'a picture of a pink and yellow pokemon with a sword'
 
 unet = UNet2DConditionModel.from_pretrained(
-    checkpoint, torch_dtype=torch.float16)
+    checkpoint, torch_dtype=torch.bfloat16
+    )
 vae = AutoencoderKL.from_pretrained(
     'madebyollin/sdxl-vae-fp16-fix',
-    torch_dtype=torch.float16,
+    torch_dtype=torch.bfloat16,
 )
 pipe = DiffusionPipeline.from_pretrained(
-    'stabilityai/stable-diffusion-xl-base-1.0', unet=unet, vae=vae, torch_dtype=torch.float16)
+    'stabilityai/stable-diffusion-xl-base-1.0', unet=unet, vae=vae, torch_dtype=torch.bfloat16
+    )
 pipe.to('cuda')
 
 image = pipe(
@@ -49,6 +51,6 @@ image.save('demo.png')
 
 # Example result
 
-prompt = 'robotic cat with wings'
+prompt = 'a picture of a pink and yellow pokemon with a sword'
 
 ![image](demo.png)