Update app.py

app.py

@@ -144,8 +144,10 @@ def embed_style(prompt, style_embed, style_seed):
     replacement_token_embedding = replacement_token_embedding.unsqueeze(0)  # Make it [1, 768] if necessary
 
     # Insert this into the token embeddings
-    token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)
-
+   # token_embeddings[0, torch.where(input_ids[0]==6829)] = replacement_token_embedding.to(torch_device)
+    indices = torch.where(input_ids[0] == 6829)[0]
+    for index in indices:
+        token_embeddings[0, index] = replacement_token_embedding.to(torch_device)
     # Combine with pos embs
     input_embeddings = token_embeddings + position_embeddings