Update app.py

app.py

@@ -16,10 +16,10 @@ def load_artifacts():
     '''
     A function to load all diffusion artifacts
     '''
-    vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float32,use_auth_token=auth_token).to(device)
-    unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float32, use_auth_token=auth_token).to(device)
-    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float32, use_auth_token=auth_token)
-    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float32, use_auth_token=auth_token).to(device)
+    vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", torch_dtype=torch.float16,use_auth_token=auth_token).to(device)
+    unet = UNet2DConditionModel.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="unet", torch_dtype=torch.float16, use_auth_token=auth_token).to(device)
+    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16, use_auth_token=auth_token)
+    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=torch.float16, use_auth_token=auth_token).to(device)
     scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)    
     return vae, unet, tokenizer, text_encoder, scheduler
 
@@ -34,7 +34,7 @@ def pil_to_latents(image):
     Function to convert image to latents
     '''
     init_image = tfms.ToTensor()(image).unsqueeze(0) * 2.0 - 1.0
-    init_image = init_image.to(device=device, dtype=torch.float32) 
+    init_image = init_image.to(device=device, dtype=torch.float16) 
     init_latent_dist = vae.encode(init_image).latent_dist.sample() * 0.18215
     return init_latent_dist
 
@@ -57,7 +57,7 @@ def text_enc(prompts, maxlen=None):
     '''
     if maxlen is None: maxlen = tokenizer.model_max_length
     inp = tokenizer(prompts, padding="max_length", max_length=maxlen, truncation=True, return_tensors="pt") 
-    return text_encoder(inp.input_ids.to(device))[0]
+    return text_encoder(inp.input_ids.to(device))[0].half()
 
 def prompt_2_img_i2i_fast(prompts, init_img, g=7.5, seed=100, strength =0.5, steps=50, dim=512):
     """
@@ -140,7 +140,8 @@ def improve_mask(mask):
 vae, unet, tokenizer, text_encoder, scheduler = load_artifacts()
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "runwayml/stable-diffusion-inpainting",
-    torch_dtype=torch.float32,
+    revision="fp16",
+    torch_dtype=torch.float16,
     use_auth_token=auth_token
 ).to(device)