victorgeek commited on
Commit
73eea52
ยท
verified ยท
1 Parent(s): 26e6741

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +118 -0
app.py ADDED
@@ -0,0 +1,118 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import gradio as gr
3
+ from diffusers import DiffusionPipeline
4
+ import diffusers
5
+ import numpy as np
6
+ import random
7
+
8
+ # =========================================================
9
+ # MODEL CONFIGURATION
10
+ # =========================================================
11
+ MAX_SEED = np.iinfo(np.int32).max
12
+ # Turbo model แ€–แ€ผแ€…แ€บแ CPU แ€•แ€ฑแ€ซแ€บแ€แ€ฝแ€„แ€บ 1-4 steps แ€žแ€ฌ แ€žแ€ฏแ€ถแ€ธแ€›แ€”แ€บ แ€กแ€€แ€ผแ€ถแ€•แ€ผแ€ฏแ€•แ€ซแ€žแ€Šแ€บ (แ€™แ€ผแ€”แ€บแ€†แ€”แ€บแ€…แ€ฑแ€›แ€”แ€บ)
13
+ DEFAULT_STEPS = 4
14
+
15
+ # =========================================================
16
+ # LOAD PIPELINE (CPU Optimized)
17
+ # =========================================================
18
+ print("Loading Z-Image-Turbo pipeline to CPU...")
19
+
20
+ # CPU แ€•แ€ฑแ€ซแ€บแ€แ€ฝแ€„แ€บ Error แ€€แ€„แ€บแ€ธแ€…แ€ฑแ€›แ€”แ€บ float32 แ€žแ€ฏแ€ถแ€ธแ€›แ€•แ€ซแ€™แ€Šแ€บ
21
+ pipe = DiffusionPipeline.from_pretrained(
22
+ "Tongyi-MAI/Z-Image-Turbo",
23
+ torch_dtype=torch.float32,
24
+ low_cpu_mem_usage=True
25
+ )
26
+
27
+ # Memory แ€แ€ปแ€ฝแ€ฑแ€แ€ฌแ€›แ€”แ€บ (CPU แ€กแ€แ€ฝแ€€แ€บ แ€กแ€›แ€ฑแ€ธแ€€แ€ผแ€ฎแ€ธแ€•แ€ซแ€žแ€Šแ€บ)
28
+ pipe.enable_attention_slicing()
29
+ pipe.to("cpu")
30
+
31
+ # =========================================================
32
+ # PROMPT EXAMPLES (User แ€•แ€ฑแ€ธแ€‘แ€ฌแ€ธแ€žแ€ฑแ€ฌ list แ€‘แ€ฒแ€™แ€พ แ€กแ€แ€ปแ€ญแ€ฏแ€ทแ€€แ€ญแ€ฏ แ€”แ€™แ€ฐแ€”แ€ฌแ€šแ€ฐแ€‘แ€ฌแ€ธแ€žแ€Šแ€บ)
33
+ # =========================================================
34
+ prompt_examples = [
35
+ "Moody mature anime scene of two lovers kissing under neon rain, sensual atmosphere",
36
+ "A woman in a blue hanbok sits on a wooden floor, gazing out of a window.",
37
+ "A traditional Japanese onsen, with steam rising, a young woman in a colorful kimono."
38
+ ]
39
+
40
+ def get_random_prompt():
41
+ return random.choice(prompt_examples)
42
+
43
+ # =========================================================
44
+ # IMAGE GENERATOR
45
+ # =========================================================
46
+ def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed, num_images):
47
+ if not prompt:
48
+ raise gr.Error("Please enter a prompt.")
49
+
50
+ if randomize_seed:
51
+ seed = random.randint(0, MAX_SEED)
52
+
53
+ # CPU แ€•แ€ฑแ€ซแ€บแ€แ€ฝแ€„แ€บ RAM แ€™แ€•แ€ผแ€Šแ€ทแ€บแ€…แ€ฑแ€›แ€”แ€บ แ€•แ€ฏแ€ถแ€กแ€›แ€ฑแ€กแ€แ€ฝแ€€แ€บแ€€แ€ญแ€ฏ แ€€แ€”แ€ทแ€บแ€žแ€แ€บแ€แ€ผแ€„แ€บแ€ธ
54
+ num_images = min(max(1, int(num_images)), 2)
55
+
56
+ generator = torch.Generator("cpu").manual_seed(int(seed))
57
+
58
+ # CPU inference แ€–แ€ผแ€…แ€บแ แ€กแ€แ€ปแ€ญแ€”แ€บแ€€แ€ผแ€ฌแ€”แ€ญแ€ฏแ€„แ€บแ€€แ€ผแ€ฑแ€ฌแ€„แ€บแ€ธ แ€žแ€แ€ญแ€•แ€ผแ€ฏแ€•แ€ซ
59
+ result = pipe(
60
+ prompt=prompt,
61
+ height=int(height),
62
+ width=int(width),
63
+ num_inference_steps=int(num_inference_steps),
64
+ guidance_scale=0.0,
65
+ generator=generator,
66
+ max_sequence_length=512, # CPU แ€กแ€แ€ฝแ€€แ€บ length แ€œแ€ปแ€พแ€ฑแ€ฌแ€ทแ€‘แ€ฌแ€ธแ€แ€ผแ€„แ€บแ€ธแ€€ แ€•แ€ญแ€ฏแ€™แ€ผแ€”แ€บแ€…แ€ฑแ€žแ€Šแ€บ
67
+ num_images_per_prompt=num_images,
68
+ output_type="pil",
69
+ )
70
+
71
+ return result.images, seed
72
+
73
+ # ============================================
74
+ # ๐ŸŽจ UI Design (Original CSS and Layout)
75
+ # ============================================
76
+ css = """
77
+ /* User แ€•แ€ฑแ€ธแ€‘แ€ฌแ€ธแ€žแ€ฑแ€ฌ CSS แ€€แ€ญแ€ฏ แ€คแ€”แ€ฑแ€›แ€ฌแ€แ€ฝแ€„แ€บ แ€‘แ€Šแ€ทแ€บแ€žแ€ฝแ€„แ€บแ€ธแ€‘แ€ฌแ€ธแ€žแ€Šแ€บ */
78
+ @import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');
79
+ .gradio-container { background-color: #FEF9C3 !important; font-family: 'Comic Neue', cursive !important; }
80
+ .header-text h1 { font-family: 'Bangers', cursive !important; text-align: center; font-size: 3rem; }
81
+ .warning-box { background: #FEF3C7; border: 3px solid #F59E0B; padding: 10px; text-align: center; }
82
+ """
83
+
84
+ with gr.Blocks(css=css) as demo:
85
+ gr.Markdown("# ๐Ÿ–ผ๏ธ AI Image Generator (CPU Version)", elem_classes="header-text")
86
+
87
+ with gr.Row():
88
+ with gr.Column():
89
+ prompt_input = gr.Textbox(label="โœ๏ธ Prompt", lines=3)
90
+ random_button = gr.Button("๐ŸŽฒ RANDOM PROMPT")
91
+
92
+ with gr.Row():
93
+ height_input = gr.Slider(256, 1024, 512, step=64, label="Height")
94
+ width_input = gr.Slider(256, 1024, 512, step=64, label="Width")
95
+
96
+ num_images_input = gr.Slider(1, 2, 1, step=1, label="Images Count")
97
+
98
+ with gr.Accordion("โš™๏ธ Settings", open=False):
99
+ steps_slider = gr.Slider(1, 10, DEFAULT_STEPS, step=1, label="Steps (Keep low for CPU)")
100
+ seed_input = gr.Number(value=42, label="Seed")
101
+ randomize_seed_checkbox = gr.Checkbox(label="Randomize Seed", value=True)
102
+
103
+ generate_button = gr.Button("โœจ GENERATE", variant="primary")
104
+
105
+ with gr.Column():
106
+ output_gallery = gr.Gallery(label="Output", columns=1)
107
+ used_seed_output = gr.Number(label="Seed Used")
108
+
109
+ random_button.click(fn=get_random_prompt, outputs=[prompt_input])
110
+ generate_button.click(
111
+ fn=generate_image,
112
+ inputs=[prompt_input, height_input, width_input, steps_slider, seed_input, randomize_seed_checkbox, num_images_input],
113
+ outputs=[output_gallery, used_seed_output]
114
+ )
115
+
116
+ if __name__ == "__main__":
117
+ demo.launch()
118
+