basharat8763 commited on
Commit
6f6b160
·
verified ·
1 Parent(s): 72bee75

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +51 -0
app.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from diffusers import DiffusionPipeline
2
+ import torch
3
+ import streamlit as st
4
+
5
+ sdxl_base_model_path = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
6
+ sdxl_refiner_model_path = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0")
7
+
8
+ @st.cache_resource
9
+ def load_pipeline():
10
+ device = "cuda" if torch.cuda.is_available() else "cpu"
11
+ try:
12
+ # Load the base pipeline
13
+ pipe = DiffusionPipeline.from_pretrained(
14
+ sdxl_base_model_path,
15
+ torch_dtype=torch.float16 if device == "cuda" else torch.float32
16
+ )
17
+ pipe.to(device)
18
+ return pipe
19
+ except ValueError as e:
20
+ st.error(f"Error loading the pipeline: {str(e)}")
21
+ return None
22
+
23
+
24
+ def image_generation(pipe, prompt):
25
+ if not pipe:
26
+ return None
27
+ try:
28
+ # Generate the image
29
+ image = pipe(
30
+ prompt=prompt,
31
+ negative_prompt="blurred, ugly, watermark, low resolution",
32
+ num_inference_steps=40,
33
+ guidance_scale=9.0
34
+ ).images[0]
35
+ return image
36
+ except Exception as e:
37
+ st.error(f"Error generating image: {str(e)}")
38
+ return None
39
+
40
+
41
+ # Streamlit app interface
42
+ st.title("Project 11: Image Generation using SD XL")
43
+ prompt = st.text_input("Enter your prompt", value="A futuristic superhero cat")
44
+
45
+ pipeline = load_pipeline()
46
+
47
+ if pipeline and st.button("Generate Image"):
48
+ with st.spinner("Generating your Image..."):
49
+ image = image_generation(pipeline, prompt)
50
+ if image:
51
+ st.image(image)