Update app/inference.py

app/inference.py

@@ -20,11 +20,11 @@ def generate_image_with_clip_score(prompt, num_inference_steps=50, guidance_scal
     model_path = os.path.join(model_path, "epoch_4")
 
     unet = UNet2DConditionModel.from_pretrained(os.path.join(model_path, "unet"), use_safetensors=True).to(device)
-    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
-    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to(device)
+    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32",cache_dir="/tmp/huggingface").to(device)
+    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32",cache_dir="/tmp/huggingface")
+    vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse",cache_dir="/tmp/huggingface").to(device)
     scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="linear")
-    clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
+    clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch32",cache_dir="/tmp/huggingface").to(device)
 
     checkpoint = torch.load(os.path.join(model_path, "training_state.pth"), map_location=device)
     unet.load_state_dict(checkpoint['model_state_dict'])