Nekochu commited on
Commit
1620b49
·
verified ·
1 Parent(s): ac9fb33

Fix: unsqueeze prompt embeddings

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -66,8 +66,8 @@ def _load_all():
66
 
67
  print("[init] Loading empty prompt cache...")
68
  data = np.load(EMPTY_PROMPT_CACHE, allow_pickle=True)
69
- embeds = torch.from_numpy(data["embeds"])
70
- masks = torch.from_numpy(data["masks"])
71
 
72
  print("[init] Ready.")
73
  return dit_session, vae, embeds, masks
 
66
 
67
  print("[init] Loading empty prompt cache...")
68
  data = np.load(EMPTY_PROMPT_CACHE, allow_pickle=True)
69
+ embeds = torch.from_numpy(data["embeds"]).unsqueeze(0)
70
+ masks = torch.from_numpy(data["masks"]).unsqueeze(0)
71
 
72
  print("[init] Ready.")
73
  return dit_session, vae, embeds, masks