Update app.py

app.py

@@ -11,13 +11,11 @@ else:
     print("Using CPU")
 
 tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
-    
 
 def generate(prompt, model_precision_type, max_new_tokens, repetition_penalty, temperature, top_p, top_k, seed):
-
     model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", device_map="auto", torch_dtype=model_precision_type)
     model.to(device)
-    
+
     input_text = f"Expand the following prompt to add more detail: {prompt}"
     input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
 
@@ -26,7 +24,7 @@ def generate(prompt, model_precision_type, max_new_tokens, repetition_penalty, t
         torch.manual_seed(seed)
     else:
         torch.manual_seed(seed)
-        
+
     outputs = model.generate(
         input_ids,
         max_new_tokens=max_new_tokens,
@@ -38,18 +36,17 @@ def generate(prompt, model_precision_type, max_new_tokens, repetition_penalty, t
     )
 
     better_prompt = tokenizer.decode(outputs[0])
-    better_prompt = better_prompt.replace("<pad>", "").replace("</s>", "")
+    better_prompt = better_prompt.replace("<pad>", "").replace("<|endoftext|>", "")
     return better_prompt
 
-
 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"),
+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")
-    
+
 temperature = gr.Slider(value=0.5, minimum=0, maximum=1, step=0.05, interactive=True, label="Temperature", info="Higher values produce more diverse outputs")
 
 top_p = gr.Slider(value=1, minimum=0, maximum=2, step=0.05, interactive=True, label="Top P", info="Higher values sample more low-probability tokens")
@@ -59,15 +56,7 @@ top_k = gr.Slider(value=1, minimum=1, maximum=100, step=1, interactive=True, lab
 seed = gr.Number(value=42, interactive=True, label="Seed", info="A starting point to initiate the generation process, put 0 for a random one")
 
 examples = [
-    [
-        "A storefront with 'Text to Image' written on it.",
-        512,
-        1.2,
-        0.5,
-        1,
-        50,
-        42,
-    ]
+    ["A storefront with 'Text to Image' written on it.", "fp16", 512, 1.2, 0.5, 1, 50, 42]
 ]
 
 gr.Interface(