Taizun commited on
Commit
da0a826
·
verified ·
1 Parent(s): ef12d1e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +96 -30
app.py CHANGED
@@ -1,51 +1,117 @@
1
- import streamlit as st
2
- from diffusers import StableDiffusionPipeline
3
  import torch
 
 
4
 
 
5
  @st.cache_resource
6
- def load_lightweight_pipeline():
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
- pipeline = StableDiffusionPipeline.from_pretrained(
9
- "CompVis/stable-diffusion-v1-4",
10
- torch_dtype=torch.float16 if device == "cuda" else torch.float32
11
- ).to(device)
12
- return pipeline
13
 
14
- def generate_image(pipeline, prompt, negative_prompt):
 
15
  try:
16
- image = pipeline(
17
  prompt=prompt,
18
- negative_prompt=negative_prompt,
19
- num_inference_steps=25,
20
- guidance_scale=7.5
21
  ).images[0]
22
  return image
23
  except Exception as e:
24
  st.error(f"Error generating image: {str(e)}")
25
  return None
26
 
27
- st.title("Image Generator")
28
- st.markdown('<small>By Taizun</small>', unsafe_allow_html=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
30
- prompt = st.text_input("Enter your prompt")
31
- pipeline = load_lightweight_pipeline()
32
 
33
- styles_dict = {
34
- "Neon Vibes": {"prompt": "neon lights in a futuristic city", "negative_prompt": "blurred, dull"},
35
- "Retro Sci-Fi": {"prompt": "retro futuristic landscape", "negative_prompt": "realistic, modern"},
36
- "Mystic Forest": {"prompt": "dark mystical forest with glowing lights", "negative_prompt": "sunny, bright"},
37
- "Abstract Art": {"prompt": "abstract colorful shapes", "negative_prompt": "realistic, detailed"}
38
- }
39
 
 
40
  style_name = st.selectbox("Select a Style", options=list(styles_dict.keys()))
41
 
 
42
  if style_name:
43
- selected_style = styles_dict[style_name]
44
- style_prompt = selected_style["prompt"]
45
- negative_prompt = selected_style["negative_prompt"]
46
 
47
- if st.button("Generate Image"):
48
- with st.spinner("Generating your image..."):
49
- image = generate_image(pipeline, prompt + " " + style_prompt, negative_prompt)
 
50
  if image:
51
- st.image(image, caption="Generated Image", use_column_width=True)
 
 
 
1
  import torch
2
+ import streamlit as st
3
+ from diffusers import DiffusionPipeline
4
 
5
+ # Cache resource to load the pipeline
6
  @st.cache_resource
7
+ def load_pipeline():
8
  device = "cuda" if torch.cuda.is_available() else "cpu"
9
+ pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0",
10
+ torch_dtype=torch.float16 if device == "cuda" else torch.float32,
11
+ use_safetensors=True,
12
+ variant="fp16" if device == "cuda" else None)
13
+ return pipe
14
 
15
+ # Image generation function
16
+ def image_generation(pipe, prompt, negative_prompt):
17
  try:
18
+ image = pipe(
19
  prompt=prompt,
20
+ negative_prompt="blurred, ugly, watermark, low resolution" + negative_prompt,
21
+ num_inference_steps=20,
22
+ guidance_scale=9.0
23
  ).images[0]
24
  return image
25
  except Exception as e:
26
  st.error(f"Error generating image: {str(e)}")
27
  return None
28
 
