TheEpic commited on
Commit
15e6751
1 Parent(s): 2913160

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -0
app.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import torch
3
+ torch.cuda.empty_cache()
4
+
5
+ st.set_page_config(page_title="AI Image Generator", page_icon="🖼️", layout="centered", initial_sidebar_state="auto", menu_items=None)
6
+ st.title("AI Image Generator")
7
+ st.write("Enter a phrase below to have the AI generate an image depicting your phrase")
8
+ prompt = st.text_input("Enter phrase here",value="", max_chars=100)
9
+
10
+ if prompt:
11
+ with st.spinner("Generation in progress..."):
12
+ import torch
13
+ from torch import autocast
14
+ from diffusers import StableDiffusionPipeline
15
+ model_id = "CompVis/stable-diffusion-v1-4"
16
+ device = "cuda"
17
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", use_auth_token=True)
18
+ pipe = pipe.to(device)
19
+ with autocast("cuda"):
20
+ image = pipe(prompt, guidance_scale=7.5).images[0]
21
+ image.save("generated_image.png")
22
+ st.image('generated_image.png')
23
+ st.download_button('Download file', 'generated_image.png')