gameui/drawing
2
0
Fork 0
Code Issues 1 Pull Requests Packages Projects Releases Wiki Activity

py递归未优化, 更换为while

Browse Source
This commit is contained in:
桜華
2023-02-18 23:50:16 +08:00
parent 2465716169
commit 87a9c5bd35
2 changed files with 37 additions and 140 deletions
Download Patch File Download Diff File Expand all files Collapse all files

176
server.py
View File

@@ -216,11 +216,21 @@ def put_watermark(img, wm_encoder=None):
# 对任务状态的修改 # 对任务状态的修改
def update_task_status(task: dict, status: str, progress: int): def update_task_status(task: dict, status: str, progress: int, data: list = []):
task["status"] = status task["status"] = status
task["progress"] = progress task["progress"] = progress
task["data"] = data
requests.put(f"http://localhost:3000/api/drawing/{task['id']}", json=task) requests.put(f"http://localhost:3000/api/drawing/{task['id']}", json=task)
# 从局域网中获取一组任务(如果列表为空，等待2s后重新获取)
def get_tasks():
tasks = requests.get("http://localhost:3000/api/drawing?status=waiting").json()
if len(tasks) == 0:
while len(tasks) == 0:
print('no task, wait 2s...')
time.sleep(2)
tasks = requests.get("http://localhost:3000/api/drawing?status=waiting").json()
return tasks
def main_dev(opt): def main_dev(opt):
model_name = '' # 默认模型 model_name = '' # 默认模型
@@ -228,29 +238,14 @@ def main_dev(opt):
config = None # 默认配置 config = None # 默认配置
device = None # 默认设备 device = None # 默认设备
while True: while True:
time.sleep(2) # 延时1s执行, 避免cpu占用过高 for task in get_tasks(): # 遍历 tasks 返回 dict
data = requests.get("http://localhost:3000/api/drawing").json() # 从局域网中获取一组参数 print('task:', task) # 打印任务
print(data) update_task_status(task, "running", 0) # 更新任务状态为运行中
# 遍历 data 返回dict
for item in data:
print(item)
update_task_status(item, "running", 0) # 更新任务状态为运行中
# 设置参数
if 'prompt' in item:
opt.prompt = item['prompt'] # 描述
if 'number' in item:
opt.n_samples = item['number'] # 列数
print(f"n_samples: {opt.n_samples}")
#if 'n_rows' in item:
# opt.n_rows = item['n_rows'] # 行数
if 'scale' in item:
opt.scale = item['scale'] # 比例
# 如果模型不同，重新加载模型(注意释放内存) # 如果模型不同，重新加载模型(注意释放内存)
if item['ckpt'] != model_name: if task['ckpt'] != model_name:
# 获取环境配置 # 获取环境配置
model_name = item['ckpt'] model_name = task['ckpt']
opt.config = f'/data/{model_name}.yaml' opt.config = f'/data/{model_name}.yaml'
opt.ckpt = f'/data/{model_name}.ckpt' opt.ckpt = f'/data/{model_name}.ckpt'
opt.device = 'cuda' opt.device = 'cuda'
@@ -261,6 +256,7 @@ def main_dev(opt):
print(f"加载模型到显存: {model_name}..") print(f"加载模型到显存: {model_name}..")
model = load_model_from_config(config, f"{opt.ckpt}", device) model = load_model_from_config(config, f"{opt.ckpt}", device)
print(f"加载到显存完成: {model_name}") print(f"加载到显存完成: {model_name}")
# 使用指定的模型和配置文件进行推理一组参数 # 使用指定的模型和配置文件进行推理一组参数
if opt.plms: if opt.plms:
sampler = PLMSSampler(model, device=device) sampler = PLMSSampler(model, device=device)
@@ -268,104 +264,46 @@ def main_dev(opt):
sampler = DPMSolverSampler(model, device=device) sampler = DPMSolverSampler(model, device=device)
else: else:
sampler = DDIMSampler(model, device=device) sampler = DDIMSampler(model, device=device)
# 检查输出目录是否存在 # 检查输出目录是否存在
os.makedirs(opt.outdir, exist_ok=True) os.makedirs(opt.outdir, exist_ok=True)
outpath = opt.outdir outpath = opt.outdir
# 创建水印编码器 # 创建水印编码器
wm = "SDV2" wm = "SDV2"
wm_encoder = WatermarkEncoder() wm_encoder = WatermarkEncoder()
wm_encoder.set_watermark('bytes', wm.encode('utf-8')) wm_encoder.set_watermark('bytes', wm.encode('utf-8'))
# x # x
batch_size = opt.n_samples batch_size = task['number']
#n_rows = opt.n_rows if opt.n_rows > 0 else batch_size
if not opt.from_file: if not opt.from_file:
prompt = opt.prompt prompt = task['prompt']
assert prompt is not None assert prompt is not None
data = [batch_size * [prompt]] data = [batch_size * [prompt]]
print("data:", data)
else: else:
print(f"reading prompts from {opt.from_file}") print(f"reading prompts from {opt.from_file}")
with open(opt.from_file, "r") as f: with open(opt.from_file, "r") as f:
data = f.read().splitlines() data = f.read().splitlines()
data = [p for p in data for i in range(opt.repeat)] data = [p for p in data for i in range(opt.repeat)]
data = list(chunk(data, batch_size)) data = list(chunk(data, batch_size))
print("data:", data)
# x # x
sample_path = os.path.join(outpath, "samples") sample_path = os.path.join(outpath, "samples")
os.makedirs(sample_path, exist_ok=True) os.makedirs(sample_path, exist_ok=True)
sample_count = 0 sample_count = 0
base_count = len(os.listdir(sample_path)) base_count = len(os.listdir(sample_path))
grid_count = len(os.listdir(outpath)) - 1
# x # x
start_code = None start_code = None
if opt.fixed_code: if opt.fixed_code:
start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device) start_code = torch.randn([task['number'], opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
# 生成图片
'''
# 切换模型
if opt.torchscript or opt.ipex:
transformer = model.cond_stage_model.model
unet = model.model.diffusion_model
decoder = model.first_stage_model.decoder
additional_context = torch.cpu.amp.autocast() if opt.bf16 else nullcontext()
shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
if opt.bf16 and not opt.torchscript and not opt.ipex:
raise ValueError('Bfloat16 is supported only for torchscript+ipex')
if opt.bf16 and unet.dtype != torch.bfloat16:
raise ValueError("Use configs/stable-diffusion/intel/ configs with bf16 enabled if you'd like to use bfloat16 with CPU.")
if unet.dtype == torch.float16 and device == torch.device("cpu"):
raise ValueError("Use configs/stable-diffusion/intel/ configs for your model if you'd like to run it on CPU.")
if opt.ipex:
import intel_extension_for_pytorch as ipex
bf16_dtype = torch.bfloat16 if opt.bf16 else None
transformer = transformer.to(memory_format=torch.channels_last)
transformer = ipex.optimize(transformer, level="O1", inplace=True)
unet = unet.to(memory_format=torch.channels_last)
unet = ipex.optimize(unet, level="O1", auto_kernel_selection=True, inplace=True, dtype=bf16_dtype)
decoder = decoder.to(memory_format=torch.channels_last)
decoder = ipex.optimize(decoder, level="O1", auto_kernel_selection=True, inplace=True, dtype=bf16_dtype)
if opt.torchscript:
with torch.no_grad(), additional_context:
# get UNET scripted
if unet.use_checkpoint:
raise ValueError("Gradient checkpoint won't work with tracing. Use configs/stable-diffusion/intel/ configs for your model or disable checkpoint in your config.")
img_in = torch.ones(2, 4, 96, 96, dtype=torch.float32)
t_in = torch.ones(2, dtype=torch.int64)
context = torch.ones(2, 77, 1024, dtype=torch.float32)
scripted_unet = torch.jit.trace(unet, (img_in, t_in, context))
scripted_unet = torch.jit.optimize_for_inference(scripted_unet)
print(type(scripted_unet))
model.model.scripted_diffusion_model = scripted_unet
# get Decoder for first stage model scripted
samples_ddim = torch.ones(1, 4, 96, 96, dtype=torch.float32)
scripted_decoder = torch.jit.trace(decoder, (samples_ddim))
scripted_decoder = torch.jit.optimize_for_inference(scripted_decoder)
print(type(scripted_decoder))
model.first_stage_model.decoder = scripted_decoder
prompts = data[0]
print("Running a forward pass to initialize optimizations")
uc = None
if opt.scale != 1.0:
uc = model.get_learned_conditioning(batch_size * [""])
if isinstance(prompts, tuple):
prompts = list(prompts)
with torch.no_grad(), additional_context:
for _ in range(3):
c = model.get_learned_conditioning(prompts)
samples_ddim, _ = sampler.sample(S=5,
conditioning=c,
batch_size=batch_size,
shape=shape,
verbose=False,
unconditional_guidance_scale=opt.scale,
unconditional_conditioning=uc,
eta=opt.ddim_eta,
x_T=start_code)
print("Running a forward pass for decoder")
for _ in range(3):
x_samples_ddim = model.decode_first_stage(samples_ddim)
precision_scope = autocast if opt.precision == "autocast" or opt.bf16 else nullcontext precision_scope = autocast if opt.precision == "autocast" or opt.bf16 else nullcontext
with torch.no_grad(), precision_scope(opt.device), model.ema_scope(): with torch.no_grad(), precision_scope(opt.device), model.ema_scope():
#all_samples = list() images = []
# 执行指定的任务批次 (row)(item['number']) # 执行指定的任务批次 (row)(task['number'])
for n in trange(1, desc="Sampling"): for n in trange(1, desc="Sampling"):
print("Sampling:", data) print("Sampling:", data)
for prompts in tqdm(data, desc="data"): for prompts in tqdm(data, desc="data"):
@@ -378,7 +316,7 @@ def main_dev(opt):
shape = [opt.C, opt.H // opt.f, opt.W // opt.f] shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
samples, _ = sampler.sample(S=opt.steps, samples, _ = sampler.sample(S=opt.steps,
conditioning=c, conditioning=c,
batch_size=opt.n_samples, batch_size=task['number'],
shape=shape, shape=shape,
verbose=False, verbose=False,
unconditional_guidance_scale=opt.scale, unconditional_guidance_scale=opt.scale,
@@ -389,58 +327,18 @@ def main_dev(opt):
x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0) x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
for x_sample in x_samples: for x_sample in x_samples:
print("Sample count:", sample_count) print("Sample count:", sample_count)
imge_path = os.path.join(sample_path, f"{base_count:05}.png")
x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c') x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
img = Image.fromarray(x_sample.astype(np.uint8)) img = Image.fromarray(x_sample.astype(np.uint8))
img = put_watermark(img, wm_encoder) img = put_watermark(img, wm_encoder)
img.save(os.path.join(sample_path, f"{base_count:05}.png")) img.save(imge_path)
base_count += 1 base_count += 1
sample_count += 1 sample_count += 1
#all_samples.append(x_samples) images.append(imge_path)
print("Sample count:", sample_count) print("Sample count:", sample_count)
# for n in trange(opt.n_iter, desc="Sampling"): update_task_status(task=task, status='done', progress=1, data=images) # 修改任务状态为完成
# for prompts in tqdm(data, desc="data"): print("批次任务结束..")
# uc = None #break
# if opt.scale != 1.0:
# uc = model.get_learned_conditioning(batch_size * [""])
# if isinstance(prompts, tuple):
# prompts = list(prompts)
# c = model.get_learned_conditioning(prompts)
# shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
# samples, _ = sampler.sample(S=opt.steps,
# conditioning=c,
# batch_size=opt.n_samples,
# shape=shape,
# verbose=False,
# unconditional_guidance_scale=opt.scale,
# unconditional_conditioning=uc,
# eta=opt.ddim_eta,
# x_T=start_code)
# x_samples = model.decode_first_stage(samples)
# x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
# for x_sample in x_samples:
# x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
# img = Image.fromarray(x_sample.astype(np.uint8))
# img = put_watermark(img, wm_encoder)
# img.save(os.path.join(sample_path, f"{base_count:05}.png"))
# base_count += 1
# sample_count += 1
# all_samples.append(x_samples)
# additionally, save as grid
#grid = torch.stack(all_samples, 0)
#grid = rearrange(grid, 'n b c h w -> (n b) c h w')
#grid = make_grid(grid, nrow=n_rows)
# to image
#grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
#grid = Image.fromarray(grid.astype(np.uint8))
#grid = put_watermark(grid, wm_encoder)
#grid.save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
#grid_count += 1
print(f"Your samples are ready and waiting for you here: \n{outpath} \n", f" \nEnjoy.")
update_task_status(task=item, status='done', progress=1) # 修改任务状态为完成
'''
print("任务结束, 等待10s后退出..")
#time.sleep(10)
break
if __name__ == "__main__": if __name__ == "__main__":

1
venv
View File

@@ -1 +0,0 @@
/data/stablediffusion/venv