ageraustine commited on
Commit
cd01352
·
1 Parent(s): 9067e27
Files changed (2) hide show
  1. app.py +53 -0
  2. requirements.txt +3 -0
app.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import torch
3
+ from diffusers import DiffusionPipeline
4
+
5
+ # Load both base & refiner
6
+ base = DiffusionPipeline.from_pretrained(
7
+ "ageraustine/stable-diffusion", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
8
+ )
9
+ base.to("cuda")
10
+ refiner = DiffusionPipeline.from_pretrained(
11
+ "ageraustine/stable-diffusion",
12
+ text_encoder_2=base.text_encoder_2,
13
+ vae=base.vae,
14
+ torch_dtype=torch.float16,
15
+ use_safetensors=True,
16
+ variant="fp16",
17
+ )
18
+ refiner.to("cuda")
19
+
20
+ # Define how many steps and what % of steps to be run on each expert (80/20) here
21
+ n_steps = 40
22
+ high_noise_frac = 0.8
23
+
24
+ # Streamlit app
25
+ st.title("Text-to-Image Generation App")
26
+
27
+ # Text input
28
+ user_text = st.text_input("Enter a text prompt for image generation")
29
+
30
+ # Generate image based on user input
31
+ if st.button("Generate Image"):
32
+ if user_text:
33
+ # Run both experts
34
+ image = base(
35
+ prompt=user_text,
36
+ num_inference_steps=n_steps,
37
+ denoising_end=high_noise_frac,
38
+ output_type="latent",
39
+ ).images
40
+ image = refiner(
41
+ prompt=user_text,
42
+ num_inference_steps=n_steps,
43
+ denoising_start=high_noise_frac,
44
+ image=image,
45
+ ).images[0]
46
+
47
+ # Display the generated image
48
+ st.image(image, caption="Generated Image", use_column_width=True)
49
+ else:
50
+ st.warning("Please enter a text prompt for image generation.")
51
+
52
+
53
+
requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ torch
2
+ diffusers
3
+ streamlit