Spaces:
Runtime error
Runtime error
import streamlit as st | |
import torch | |
torch.cuda.empty_cache() | |
st.set_page_config(page_title="AI Image Generator", page_icon="🖼️", layout="centered", initial_sidebar_state="auto", menu_items=None) | |
st.title("AI Image Generator") | |
st.write("Enter a phrase below to have the AI generate an image depicting your phrase") | |
prompt = st.text_input("Enter phrase here",value="", max_chars=100) | |
if prompt: | |
with st.spinner("Generation in progress..."): | |
import torch | |
from torch import autocast | |
from diffusers import StableDiffusionPipeline | |
model_id = "CompVis/stable-diffusion-v1-4" | |
device = "cuda" | |
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=True) | |
pipe = pipe.to(device) | |
with autocast("cuda"): | |
image = pipe(prompt, guidance_scale=7.5).images[0] | |
image.save("generated_image.png") | |
st.image('generated_image.png') | |
st.download_button('Download file', 'generated_image.png') |