Update app.py

app.py

@@ -11,8 +11,8 @@ else:
     print("Using CPU")
 
 tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
-model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", torch_dtype=torch.float16)
-
+model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", device_map="auto", torch_dtype=model_precision_type
+                                                   
 model.to(device)
     
 
@@ -51,6 +51,8 @@ def generate(
 
 prompt = gr.Textbox(label="Prompt", interactive=True)
 
+model_precision_type = gr.Dropdown(choices=[('fp16', torch.float16), ('fp32', torch.float32)], type="value", value=torch.float16, label="Model Precision", info="fp16 is faster, fp32 is more precise"),
+
 max_new_tokens = gr.Slider(value=512, minimum=250, maximum=512, step=1, interactive=True, label="Max New Tokens", info="The maximum numbers of new tokens, controls how long is the output")
     
 repetition_penalty = gr.Slider(value=1.2, minimum=0, maximum=2, step=0.05, interactive=True, label="Repetition Penalty", info="Penalize repeated tokens, making the AI repeat less itself")
@@ -77,7 +79,7 @@ examples = [
 
 gr.Interface(
     fn=generate,
-    inputs=[prompt, max_new_tokens, repetition_penalty, temperature, top_p, top_k, seed],
+    inputs=[prompt, model_precision_type, max_new_tokens, repetition_penalty, temperature, top_p, top_k, seed],
     outputs=gr.Textbox(label="Better Prompt"),
     title="SuperPrompt-v1",
     description='Make your prompts more detailed! <br> <a href="https://huggingface.co/roborovski/superprompt-v1">Model used</a> <br> <a href="https://brianfitzgerald.xyz/prompt-augmentation/">Model Blog</a> <br> Task Prefix: "Expand the following prompt to add more detail:" is already setted! <br> Hugging Face Space made by [Nick088](https://linktr.ee/Nick088)',