saniaE commited on
Commit
b172d7f
·
1 Parent(s): 1cca07b

added debugging steps

Browse files
Files changed (1) hide show
  1. app.py +10 -3
app.py CHANGED
@@ -7,7 +7,7 @@ import torchvision.utils as vutils
7
  from fastapi import FastAPI, Response, HTTPException, Query
8
  from fastapi.responses import StreamingResponse
9
  from fastapi.middleware.cors import CORSMiddleware
10
- from huggingface_hub import hf_hub_download
11
 
12
  from models import Generator
13
 
@@ -35,14 +35,21 @@ def load_model():
35
  global gen_model
36
  try:
37
  token = os.getenv("HF_TOKEN")
 
 
 
 
 
 
38
  model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, token=token)
39
-
 
40
  checkpoint = torch.load(model_path, map_location=DEVICE)
41
 
42
  gen_model = Generator(z_dim=Z_DIM).to(DEVICE)
43
  gen_model.load_state_dict(checkpoint["gen_state_dict"])
44
  gen_model.eval()
45
- print("Model loaded successfully.")
46
  except Exception as e:
47
  print(f"Error loading model: {e}")
48
 
 
7
  from fastapi import FastAPI, Response, HTTPException, Query
8
  from fastapi.responses import StreamingResponse
9
  from fastapi.middleware.cors import CORSMiddleware
10
+ from huggingface_hub import hf_hub_download, login
11
 
12
  from models import Generator
13
 
 
35
  global gen_model
36
  try:
37
  token = os.getenv("HF_TOKEN")
38
+ if token:
39
+ login(token=token)
40
+ print("Login successful.")
41
+ else:
42
+ print("No HF_TOKEN found - attempting public download.")
43
+
44
  model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME, token=token)
45
+ print(f"File downloaded to: {model_path}")
46
+
47
  checkpoint = torch.load(model_path, map_location=DEVICE)
48
 
49
  gen_model = Generator(z_dim=Z_DIM).to(DEVICE)
50
  gen_model.load_state_dict(checkpoint["gen_state_dict"])
51
  gen_model.eval()
52
+ print(" SUCCESS: Petrol Pump GAN is live!")
53
  except Exception as e:
54
  print(f"Error loading model: {e}")
55