Add CUDA availability check for input IDs and attention mask in prompt encoding

app.py

@@ -51,8 +51,8 @@ def encode_prompt(text_tokenizer, text_encoder, prompt):
     print(f'prompt={prompt}')
     captions = [prompt]
     tokens = text_tokenizer(text=captions, max_length=512, padding='max_length', truncation=True, return_tensors='pt')  # todo: put this into dataset
-    input_ids = tokens.input_ids.cuda(non_blocking=True)
-    mask = tokens.attention_mask.cuda(non_blocking=True)
+    input_ids = tokens.input_ids.cuda(non_blocking=True) if torch.cuda.is_available() else tokens.input_ids
+    mask = tokens.attention_mask.cuda(non_blocking=True) if torch.cuda.is_available() else tokens.attention_mask
     text_features = text_encoder(input_ids=input_ids, attention_mask=mask)['last_hidden_state'].float()
     lens: List[int] = mask.sum(dim=-1).tolist()
     cu_seqlens_k = F.pad(mask.sum(dim=-1).to(dtype=torch.int32).cumsum_(0), (1, 0))