import streamlit as st
import numpy as np
import torch
from PIL import Image
from torchvision.utils import save_image
from model import load_model, sample

# Load the pre-trained Stable Diffusion model
model = load_model()

# Function to generate hugging face images
def generate_hugging_face_image():
    z = torch.randn(1, 256, 1, 1).cuda()
    x = sample(model, z, clip_denoised=True)
    save_image(x, 'output.png')
    return 'output.png'

# Streamlit app
st.title('Hugging Face Image Generator')

if st.button('Generate Hugging Face Image'):
    image_path = generate_hugging_face_image()
    image = Image.open(image_path)
    st.image(image, caption='Hugging Face Image', use_column_width=True)