Update app.py

app.py

@@ -11,26 +11,25 @@ else:
     device = "cpu"
     print("Using CPU")
 
-tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
-    
+tokenizer = T5Tokenizer.from_pretrained("roborovski/superprompt-v1")
+model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1")
+model.to(device)
+
 @spaces.GPU()
 def generate(your_prompt, max_new_tokens, repetition_penalty, temperature, model_precision_type, top_p, top_k, seed):
+    if seed == 0:
+        seed = random.randint(1, 2**32-1)
+    transformers.set_seed(seed)
+
     if model_precision_type == "fp16":
         dtype = torch.float16
     elif model_precision_type == "fp32":
         dtype = torch.float32
 
-    model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", torch_dtype=dtype)
-    model.to(device)
+    model.to(dtype)
     
     input_text = f"Expand the following prompt to add more detail: {your_prompt}"
-    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
-
-    if seed == 0:
-        seed = random.randint(1, 100000)
-        torch.manual_seed(seed)
-    else:
-        torch.manual_seed(seed)
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device).to(dtype)
         
     outputs = model.generate(
         input_ids,
@@ -42,8 +41,7 @@ def generate(your_prompt, max_new_tokens, repetition_penalty, temperature, model
         top_k=top_k,
     )
 
-    better_prompt = tokenizer.decode(outputs[0])
-    better_prompt = better_prompt.replace("<pad>", "").replace("</s>", "")
+    better_prompt = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return better_prompt