diff --git a/.gitignore b/.gitignore index 438cb08..33b663b 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ node_modules .output .env dist +outputs diff --git a/server.py b/server.py index 9d955e3..5a6aabe 100644 --- a/server.py +++ b/server.py @@ -210,179 +210,190 @@ def put_watermark(img, wm_encoder=None): img = Image.fromarray(img[:, :, ::-1]) return img +import time +import requests -def main(opt): - seed_everything(opt.seed) +def main_dev(opt): + model_name = '' # 默认模型 + model = None # 默认模型 + config = None # 默认配置 + device = None # 默认设备 + while True: + time.sleep(2) # 延时1s执行, 避免cpu占用过高 + data = requests.get("http://localhost:3000/api/drawing").json() # 从局域网中获取一组参数 + print(data) # [{'model': '768-v-ema', 'prompt': '一只猫', 'watermark': '0'}, {'model': '768-v-ema', 'prompt': '一只狗', 'watermark': '0'}] + # 遍历 data 返回dict + for item in data: + print(item) # {'model': '768-v-ema', 'prompt': '一只猫', 'watermark': '0'} + # 设置参数 + if 'prompt' in item: opt.prompt = item['prompt'] # 描述 + if 'n_samples' in item: opt.n_samples = item['n_samples'] # 列数 + if 'n_rows' in item: opt.n_rows = item['n_rows'] # 行数 + if 'scale' in item: opt.scale = item['scale'] # 比例 - config = OmegaConf.load(f"{opt.config}") - device = torch.device("cuda") if opt.device == "cuda" else torch.device("cpu") - model = load_model_from_config(config, f"{opt.ckpt}", device) - - if opt.plms: - sampler = PLMSSampler(model, device=device) - elif opt.dpm: - sampler = DPMSolverSampler(model, device=device) - else: - sampler = DDIMSampler(model, device=device) - - os.makedirs(opt.outdir, exist_ok=True) - outpath = opt.outdir - - print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...") - wm = "SDV2" - wm_encoder = WatermarkEncoder() - wm_encoder.set_watermark('bytes', wm.encode('utf-8')) - - batch_size = opt.n_samples - n_rows = opt.n_rows if opt.n_rows > 0 else batch_size - if not opt.from_file: - prompt = opt.prompt - assert prompt is not None - data = [batch_size * [prompt]] - - else: - print(f"reading prompts from {opt.from_file}") - with open(opt.from_file, "r") as f: - data = f.read().splitlines() - data = [p for p in data for i in range(opt.repeat)] - data = list(chunk(data, batch_size)) - - sample_path = os.path.join(outpath, "samples") - os.makedirs(sample_path, exist_ok=True) - sample_count = 0 - base_count = len(os.listdir(sample_path)) - grid_count = len(os.listdir(outpath)) - 1 - - start_code = None - if opt.fixed_code: - start_code = torch.randn([opt.n_samples, 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 - 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, + # 如果模型不同，重新加载模型(注意释放内存) + if item['model'] != model_name: + # 获取环境配置 + model_name = item['model'] + opt.config = f'/data/{model_name}.yaml' + opt.ckpt = f'/data/{model_name}.ckpt' + opt.device = 'cuda' + print(f"config: {opt.config}", f"ckpt: {opt.ckpt}", f"device: {opt.device}") + config = OmegaConf.load(f"{opt.config}") + device = torch.device("cuda") if opt.device == "cuda" else torch.device("cpu") + # 加载模型(到显存) + print(f"load model: {item['model']}..") + model_name = item['model'] + model = load_model_from_config(config, f"{opt.ckpt}", device) + print(f"model_name: {model_name}") + # 使用指定的模型和配置文件进行推理一组参数 + if opt.plms: + sampler = PLMSSampler(model, device=device) + elif opt.dpm: + sampler = DPMSolverSampler(model, device=device) + else: + sampler = DDIMSampler(model, device=device) + # 检查输出目录是否存在 + os.makedirs(opt.outdir, exist_ok=True) + outpath = opt.outdir + # 创建水印编码器 + wm = "SDV2" + wm_encoder = WatermarkEncoder() + wm_encoder.set_watermark('bytes', wm.encode('utf-8')) + # x + batch_size = opt.n_samples + n_rows = opt.n_rows if opt.n_rows > 0 else batch_size + if not opt.from_file: + prompt = opt.prompt + assert prompt is not None + data = [batch_size * [prompt]] + else: + print(f"reading prompts from {opt.from_file}") + with open(opt.from_file, "r") as f: + data = f.read().splitlines() + data = [p for p in data for i in range(opt.repeat)] + data = list(chunk(data, batch_size)) + # x + sample_path = os.path.join(outpath, "samples") + os.makedirs(sample_path, exist_ok=True) + sample_count = 0 + base_count = len(os.listdir(sample_path)) + grid_count = len(os.listdir(outpath)) - 1 + # x + start_code = None + if opt.fixed_code: + start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device) + # x + 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=opt.n_samples, + batch_size=batch_size, 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.") - + 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 + 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 + print(f"Your samples are ready and waiting for you here: \n{outpath} \n" + f" \nEnjoy.") + break if __name__ == "__main__": opt = parse_args() - main(opt) + main_dev(opt) diff --git a/venv b/venv new file mode 120000 index 0000000..de9d3d0 --- /dev/null +++ b/venv @@ -0,0 +1 @@ +/data/stablediffusion/venv \ No newline at end of file