File size: 12,688 Bytes
045e583
 
 
 
 
 
 
 
 
 
 
 
 
40d5852
d122a2a
0df3aff
a76c46f
045e583
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6beb2ce
045e583
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
72ea199
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
045e583
72ea199
 
 
 
 
 
 
a205be7
72ea199
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a205be7
72ea199
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
045e583
 
 
72ea199
7d73d78
 
0df3aff
 
 
 
 
 
 
d122a2a
0df3aff
d122a2a
 
 
 
 
 
a76c46f
d122a2a
 
a76c46f
 
 
 
045e583
 
 
 
d122a2a
045e583
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
import os
import yaml
import torch
import sys
sys.path.append(os.path.abspath('./'))
from inference.utils import *
from train import WurstCoreB
from gdf import DDPMSampler
from train import WurstCore_t2i as WurstCoreC
import numpy as np
import random
import argparse
import gradio as gr
import spaces
from huggingface_hub import hf_hub_url
import subprocess
from huggingface_hub import hf_hub_download

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument( '--height', type=int, default=2560, help='image height')
    parser.add_argument('--width', type=int, default=5120, help='image width')
    parser.add_argument('--seed', type=int, default=123, help='random seed')
    parser.add_argument('--dtype', type=str, default='bf16', help=' if bf16 does not work, change it to float32 ')
    parser.add_argument('--config_c', type=str, 
    default='configs/training/t2i.yaml' ,help='config file for stage c, latent generation')
    parser.add_argument('--config_b', type=str, 
    default='configs/inference/stage_b_1b.yaml' ,help='config file for stage b, latent decoding')
    parser.add_argument( '--prompt', type=str,
     default='A photo-realistic image of a west highland white terrier in the garden, high quality, detail rich, 8K', help='text prompt')
    parser.add_argument( '--num_image', type=int, default=1, help='how many images generated')
    parser.add_argument( '--output_dir', type=str, default='figures/output_results/', help='output directory for generated image')
    parser.add_argument( '--stage_a_tiled', action='store_true', help='whther or nor to use tiled decoding for stage a to save memory')
    parser.add_argument( '--pretrained_path', type=str, default='models/ultrapixel_t2i.safetensors', help='pretrained path of newly added paramter of UltraPixel')
    args = parser.parse_args()
    return args

def clear_image():
    return None
def load_message(height, width, seed, prompt, args, stage_a_tiled):
    args.height = height
    args.width = width
    args.seed  = seed
    args.prompt = prompt + ' rich detail, 4k, high quality'
    args.stage_a_tiled = stage_a_tiled
    return args
@spaces.GPU(duration=120)
def get_image(height, width, seed, prompt, cfg, timesteps, stage_a_tiled):
    global args
    args = load_message(height, width, seed, prompt,  args, stage_a_tiled)
    torch.manual_seed(args.seed)
    random.seed(args.seed) 
    np.random.seed(args.seed)
    dtype = torch.bfloat16 if args.dtype == 'bf16' else torch.float

    captions = [args.prompt] * args.num_image
    height, width = args.height, args.width
    batch_size=1 
    height_lr, width_lr = get_target_lr_size(height / width, std_size=32)
    stage_c_latent_shape, stage_b_latent_shape = calculate_latent_sizes(height, width, batch_size=batch_size)
    stage_c_latent_shape_lr, stage_b_latent_shape_lr = calculate_latent_sizes(height_lr, width_lr, batch_size=batch_size)
   
    # Stage C Parameters
    extras.sampling_configs['cfg'] = 4
    extras.sampling_configs['shift'] = 1
    extras.sampling_configs['timesteps'] = 20
    extras.sampling_configs['t_start'] = 1.0
    extras.sampling_configs['sampler'] = DDPMSampler(extras.gdf)
    
    
    
    # Stage B Parameters
    extras_b.sampling_configs['cfg'] = 1.1
    extras_b.sampling_configs['shift'] = 1
    extras_b.sampling_configs['timesteps'] = 10
    extras_b.sampling_configs['t_start'] = 1.0

    for _, caption in enumerate(captions):

        
            batch = {'captions': [caption] * batch_size}
            #conditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
            #unconditions = core.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)    
            
            conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
            unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
            
            
            with torch.no_grad():
        
            
                models.generator.cuda()
                print('STAGE C GENERATION***************************')
                with torch.cuda.amp.autocast(dtype=dtype):
                    sampled_c = generation_c(batch, models, extras, core, stage_c_latent_shape, stage_c_latent_shape_lr, device)
                
                    
                    
                models.generator.cpu()
                torch.cuda.empty_cache()
                
                conditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=False)
                unconditions_b = core_b.get_conditions(batch, models_b, extras_b, is_eval=True, is_unconditional=True)
                conditions_b['effnet'] = sampled_c
                unconditions_b['effnet'] = torch.zeros_like(sampled_c)
                print('STAGE B + A DECODING***************************')
                
                with torch.cuda.amp.autocast(dtype=dtype):
                        sampled = decode_b(conditions_b, unconditions_b, models_b, stage_b_latent_shape, extras_b, device, stage_a_tiled=args.stage_a_tiled)
                
                torch.cuda.empty_cache()
                imgs = show_images(sampled)
                #for idx, img in enumerate(imgs):
                    #print(os.path.join(save_dir, args.prompt[:20]+'_' + str(cnt).zfill(5) + '.jpg'), idx)
                    #img.save(os.path.join(save_dir, args.prompt[:20]+'_' + str(cnt).zfill(5) + '.jpg'))
                    
    return imgs[0]           
    #print('finished! Results ')


