jacky3102 commited on
Commit
9da6ab4
·
verified ·
1 Parent(s): 0a76a4e

Upload app.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. app.py +231 -0
app.py ADDED
@@ -0,0 +1,231 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import torch
3
+ import gradio as gr
4
+ from diffusers import StableDiffusionXLPipeline, AutoencoderKL
5
+ from diffusers.schedulers import EulerAncestralDiscreteScheduler
6
+ from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
7
+ from safetensors.torch import load_file
8
+ import gc
9
+
10
+ # ─── Config ──────────────────────────────────────────────────────────
11
+ MODEL_ID = "NoobAI-XL/NoobAI-XL-1.1-SDNQ-UINT4" # HF repo
12
+ DEVICE = "cpu"
13
+ DTYPE = torch.float32 # CPU 用 fp32 更稳
14
+ DEFAULT_W, DEFAULT_H = 640, 960
15
+ DEFAULT_STEPS = 20
16
+ DEFAULT_GUIDANCE = 5.0
17
+
18
+ # ─── Prompt enhancement (Chinese → English + quality tags) ───────────
19
+ QUALITY_TAGS = (
20
+ "masterpiece, best quality, ultra high res, "
21
+ "dynamic composition, dramatic lighting, "
22
+ "detailed background, manga style, anime style, "
23
+ "sharp focus, vivid colors, professional illustration"
24
+ )
25
+
26
+ NEGATIVE_PROMPT = (
27
+ "low quality, worst quality, bad anatomy, bad composition, "
28
+ "blurry, pixelated, watermark, text, signature, logo, "
29
+ "ugly, deformed, noise, artifacts, lowres, error, "
30
+ "missing fingers, extra digits, fewer digits, cropped, "
31
+ "jpeg artifacts, out of frame, mutation, mutated, "
32
+ "extra limbs, malformed limbs, long neck, disfigured"
33
+ )
34
+
35
+ STYLE_PROMPTS = {
36
+ "Manga / Anime": "manga style, anime style, screentone, ink lines, dynamic panel composition",
37
+ "Western Comics": "western comic style, bold outlines, halftone dots, vibrant colors, dynamic action",
38
+ "Manhwa": "manhwa style, webtoon style, clean lines, soft shading, vertical composition",
39
+ "Sketch / Lineart": "sketch style, lineart, black and white, pencil drawing, detailed strokes",
40
+ }
41
+
42
+ def enhance_prompt(user_prompt: str, style: str) -> str:
43
+ """中文场景 → 英文提示词 + 质量标签 + 风格标签"""
44
+ # 简单的中文关键词映射(可扩展为更完整的翻译/扩写)
45
+ cn_en_map = {
46
+ "恐龙": "dinosaur",
47
+ "霸王龙": "T-Rex",
48
+ "暴雨": "heavy rain",
49
+ "史前森林": "prehistoric forest",
50
+ "仰头": "looking up",
51
+ "怒吼": "roaring",
52
+ "闪电": "lightning",
53
+ "照亮": "illuminating",
54
+ "天空": "sky",
55
+ "低角度": "low angle",
56
+ "镜头": "camera angle",
57
+ "强烈": "intense",
58
+ "漫画": "manga",
59
+ "构图": "composition",
60
+ "动态": "dynamic",
61
+ "线条": "lines",
62
+ "丰富": "rich",
63
+ "背景": "background",
64
+ "细节": "details",
65
+ "愤怒": "angry",
66
+ "站在": "standing in",
67
+ "里": "in",
68
+ "的": "",
69
+ "一只": "a",
70
+ "在": "in",
71
+ "中": "in",
72
+ "和": "and",
73
+ "与": "and",
74
+ }
75
+ # 粗略翻译:把常见中文词替换成英文
76
+ translated = user_prompt
77
+ for cn, en in cn_en_map.items():
78
+ translated = translated.replace(cn, en)
79
+ # 如果全是中文,做简单拼接;如果已经有英文,保留
80
+ if any('\u4e00' <= c <= '\u9fff' for c in translated):
81
+ # 还含有中文,简单处理
82
+ prompt_en = translated
83
+ else:
84
+ prompt_en = translated
85
+ style_tag = STYLE_PROMPTS.get(style, STYLE_PROMPTS["Manga / Anime"])
86
+ return f"{prompt_en}, {style_tag}, {QUALITY_TAGS}"
87
+
88
+ # ─── Load pipeline (lazy, cached) ────────────────────────────────────
89
+ _pipe = None
90
+
91
+ def get_pipeline():
92
+ global _pipe
93
+ if _pipe is not None:
94
+ return _pipe
95
+
96
+ print("Loading NoobAI-XL SDNQ pipeline on CPU...")
97
+
98
+ # 1) VAE (单独加载,启用 slicing)
99
+ vae = AutoencoderKL.from_pretrained(
100
+ MODEL_ID,
101
+ subfolder="vae",
102
+ torch_dtype=DTYPE,
103
+ low_cpu_mem_usage=True,
104
+ )
105
+ vae.enable_slicing()
106
+ vae.enable_tiling()
107
+
108
+ # 2) Text encoders & tokenizers (SDXL 双编码器)
109
+ tokenizer_1 = CLIPTokenizer.from_pretrained(MODEL_ID, subfolder="tokenizer")
110
+ tokenizer_2 = CLIPTokenizer.from_pretrained(MODEL_ID, subfolder="tokenizer_2")
111
+ text_encoder_1 = CLIPTextModel.from_pretrained(
112
+ MODEL_ID, subfolder="text_encoder", torch_dtype=DTYPE, low_cpu_mem_usage=True
113
+ )
114
+ text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
115
+ MODEL_ID, subfolder="text_encoder_2", torch_dtype=DTYPE, low_cpu_mem_usage=True
116
+ )
117
+
118
+ # 3) Pipeline
119
+ pipe = StableDiffusionXLPipeline.from_pretrained(
120
+ MODEL_ID,
121
+ vae=vae,
122
+ text_encoder=text_encoder_1,
123
+ text_encoder_2=text_encoder_2,
124
+ tokenizer=tokenizer_1,
125
+ tokenizer_2=tokenizer_2,
126
+ scheduler=EulerAncestralDiscreteScheduler.from_pretrained(MODEL_ID, subfolder="scheduler"),
127
+ torch_dtype=DTYPE,
128
+ low_cpu_mem_usage=True,
129
+ variant="fp16" if DTYPE == torch.float16 else None,
130
+ )
131
+
132
+ # CPU 优化
133
+ pipe.enable_model_cpu_offload() # 逐层卸载到 CPU
134
+ pipe.enable_vae_slicing()
135
+ pipe.enable_vae_tiling()
136
+ pipe.enable_attention_slicing("auto") # 注意力切片
137
+ # 关闭 xformers / flash attention(CPU 不支持)
138
+ try:
139
+ pipe.disable_xformers_memory_efficient_attention()
140
+ except Exception:
141
+ pass
142
+
143
+ pipe.to(DEVICE)
144
+ _pipe = pipe
145
+ print("Pipeline ready.")
146
+ return pipe
147
+
148
+ # ─── Generation ──────────────────────────────────────────────────────
149
+ def generate(prompt, style, width, height, steps, guidance, progress=gr.Progress()):
150
+ if not prompt.strip():
151
+ return None, "请输入场景描述"
152
+
153
+ progress(0.1, desc="增强提示词...")
154
+ full_prompt = enhance_prompt(prompt, style)
155
+
156
+ progress(0.2, desc="加载模型...")
157
+ pipe = get_pipeline()
158
+
159
+ progress(0.4, desc="生成中(CPU 慢,请耐心等待)...")
160
+ with torch.inference_mode():
161
+ image = pipe(
162
+ prompt=full_prompt,
163
+ negative_prompt=NEGATIVE_PROMPT,
164
+ width=width,
165
+ height=height,
166
+ num_inference_steps=steps,
167
+ guidance_scale=guidance,
168
+ num_images_per_prompt=1,
169
+ ).images[0]
170
+
171
+ progress(0.9, desc="保存 PNG...")
172
+ # 保存到临时文件供 Gradio 返回
173
+ out_path = "/tmp/comic_scene.png"
174
+ image.save(out_path, "PNG")
175
+ progress(1.0, desc="完成")
176
+ gc.collect()
177
+ torch.cuda.empty_cache() if torch.cuda.is_available() else None
178
+ return out_path, f"✅ 生成完成\nPrompt: {full_prompt[:200]}..."
179
+
180
+ # ─── Gradio UI ───────────────────────────────────────────────────────
181
+ css = """
182
+ #title { text-align: center; font-size: 1.5rem; font-weight: 700; margin-bottom: 0.5rem; }
183
+ #subtitle { text-align: center; color: #666; margin-bottom: 1rem; }
184
+ .generate-btn { background: linear-gradient(90deg, #667eea 0%, #764ba2 100%) !important; }
185
+ """
186
+
187
+ with gr.Blocks(css=css, title="Comic Scene Generator") as demo:
188
+ gr.HTML('<div id="title">🎨 Comic Scene Generator</div>')
189
+ gr.HTML('<div id="subtitle">输入一句中文场景 → 自动扩写 → NoobAI-XL 生成漫画场景图</div>')
190
+
191
+ with gr.Row():
192
+ with gr.Column(scale=1):
193
+ prompt = gr.Textbox(
194
+ label="Scene",
195
+ placeholder="例:一只愤怒的霸王龙站在暴雨中的史前森林里仰头怒吼,闪电照亮天空,低角度镜头,强烈的漫画构图,动态线条,丰富背景细节",
196
+ lines=4,
197
+ )
198
+ style = gr.Dropdown(
199
+ label="Style",
200
+ choices=list(STYLE_PROMPTS.keys()),
201
+ value="Manga / Anime",
202
+ )
203
+ with gr.Row():
204
+ width = gr.Number(label="Width", value=DEFAULT_W, precision=0, minimum=320, maximum=1024, step=64)
205
+ height = gr.Number(label="Height", value=DEFAULT_H, precision=0, minimum=320, maximum=1536, step=64)
206
+ steps = gr.Slider(label="Steps", minimum=10, maximum=50, value=DEFAULT_STEPS, step=1)
207
+ guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=15.0, value=DEFAULT_GUIDANCE, step=0.5)
208
+ btn = gr.Button("Generate", elem_classes="generate-btn", variant="primary")
209
+
210
+ with gr.Column(scale=1):
211
+ output_img = gr.Image(label="生成的漫画场景", type="filepath", interactive=False)
212
+ status = gr.Textbox(label="Status", interactive=False)
213
+
214
+ btn.click(
215
+ generate,
216
+ inputs=[prompt, style, width, height, steps, guidance],
217
+ outputs=[output_img, status],
218
+ )
219
+
220
+ # 示例
221
+ gr.Examples(
222
+ examples=[
223
+ ["一只愤怒的霸王龙站在暴雨中的史前森林里仰头怒吼,闪电照亮天空,低角度镜头,强烈的漫画构图,动态线条,丰富背景细节", "Manga / Anime", 640, 960, 20, 5.0],
224
+ ["赛博朋克城市夜晚,霓虹灯招牌倒映在湿漉漉的街道上,孤独的少女撑着透明伞,电影感构图", "Manga / Anime", 640, 960, 20, 5.0],
225
+ ["古风仙侠,白衣剑仙御剑飞行穿越云海,日出金光洒在云层上,宏大壮观,工笔画风", "Manga / Anime", 640, 960, 25, 5.5],
226
+ ],
227
+ inputs=[prompt, style, width, height, steps, guidance],
228
+ )
229
+
230
+ if __name__ == "__main__":
231
+ demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)