Update app.py

app.py

@@ -6,36 +6,6 @@ import torch
 import torchvision.transforms as transforms
 import requests
 
-# Function to download the model from Google Drive
-def download_file_from_google_drive(id, destination):
-    URL = "https://drive.google.com/uc?export=download"
-    session = requests.Session()
-    response = session.get(URL, params={'id': id}, stream=True)
-    token = get_confirm_token(response)
-
-    if token:
-        params = {'id': id, 'confirm': token}
-        response = session.get(URL, params=params, stream=True)
-
-    save_response_content(response, destination)
-
-def get_confirm_token(response):
-    for key, value in response.cookies.items():
-        if key.startswith('download_warning'):
-            return value
-    return None
-
-def save_response_content(response, destination):
-    CHUNK_SIZE = 32768
-    with open(destination, "wb") as f:
-        for chunk in response.iter_content(CHUNK_SIZE):
-            if chunk:  # filter out keep-alive new chunks
-                f.write(chunk)
-
-# Replace 'YOUR_FILE_ID' with your actual file ID from Google Drive
-file_id = '1WJ33nys02XpPDsMO5uIZFiLqTuAT_iuV'
-destination = 'ema_ckpt_cond.pt'
-download_file_from_google_drive(file_id, destination)
 
 # Preprocessing
 from modules import PaletteModelV2
@@ -44,7 +14,7 @@ from diffusion import Diffusion_cond
 device = 'cuda'
 
 model = PaletteModelV2(c_in=2, c_out=1, num_classes=5, image_size=256, true_img_size=64).to(device)
-ckpt = torch.load(destination, map_location=device)
+ckpt = torch.load('ema_ckpt_cond.pt')
 model.load_state_dict(ckpt)
 
 diffusion = Diffusion_cond(noise_steps=1000, img_size=256, device=device)