29
+ # Define the table as a list of dictionaries with more artistic styles
30
+ table = [
31
+ {
32
+ "name": "sai-neonpunk",
33
+ "prompt": "neonpunk style . cyberpunk, vaporwave, neon, vibes, vibrant, stunningly beautiful, crisp, detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic, ultra detailed, intricate, professional",
34
+ "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured"
35
+ },
36
+ {
37
+ "name": "futuristic-retro cyberpunk",
38
+ "prompt": "retro cyberpunk. 80's inspired, synthwave, neon, vibrant, detailed, retro futurism",
39
+ "negative_prompt": "modern, desaturated, black and white, realism, low contrast"
40
+ },
41
+ {
42
+ "name": "Dark Fantasy",
43
+ "prompt": "Dark Fantasy Art, dark, moody, dark fantasy style",
44
+ "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, bright, sunny"
45
+ },
46
+ {
47
+ "name": "Double Exposure",
48
+ "prompt": "Double Exposure Style, double image ghost effect, image combination, double exposure style",
49
+ "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
50
+ },
51
+ {
52
+ "name": "Vintage Retro",
53
+ "prompt": "vintage retro style, old-school, warm colors, nostalgic, retro textures, faded colors, cozy atmosphere, worn, classic film",
54
+ "negative_prompt": "modern, bright, high definition, sharp"
55
+ },
56
+ {
57
+ "name": "Sci-Fi Futuristic",
58
+ "prompt": "Sci-Fi futuristic, cybernetic, sleek metallic designs, neon, artificial intelligence, high-tech, future cities, robots, alien landscapes",
59
+ "negative_prompt": "low-tech, blurry, old-fashioned, dull, dark"
60
+ },
61
+ {
62
+ "name": "Surrealism",
63
+ "prompt": "Surrealism, dreamlike, abstract, bizarre, imaginative, fantastical, magical realism, melting objects, floating landscapes",
64
+ "negative_prompt": "realistic, plain, dull, conventional"
65
+ },
66
+ {
67
+ "name": "Pop Art",
68
+ "prompt": "Pop Art, bold colors, comic book style, strong lines, graphic, playful, contemporary, inspired by Andy Warhol and Roy Lichtenstein",
69
+ "negative_prompt": "minimalistic, muted colors, traditional, abstract"
70
+ },
71
+ {
72
+ "name": "Gothic Horror",
73
+ "prompt": "Gothic horror, dark, eerie, vintage horror, Victorian architecture, spooky ambiance, haunted, gothic castles, fog, grim landscapes",
74
+ "negative_prompt": "light, cheerful, happy, bright"
75
+ },
76
+ {
77
+ "name": "Minimalism",
78
+ "prompt": "Minimalism, clean lines, simple forms, neutral tones, white space, simple color palette, elegant, modern design",
79
+ "negative_prompt": "cluttered, bright, complex, busy"
80
+ },
81
+ {
82
+ "name": "Steampunk",
83
+ "prompt": "Steampunk, Victorian industrial, brass gears, steam-powered machinery, fantasy airships, retro-futuristic, adventure, mechanical details",
84
+ "negative_prompt": "futuristic, modern, minimalistic, digital"
85
+ }
86
+ ]
87
+
88
+ # Convert to dictionary for easy lookup
89
+ styles_dict = {entry["name"]: entry for entry in table}
90
+
91
+ # Streamlit UI
92
+ st.title("Image Generation")
93
+
94
+ # Display credit after the title in smaller font
95
+ st.markdown("<h3 style='font-size: 18px;'>Made by Taizun</h3>", unsafe_allow_html=True)
96
 
97
+ # User input for the custom prompt
98
+ prompt = st.text_input("Enter your Prompt")
99
 
100
+ # Load the pre-trained model
101
+ pipeline = load_pipeline()
 
 
 
 
102
 
103
+ # Dropdown for selecting a style
104
  style_name = st.selectbox("Select a Style", options=list(styles_dict.keys()))
105
 
106
+ # Display the selected style's prompt and negative prompt
107
  if style_name:
108
+ selected_entry = styles_dict[style_name]
109
+ selected_style_prompt = selected_entry["prompt"]
110
+ selected_style_negative_prompt = selected_entry["negative_prompt"]
111
 
112
+ # Button to trigger image generation
113
+ if st.button("Generate Awesome Image"):
114
+ with st.spinner("Generating your awesome image..."):
115
+ image = image_generation(pipeline, prompt + selected_style_prompt, selected_style_negative_prompt)
116
  if image:
117
+ st.image(image)