DSnake333 commited on
Commit
513d117
·
verified ·
1 Parent(s): 1f57593

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, AutoencoderKL
3
+ import torch
4
+
5
+ # Load the model and scheduler
6
+ model_id = "sam749/Photon-v1"
7
+ vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse", torch_dtype=torch.float16).to("cuda")
8
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
9
+ pipe.vae = vae
10
+ pipe.to("cuda")
11
+
12
+ pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
13
+
14
+ # Streamlit interface
15
+ st.title("Text-to-Image Generator")
16
+
17
+ # Prompt inputs
18
+ prompt = st.text_input("Enter the prompt for the image:")
19
+ negative_prompt = "cartoon, painting, illustration, (worst quality, low quality, normal quality:2)"
20
+
21
+ cfg_scale = 5.5
22
+ width = 512
23
+ height = 768
24
+
25
+ num_inference_steps = 24
26
+
27
+ if st.button("Generate Image"):
28
+ with st.spinner("Generating..."):
29
+ # Generate image with the given parameters
30
+ generator = torch.Generator("cuda").manual_seed(0)
31
+ image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps,
32
+ guidance_scale=cfg_scale, width=width, height=height, generator=generator).images[0]
33
+
34
+ # Display the generated image
35
+ st.image(image, caption="Generated Image", use_column_width=True)
36
+ image.save("generated_image.png")