ginipick commited on
Commit
866ecd0
1 Parent(s): ba28a5e

Delete uhdimage.cod

Browse files
Files changed (1) hide show
  1. uhdimage.cod +0 -271
uhdimage.cod DELETED
@@ -1,271 +0,0 @@
1
- import os
2
- import yaml
3
- import torch
4
- import sys
5
- sys.path.append(os.path.abspath('./'))
6
- from inference.utils import *
7
- from train import WurstCoreB
8
- from gdf import DDPMSampler
9
- from train import WurstCore_t2i as WurstCoreC
10
- import numpy as np
11
- import random
12
- import argparse
13
- import gradio as gr
14
- import spaces
15
- from huggingface_hub import hf_hub_url
16
- import subprocess
17
- from huggingface_hub import hf_hub_download
18
- from transformers import pipeline
19
-
20
- # Initialize the translation pipeline
21
- translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
22
-
23
- def parse_args():
24
- parser = argparse.ArgumentParser()
25
- parser.add_argument('--height', type=int, default=2560, help='image height')
26
- parser.add_argument('--width', type=int, default=5120, help='image width')
27
- parser.add_argument('--seed', type=int, default=123, help='random seed')
28
- parser.add_argument('--dtype', type=str, default='bf16', help='if bf16 does not work, change it to float32')
29
- parser.add_argument('--config_c', type=str,
30
- default='configs/training/t2i.yaml', help='config file for stage c, latent generation')
31
- parser.add_argument('--config_b', type=str,
32
- default='configs/inference/stage_b_1b.yaml', help='config file for stage b, latent decoding')
33
- parser.add_argument('--prompt', type=str,
34
- default='A photo-realistic image of a west highland white terrier in the garden, high quality, detail rich, 8K', help='text prompt')
35
- parser.add_argument('--num_image', type=int, default=1, help='how many images generated')
36
- parser.add_argument('--output_dir', type=str, default='figures/output_results/', help='output directory for generated image')
37
- parser.add_argument('--stage_a_tiled', action='store_true', help='whether or not to use tiled decoding for stage a to save memory')
38
- parser.add_argument('--pretrained_path', type=str, default='models/ultrapixel_t2i.safetensors', help='pretrained path of newly added parameter of UltraPixel')
39
- args = parser.parse_args()
40
- return args
41
-
42
- def clear_image():
43
- return None
44
-
45
- def load_message(height, width, seed, prompt, args, stage_a_tiled):
46
- args.height = height
47
- args.width = width
48
- args.seed = seed
49
- args.prompt = prompt + ' rich detail, 4k, high quality'
50
- args.stage_a_tiled = stage_a_tiled
51
- return args
52
-
53
- def is_korean(text):
54
- return any('\uac00' <= char <= '\ud7a3' for char in text)
55
-
56
- def translate_if_korean(text):
57
- if is_korean(text):
58
- translated = translator(text, max_length=512)[0]['translation_text']
59
- print(f"Translated from Korean: {text} -> {translated}")
60
- return translated
61
- return text
62
-
63
- @spaces.GPU(duration=120)
64
- def get_image(height, width, seed, prompt, cfg, timesteps, stage_a_tiled):
65
- global args
66
-
67
- # Translate the prompt if it's in Korean
68
- prompt = translate_if_korean(prompt)
69
-
70
- args = load_message(height, width, seed, prompt, args, stage_a_tiled)
71
- torch.manual_seed(args.seed)
72
- random.seed(args.seed)
73
- np.random.seed(args.seed)
74
- dtype = torch.bfloat16 if args.dtype == 'bf16' else torch.float
75
-
76
- captions = [args.prompt] * args.num_image
77
- height, width = args.height, args.width
78
- batch_size = 1
79
- height_lr, width_lr = get_target_lr_size(height / width, std_size=32)
80
- stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)
81
- stage_c_latent_shape_lr, stage_b_latent_shape_lr = calculate_latent_sizes(height_lr, width_lr, batch_size=batch_size)
82
-
83
- # Stage C Parameters
84
- extras.sampling_configs['cfg'] = 4
85
- extras.sampling_configs['shift'] = 1
86
- extras.sampling_configs['timesteps'] = 20
87
- extras.sampling_configs['t_start'] = 1.0
88
- extras.sampling_configs['sampler'] = DDPMSampler(extras.gdf)
89
-
90
- # Stage B Parameters
91
- extras_b.sampling_configs['cfg'] = 1.1
92
- extras_b.sampling_configs['shift'] = 1
93
- extras_b.sampling_configs['timesteps'] = 10
94
- extras_b.sampling_configs['t_start'] = 1.0
95
-
96
- for _, caption in enumerate(captions):
97
- batch = {'captions': [caption] * batch_size}
98
- conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
99
- unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
100
-
101
- with torch.no_grad():
102
- models.generator.cuda()
103
- print('STAGE C GENERATION***************************')
104
- with torch.cuda.amp.autocast(dtype=dtype):
105
- sampled_c = generation_c(batch, models, extras, core, stage_c_latent_shape, stage_c_latent_shape_lr, device)
106
-
107
- models.generator.cpu()
108
- torch.cuda.empty_cache()
109
-
110
- conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
111
- unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
112
- conditions_b['effnet'] = sampled_c
113
- unconditions_b['effnet'] = torch.zeros_like(sampled_c)
114
- print('STAGE B + A DECODING***************************')
115
-
116
- with torch.cuda.amp.autocast(dtype=dtype):
117
- sampled = decode_b(conditions_b, unconditions_b, models_b, stage_b_latent_shape, extras_b, device, stage_a_tiled=args.stage_a_tiled)
118
-
119
- torch.cuda.empty_cache()
120
- imgs = show_images(sampled)
121
-
122
- return imgs[0]
123
-
124
- css = """
125
- footer {
126
- visibility: hidden;
127
- }
128
- """
129
-
130
- with gr.Blocks(theme="Nymbo/Nymbo_Theme", css=css) as demo:
131
- with gr.Column(elem_id="col-container"):
132
- gr.Markdown("<h1><center>초고해상도 UHD 이미지(최대 5120 X 4096 픽셀) 생성</center></h1>")
133
-
134
- with gr.Row():
135
- prompt = gr.Textbox(
136
- label="Text Prompt (한글 또는 영어로 입력하세요)",
137
- show_label=False,
138
- max_lines=1,
139
- placeholder="프롬프트를 입력하세요 (Enter your prompt in Korean or English)",
140
- container=False
141
- )
142
- polish_button = gr.Button("제출! (Submit!)", scale=0)
143
-
144
- output_img = gr.Image(label="Output Image", show_label=False)
145
-
146
- with gr.Accordion("Advanced Settings", open=False):
147
- seed = gr.Number(
148
- label="Random Seed",
149
- value=123,
150
- step=1,
151
- minimum=0,
152
- )
153
-
154
- with gr.Row():
155
- width = gr.Slider(
156
- label="Width",
157
- minimum=1536,
158
- maximum=5120,
159
- step=32,
160
- value=4096
161
- )
162
-
163
- height = gr.Slider(
164
- label="Height",
165
- minimum=1536,
166
- maximum=4096,
167
- step=32,
168
- value=2304
169
- )
170
-
171
- with gr.Row():
172
- cfg = gr.Slider(
173
- label="CFG",
174
- minimum=3,
175
- maximum=10,
176
- step=0.1,
177
- value=4
178
- )
179
-
180
- timesteps = gr.Slider(
181
- label="Timesteps",
182
- minimum=10,
183
- maximum=50,
184
- step=1,
185
- value=20
186
- )
187
-
188
- stage_a_tiled = gr.Checkbox(label="Stage_a_tiled", value=False)
189
-
190
- clear_button = gr.Button("Clear!")
191
-
192
- gr.Examples(
193
- examples=[
194
- "A detailed view of a blooming magnolia tree, with large, white flowers and dark green leaves, set against a clear blue sky.",
195
- "눈 덮인 산맥의 장엄한 전경, 푸른 하늘을 배경으로 한 고요한 호수가 있는 모습",
196
- "The image features a snow-covered mountain range with a large, snow-covered mountain in the background. The mountain is surrounded by a forest of trees, and the sky is filled with clouds. The scene is set during the winter season, with snow covering the ground and the trees.",
197
- "스웨터를 입은 악어",
198
- "A vibrant anime scene of a young girl with long, flowing pink hair, big sparkling blue eyes, and a school uniform, standing under a cherry blossom tree with petals falling around her. The background shows a traditional Japanese school with cherry blossoms in full bloom.",
199
- "골든 리트리버 강아지가 푸른 잔디밭에서 빨간 공을 쫓는 귀여운 모습",
200
- "A cozy, rustic log cabin nestled in a snow-covered forest, with smoke rising from the stone chimney, warm lights glowing from the windows, and a path of footprints leading to the front door.",
201
- "캐나다 밴프 국립공원의 아름다운 풍경, 청록색 호수와 눈 덮인 산들, 울창한 소나무 숲이 어우러진 모습",
202
- "귀여운 시츄가 욕조에서 목욕하는 모습, 거품에 둘러싸인 채 살짝 젖은 모습으로 카메라를 바라보고 있음",
203
- ],
204
- inputs=[prompt],
205
- outputs=[output_img],
206
- examples_per_page=5
207
- )
208
-
209
- polish_button.click(get_image, inputs=[height, width, seed, prompt, cfg, timesteps, stage_a_tiled], outputs=output_img)
210
- polish_button.click(clear_image, inputs=[], outputs=output_img)
211
-
212
- def download_with_wget(url, save_path):
213
- try:
214
- subprocess.run(['wget', url, '-O', save_path], check=True)
215
- print(f"Downloaded to {save_path}")
216
- except subprocess.CalledProcessError as e:
217
- print(f"Error downloading file: {e}")
218
-
219
- def download_model():
220
- urls = [
221
- 'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_a.safetensors',
222
- 'https://huggingface.co/stabilityai/StableWurst/resolve/main/previewer.safetensors',
223
- 'https://huggingface.co/stabilityai/StableWurst/resolve/main/effnet_encoder.safetensors',
224
- 'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_b_lite_bf16.safetensors',
225
- 'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_c_bf16.safetensors',
226
- ]
227
- for file_url in urls:
228
- hf_hub_download(repo_id="stabilityai/stable-cascade", filename=file_url.split('/')[-1], local_dir='models')
229
- hf_hub_download(repo_id="roubaofeipi/UltraPixel", filename='ultrapixel_t2i.safetensors', local_dir='models')
230
-
231
- if __name__ == "__main__":
232
- args = parse_args()
233
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
234
- download_model()
235
- config_file = args.config_c
236
- with open(config_file, "r", encoding="utf-8") as file:
237
- loaded_config = yaml.safe_load(file)
238
-
239
- core = WurstCoreC(config_dict=loaded_config, device=device, training=False)
240
-
241
- # SETUP STAGE B
242
- config_file_b = args.config_b
243
- with open(config_file_b, "r", encoding="utf-8") as file:
244
- config_file_b = yaml.safe_load(file)
245
-
246
- core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False)
247
-
248
- extras = core.setup_extras_pre()
249
- models = core.setup_models(extras)
250
- models.generator.eval().requires_grad_(False)
251
- print("STAGE C READY")
252
-
253
- extras_b = core_b.setup_extras_pre()
254
- models_b = core_b.setup_models(extras_b, skip_clip=True)
255
- models_b = WurstCoreB.Models(
256
- **{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model}
257
- )
258
- models_b.generator.bfloat16().eval().requires_grad_(False)
259
- print("STAGE B READY")
260
-
261
- pretrained_path = args.pretrained_path
262
- sdd = torch.load(pretrained_path, map_location='cpu')
263
- collect_sd = {}
264
- for k, v in sdd.items():
265
- collect_sd[k[7:]] = v
266
-
267
- models.train_norm.load_state_dict(collect_sd)
268
- models.generator.eval()
269
- models.train_norm.eval()
270
-
271
- demo.launch(debug=True, share=True, auth=("gini","pick"))