app.py

import transformers
import streamlit as st

from transformers import AutoTokenizer, AutoModelWithLMHead
from transformers import pipeline

#tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")

@st.cache
def load_model(model_name):
    model = AutoModelWithLMHead.from_pretrained(model_name)
    return model
    
def load_text_gen_model():
    generator = pipeline("text-generation", model="gpt2-medium")
    return generator 
    
#model = load_model("gpt2-medium")

text_generator = load_text_gen_model()

action = st.sidebar.selectbox("Pick an Action", ["Generate an Article","Create an Image"])

if action == "Generate an Article":
    prompt = st.text_input("Enter a prompt")
    if prompt:
        res = text_generator( prompt, max_length=100, temperature=0.7)
        st.write(res)