NSFW-ToyWorld / prompt_extend.py
Yntec's picture
Move to Gradio 4.x
73d7fa1 verified
raw
history blame
203 Bytes
from transformers import pipeline
pipe = pipeline('text-generation', model='daspartho/prompt-extend')
def extend_prompt(prompt):
return pipe(prompt+',', num_return_sequences=1)[0]["generated_text"]