import torch
import streamlit as st
from PIL import Image
from megatron import MegatronLM, download_pretrained_model
 
#@st.cache(allow_output_mutation=True) # this line is important! without it, the app will crash on first run after a restart. 
def load_model():
    model = MegatronLM.from_pretrained('google/megatron-lm-1b')    
    return model.cuda()

 

 

 

     #@st.cache(allow_output_mutation=True) # this line is important! without it, the app will crash on first run after a restart. 

     def load_image(filename):        

          image = Image.open(filename).convert("RGB")        

          return image