Update app.py

app.py

@@ -118,7 +118,11 @@ def infer_30(
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids
+    max_length = 77
+    if input_ids.shape[1] > max_length:
+        input_ids = input_ids[:, :max_length]
+    input_ids = input_ids.to(device)
     print('-- generating image --')
     sd_image = pipe(
             prompt_embeds = pipe.text_encoder(input_ids)[0], #ensure that the input_ids are on the correct device.
@@ -170,7 +174,11 @@ def infer_60(
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids
+    max_length = 77
+    if input_ids.shape[1] > max_length:
+        input_ids = input_ids[:, :max_length]
+    input_ids = input_ids.to(device)
     print('-- generating image --')
     sd_image = pipe(
             prompt_embeds = pipe.text_encoder(input_ids)[0], #ensure that the input_ids are on the correct device.
@@ -222,7 +230,11 @@ def infer_90(
     torch.set_float32_matmul_precision("highest")
     seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device='cuda').manual_seed(seed)
-    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+    input_ids = pipe.tokenizer(prompt, return_tensors="pt").input_ids
+    max_length = 77
+    if input_ids.shape[1] > max_length:
+        input_ids = input_ids[:, :max_length]
+    input_ids = input_ids.to(device)
     print('-- generating image --')
     sd_image = pipe(
             prompt_embeds = pipe.text_encoder(input_ids)[0], #ensure that the input_ids are on the correct device.