Spaces:
Running
on
Zero
Running
on
Zero
File size: 8,513 Bytes
045e583 40d5852 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 |
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
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():
with gr.Row():
with gr.Column():
height = gr.Slider(value=2304, step=32, minimum=1536, maximum=4096, label='Height')
width = gr.Slider(value=4096, step=32, minimum=1536, maximum=5120, label='Width')
seed = gr.Number(value=123, step=1, label='Random Seed')
prompt = gr.Textbox(value='', max_lines=4, label='Text Prompt')
cfg = gr.Slider(value=4, step=0.1, minimum=3, maximum=10, label='CFG')
timesteps = gr.Slider(value=20, step=1, minimum=10, maximum=50, label='Timesteps')
stage_a_tiled = gr.Checkbox(value=False, label='Stage_a_tiled')
with gr.Row():
clear_button = gr.Button("Clear!")
polish_button = gr.Button("Submit!")
with gr.Column():
output_img = gr.Image(label='Output Image', sources=None)
with gr.Column():
prompt2 = gr.Textbox(
value='''
1. a happy cat
2. a happy girl
''', label='Text prompt examples'
)
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)
if __name__ == "__main__":
args = parse_args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
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,
) |