Update app.py

app.py

@@ -1,6 +1,10 @@
 import gradio as gr
 from transformers import pipeline
 from gtts import gTTS
+import torch
+
+# GPU control
+device = 0 if torch.cuda.is_available() else -1
 
 # gpt2-large model text generation
 def text_gen(prompt):
@@ -11,7 +15,7 @@ def text_gen(prompt):
   #res_k = generator(prompt, max_new_tokens=120, num_return_sequences=1, do_sample=True, top_k=4)
   #res_c = generator(prompt, max_new_tokens=100, num_return_sequences=1, penalty_alpha=0.6, top_k=4)
   #<span style="font-family:Papyrus; font-size:2em;">👽 0xZee</span>
-  res = generator(prompt, max_new_tokens=100, num_return_sequences=1)
+  res = generator(prompt, max_lenght=120, num_return_sequences=1, device=device)
   res = res[0]['generated_text']
   # text_to_audio
   audio = gTTS(text=res, lang='en', slow=False)
@@ -20,7 +24,7 @@ def text_gen(prompt):
 
 def complete_with_gpt(text):
   complete = pipeline("text-generation", model='openai-gpt')
-  res = complete(text[-50:], max_new_tokens=100, num_return_sequences=1)
+  res = complete(text[-50:], max_new_tokens=100, num_return_sequences=1, device=device)
   return res[0]['generated_text']