hermanda commited on
Commit
3b3ffae
·
verified ·
1 Parent(s): d92e253

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +104 -0
app.py ADDED
@@ -0,0 +1,104 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import gradio as gr
3
+ from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
4
+ from huggingface_hub import hf_hub_download
5
+ from safetensors.torch import load_file
6
+
7
+ # Initialize model and pipeline once at startup
8
+ base = "stabilityai/stable-diffusion-xl-base-1.0"
9
+ repo = "ByteDance/SDXL-Lightning"
10
+ ckpt = "sdxl_lightning_4step_unet.safetensors"
11
+
12
+ # Load model with float32 precision for CPU compatibility
13
+ unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float32)
14
+ unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cpu"))
15
+
16
+ # Create pipeline with CPU configuration
17
+ pipe = StableDiffusionXLPipeline.from_pretrained(
18
+ base,
19
+ unet=unet,
20
+ torch_dtype=torch.float32
21
+ ).to("cpu")
22
+
23
+ # Configure scheduler
24
+ pipe.scheduler = EulerDiscreteScheduler.from_config(
25
+ pipe.scheduler.config,
26
+ timestep_spacing="trailing"
27
+ )
28
+
29
+ # Expanded list of predefined elements
30
+ elements_list = [
31
+ "Kittens", "Tea", "Home", "Snow", "Young Girl", "Stars",
32
+ "Blanket", "Books", "Candles", "Flowers", "Moon", "Cookies",
33
+ "Fireplace", "Pillows", "Mittens", "Lanterns", "Socks",
34
+ "Hot Chocolate", "Snowflakes", "Winter Scarf", "Marshmallows",
35
+ "Vintage Clock", "Knitted Sweater", "Fairy Lights", "Porcelain Cup"
36
+ ]
37
+
38
+ def generate_image(custom_text, elements, steps):
39
+ """Generate image using the provided text, selected elements, and steps"""
40
+ # Construct the prompt
41
+ prompt_parts = []
42
+ if custom_text.strip():
43
+ prompt_parts.append(custom_text)
44
+ if elements:
45
+ prompt_parts.append(", ".join(elements))
46
+
47
+ prompt = ", ".join(prompt_parts) or "a beautiful image"
48
+
49
+ image = pipe(
50
+ prompt,
51
+ num_inference_steps=int(steps),
52
+ guidance_scale=0,
53
+ width=768,
54
+ height=960
55
+ ).images[0]
56
+
57
+ return image
58
+
59
+ # Create Gradio interface
60
+ with gr.Blocks(title="Good Night Image Diffuser") as demo:
61
+ gr.Markdown("# 🌙 Generate Good Night Wish Images")
62
+ gr.Markdown("Create personalized good night images with your message and favorite elements!")
63
+
64
+ with gr.Row():
65
+ with gr.Column(scale=1):
66
+ custom_text = gr.Textbox(
67
+ label="Your Message",
68
+ value="Create a cozy and heartwarming scene. Use a warm, pastel color palette with soft shadows and subtle textures to evoke comfort and nostalgia. Additional elements to include:",
69
+ max_lines=3
70
+ )
71
+ elements = gr.CheckboxGroup(
72
+ label="Image Elements",
73
+ choices=elements_list,
74
+ value=["Kittens", "Moon"],
75
+ info="Select elements to include in your image"
76
+ )
77
+ steps_slider = gr.Slider(
78
+ label="Number of Inference Steps",
79
+ minimum=1,
80
+ maximum=8,
81
+ value=4,
82
+ step=2,
83
+ info="Adjust the number of denoising steps (more steps can improve quality but take longer)"
84
+ )
85
+ generate_btn = gr.Button("✨ Generate Image", variant="primary")
86
+
87
+ with gr.Column(scale=1):
88
+ output_image = gr.Image(
89
+ label="Generated Image",
90
+ width=768,
91
+ height=960,
92
+ elem_id="output-image"
93
+ )
94
+
95
+ # Connect components
96
+ generate_btn.click(
97
+ fn=generate_image,
98
+ inputs=[custom_text, elements, steps_slider],
99
+ outputs=output_image,
100
+ api_name="generate"
101
+ )
102
+
103
+ if __name__ == "__main__":
104
+ demo.launch()