Walid-Ahmed commited on
Commit
c5bfeea
·
verified ·
1 Parent(s): d7c7443

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +97 -0
app.py ADDED
@@ -0,0 +1,97 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import DiffusionPipeline
4
+
5
+ # Path for the LoRA model
6
+ color_book_lora_path = "artificialguybr/ColoringBookRedmond-V2"
7
+ color_book_trigger = ", ColoringBookAF, Coloring Book"
8
+
9
+ # Predefined styles for users to choose from
10
+ styles = {
11
+ "Neonpunk": {
12
+ "prompt": "neonpunk style, cyberpunk, vaporwave, neon, vibrant, stunningly beautiful, crisp, "
13
+ "detailed, sleek, ultramodern, magenta highlights, dark purple shadows, high contrast, cinematic",
14
+ "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured"
15
+ },
16
+ "Retro Cyberpunk": {
17
+ "prompt": "retro cyberpunk, 80's inspired, synthwave, neon, vibrant, detailed, retro futurism",
18
+ "negative_prompt": "modern, desaturated, black and white, realism, low contrast"
19
+ },
20
+ "Dark Fantasy": {
21
+ "prompt": "Dark Fantasy Art, dark, moody, dark fantasy style",
22
+ "negative_prompt": "ugly, deformed, noisy, blurry, low contrast, bright, sunny"
23
+ },
24
+ "Double Exposure": {
25
+ "prompt": "Double Exposure Style, double image ghost effect, image combination, double exposure style",
26
+ "negative_prompt": "ugly, deformed, noisy, blurry, low contrast"
27
+ },
28
+ "None": {
29
+ "prompt": "8K",
30
+ "negative_prompt": "painting, drawing, illustration, glitch, deformed, mutated, cross-eyed, ugly, disfigured"
31
+ }
32
+ }
33
+
34
+ # Cache the pipeline to avoid reloading it multiple times
35
+ @gr.memo
36
+ def load_pipeline(use_lora: bool):
37
+ """Load the diffusion pipeline with or without LoRA weights."""
38
+ device = "cuda" if torch.cuda.is_available() else "cpu"
39
+
40
+ # Load the base model
41
+ pipe = DiffusionPipeline.from_pretrained(
42
+ "stabilityai/stable-diffusion-xl-refiner-1.0",
43
+ torch_dtype=torch.float16 if device == "cuda" else torch.float32,
44
+ use_safetensors=True,
45
+ variant="fp16" if device == "cuda" else None
46
+ )
47
+
48
+ # Apply LoRA weights if selected
49
+ if use_lora:
50
+ pipe.load_lora_weights(color_book_lora_path)
51
+
52
+ return pipe
53
+
54
+ def generate_image(prompt: str, style_name: str, use_lora: bool):
55
+ """Generates an AI-generated image based on user input and selected style."""
56
+ try:
57
+ # Load the pre-trained pipeline
58
+ pipeline = load_pipeline(use_lora)
59
+
60
+ # Retrieve style prompts
61
+ style_prompt = styles.get(style_name, {}).get("prompt", "")
62
+ negative_prompt = styles.get(style_name, {}).get("negative_prompt", "")
63
+
64
+ # Apply LoRA trigger phrase if selected
65
+ if use_lora:
66
+ prompt += color_book_trigger
67
+
68
+ # Generate the image
69
+ image = pipeline(
70
+ prompt=prompt + " " + style_prompt,
71
+ negative_prompt="blurred, ugly, watermark, low resolution, " + negative_prompt,
72
+ num_inference_steps=20,
73
+ guidance_scale=9.0
74
+ ).images[0]
75
+
76
+ return image
77
+
78
+ except Exception as e:
79
+ return f"Error generating image: {str(e)}"
80
+
81
+ # Gradio UI
82
+ interface = gr.Interface(
83
+ fn=generate_image,
84
+ inputs=[
85
+ gr.Textbox(label="Enter Your Prompt", placeholder="A cute lion"),
86
+ gr.Dropdown(label="Select a Style", choices=list(styles.keys()), value="None"),
87
+ gr.Checkbox(label="Use Coloring Book LoRA", value=False)
88
+ ],
89
+ outputs=gr.Image(label="Generated Image"),
90
+ title="🎨 AI Coloring Book & Style Generator",
91
+ description="Generate amazing AI-powered art using Stable Diffusion XL. "
92
+ "Choose a style or apply a Coloring Book LoRA for unique results."
93
+ )
94
+
95
+ # Run the Gradio interface
96
+ if __name__ == "__main__":
97
+ interface.launch()