prompt_extend.py

from transformers import pipeline

pipe = pipeline('text-generation', model='Gustavosta/MagicPrompt-Stable-Diffusion')

def extend_prompt(prompt):
    return pipe(prompt+',', num_return_sequences=1)[0]["generated_text"]