Update app.py

app.py

@@ -1,3 +1,47 @@
+import functools as ft
 import gradio as gr
 
-gr.load("models/roborovski/superprompt-v1").launch()
+import torch
+import transformers
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+
+tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(
+    "models/roborovski/superprompt-v1", local_files_only=True
+)
+model: T5ForConditionalGeneration = T5ForConditionalGeneration.from_pretrained(
+    "models/roborovski/superprompt-v1", local_files_only=True
+)
+
+@ft.lru_cache(maxsize=1024)
+def super_prompt(text: str, seed: int, max_new_tokens: int, prompt: str) -> str:
+    transformers.set_seed(seed)
+
+    if max_new_tokens <= 0:
+        max_new_tokens = 150
+
+    with torch.inference_mode():
+        if prompt:
+            input_text = f"{prompt} {text}"
+        else:
+            input_text = f"Expand the following prompt to add more detail: {text}"
+
+        input_ids = tokenizer(input_text, return_tensors="pt").input_ids
+
+        outputs = model.generate(
+            input_ids,
+            max_length=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+        )
+
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+demo = gr.Interface(
+    fn=super_prompt,
+    inputs=["text", "slider", "slider", "text"],
+    outputs=["text"],
+)
+
+demo.launch()