with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("<h1><center>UltraPixel: Advancing Ultra-High-Resolution Image Synthesis to New Peaks </center></h1>")
        
        with gr.Row():
            prompt = gr.Textbox(
                label="Text Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False
            )
            polish_button = gr.Button("Submit!", scale=0)
        
        output_img = gr.Image(label="Output Image", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Number(
                label="Random Seed",
                value=123,
                step=1,
                minimum=0,
                #maximum=MAX_SEED
            )
            
            #randomize_seed = gr.Checkbox(label="Randomize seed", value=False)
            
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=1536,
                    maximum=5120,
                    step=32,
                    value=4096
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=1536,
                    maximum=4096,
                    step=32,
                    value=2304
                )
            
            with gr.Row():
                cfg = gr.Slider(
                    label="CFG",
                    minimum=3,
                    maximum=10,
                    step=0.1,
                    value=4
                )
                
                timesteps = gr.Slider(
                    label="Timesteps",
                    minimum=10,
                    maximum=50,
                    step=1,
                    value=20
                )
            
            stage_a_tiled = gr.Checkbox(label="Stage_a_tiled", value=False)
        
        clear_button = gr.Button("Clear!")
        
        gr.Examples(
            examples=[
                "A detailed view of a blooming magnolia tree, with large, white flowers and dark green leaves, set against a clear blue sky.",
                "A close-up portrait of a young woman with flawless skin, vibrant red lipstick, and wavy brown hair, wearing a vintage floral dress and standing in front of a blooming garden.",
                "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.",
                "Crocodile in a sweater.",
                "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.",
                "A playful Labrador retriever puppy with a shiny, golden coat, chasing a red ball in a spacious backyard, with green grass and a wooden fence.",
                "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.",
                "A highly detailed, high-quality image of the Banff National Park in Canada. The turquoise waters of Lake Louise are surrounded by snow-capped mountains and dense pine forests. A wooden canoe is docked at the edge of the lake. The sky is a clear, bright blue, and the air is crisp and fresh.",
                "A highly detailed, high-quality image of a Shih Tzu receiving a bath in a home bathroom. The dog is standing in a tub, covered in suds, with a slightly wet and adorable look. The background includes bathroom fixtures, towels, and a clean, tiled floor.",
            ],
            inputs=[prompt],
            outputs=[output_img],
            examples_per_page=5
        )
        
        polish_button.click(get_image, inputs=[height, width, seed, prompt, cfg, timesteps, stage_a_tiled], outputs=output_img)           
        polish_button.click(clear_image, inputs=[], outputs=output_img)
   


def download_with_wget(url, save_path):
    
    try:
        subprocess.run(['wget', url, '-O', save_path], check=True)
        print(f"Downloaded to {save_path}")
    except subprocess.CalledProcessError as e:
        print(f"Error downloading file: {e}")
def download_model():
    
    urls = [
        'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_a.safetensors',
        'https://huggingface.co/stabilityai/StableWurst/resolve/main/previewer.safetensors',
        'https://huggingface.co/stabilityai/StableWurst/resolve/main/effnet_encoder.safetensors',
        'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_b_lite_bf16.safetensors', 
        'https://huggingface.co/stabilityai/StableWurst/resolve/main/stage_c_bf16.safetensors', 
        
    ]
    for file_url in urls:
        hf_hub_download(repo_id="stabilityai/stable-cascade", filename=file_url.split('/')[-1], local_dir='models')
    #   'https://huggingface.co/roubaofeipi/UltraPixel/blob/main/ultrapixel_t2i.safetensors'  
    hf_hub_download(repo_id="roubaofeipi/UltraPixel", filename='ultrapixel_t2i.safetensors', local_dir='models')
    
if __name__ == "__main__":
   
    args = parse_args()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    download_model()
    config_file = args.config_c
    with open(config_file, "r", encoding="utf-8") as file:
        loaded_config = yaml.safe_load(file)
    
    core = WurstCoreC(config_dict=loaded_config, device=device, training=False)
    
    # SETUP STAGE B
    config_file_b = args.config_b
    with open(config_file_b, "r", encoding="utf-8") as file:
        config_file_b = yaml.safe_load(file)
        
    core_b = WurstCoreB(config_dict=config_file_b, device=device, training=False)
    
    extras = core.setup_extras_pre()
    models = core.setup_models(extras)
    models.generator.eval().requires_grad_(False)
    print("STAGE C READY")
    
    extras_b = core_b.setup_extras_pre()
    models_b = core_b.setup_models(extras_b, skip_clip=True)
    models_b = WurstCoreB.Models(
       **{**models_b.to_dict(), 'tokenizer': models.tokenizer, 'text_model': models.text_model}
    )
    models_b.generator.bfloat16().eval().requires_grad_(False)
    print("STAGE B READY")
    
    pretrained_path = args.pretrained_path    
    sdd = torch.load(pretrained_path, map_location='cpu')
    collect_sd = {}
    for k, v in sdd.items():
        collect_sd[k[7:]] = v
    
    models.train_norm.load_state_dict(collect_sd)
    models.generator.eval()
    models.train_norm.eval()
    
    
    demo.launch(
            debug=True, share=True, 
        )