inplace attention state, faster and less memory

2022-07-04 09:14:37 -04:00
parent aca617dc64
commit 6f617fe98f
4 changed files with 25 additions and 15 deletions
@@ -54,7 +54,8 @@ def generate_image(

    if token_count < 256:
        image_tokens = model.generate_image_tokens(text, seed, grid_size ** 2)
-        print('image tokens', image_tokens.to('cpu').detach().numpy())
+        image_tokens = image_tokens[:, :token_count].to('cpu').detach().numpy()
+        print('image tokens', image_tokens)
    else:
        image = model.generate_image(text, seed, grid_size)
        save_image(image, image_path)