Update app.py

app.py

@@ -1,8 +1,14 @@
 import gradio as gr
+from transformers import pipeline
 
-def greet(name):
-    return "Hello " + name + "!!"
+pipe = pipeline('text-generation', model='daspartho/prompt-extend')
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+def extend_prompt(input):
+    return pipe(prompt+',', num_return_sequences=1)[0]["generated_text"]
+
+input_prompt = gr.Text(label="Enter the initial prompt")
+sd2_output = gr.Text(label="Extended prompt suitable for Stable Diffusion 2")
+
+demo = gr.Interface(fn=extend_prompt, inputs=input_prompt, outputs=sd2_output)
 demo.queue(max_size=10,concurrency_count=20)
 demo.launch(enable_queue=True)