From a751623e4debb06ab3f4ff366ab613237eddff38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=A1=9C=E8=8F=AF?= Date: Sat, 18 Feb 2023 21:13:41 +0800 Subject: [PATCH] =?UTF-8?q?=E6=8C=87=E5=AE=9A=E6=89=A7=E8=A1=8C=E6=AC=A1?= =?UTF-8?q?=E6=95=B0,=20=E5=8F=96=E6=B6=88=E5=AE=AB=E6=A0=BC=E5=9B=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- server.py | 113 ++++++++++++++++++++++++++++++++++-------------------- 1 file changed, 71 insertions(+), 42 deletions(-) diff --git a/server.py b/server.py index 6dd975d..8b55b84 100644 --- a/server.py +++ b/server.py @@ -355,48 +355,77 @@ def main_dev(opt): 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 - with torch.no_grad(), \ - precision_scope(opt.device), \ - model.ema_scope(): - all_samples = list() - for n in trange(opt.n_iter, desc="Sampling"): - for prompts in tqdm(data, desc="data"): - uc = None - 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 + with torch.no_grad(), precision_scope(opt.device), model.ema_scope(): + all_samples = list() + # 执行指定的次数 + for n in trange(item['number'], desc="Sampling"): + print("Sampling:", n) + for prompts in tqdm(data, desc="data"): + uc = None + 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) + print("Sample count:", sample_count) + #for n in trange(opt.n_iter, desc="Sampling"): + # for prompts in tqdm(data, desc="data"): + # uc = None + # 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)