r
/
stable-diffusion-finetune
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

optionally fix start code for sampling

This commit is contained in:
rromb 2022-07-15 13:50:15 +02:00
parent 5c3f6795fa
commit 37e59ee487
1 changed files with 16 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
scripts/txt2img.py
View File

@ -83,6 +83,11 @@ def main():
action='store_true', action='store_true',
help="use plms sampling", help="use plms sampling",
) )
parser.add_argument(
"--fixed_code",
action='store_true',
help="if enabled, uses the same starting code across all samples ",
)
parser.add_argument( parser.add_argument(
"--ddim_eta", "--ddim_eta",
@ -155,7 +160,6 @@ def main():
type=str, type=str,
help="if specified, load prompts from this file", help="if specified, load prompts from this file",
) )
parser.add_argument( parser.add_argument(
"--config", "--config",
type=str, type=str,
@ -209,6 +213,10 @@ def main():
base_count = len(os.listdir(sample_path)) base_count = len(os.listdir(sample_path))
grid_count = len(os.listdir(outpath)) - 1 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)
with torch.no_grad(): with torch.no_grad():
with model.ema_scope(): with model.ema_scope():
tic = time.time() tic = time.time()
@ -221,7 +229,7 @@ def main():
if isinstance(prompts, tuple): if isinstance(prompts, tuple):
prompts = list(prompts) prompts = list(prompts)
c = model.get_learned_conditioning(prompts) c = model.get_learned_conditioning(prompts)
shape = [opt.C, opt.H//opt.f, opt.W//opt.f] shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
samples_ddim, _ = sampler.sample(S=opt.ddim_steps, samples_ddim, _ = sampler.sample(S=opt.ddim_steps,
conditioning=c, conditioning=c,
batch_size=opt.n_samples, batch_size=opt.n_samples,
@ -230,15 +238,17 @@ def main():
unconditional_guidance_scale=opt.scale, unconditional_guidance_scale=opt.scale,
unconditional_conditioning=uc, unconditional_conditioning=uc,
eta=opt.ddim_eta, eta=opt.ddim_eta,
dynamic_threshold=opt.dyn) dynamic_threshold=opt.dyn,
x_T=start_code)
x_samples_ddim = model.decode_first_stage(samples_ddim) x_samples_ddim = model.decode_first_stage(samples_ddim)
x_samples_ddim = torch.clamp((x_samples_ddim+1.0)/2.0, min=0.0, max=1.0) x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
if not opt.skip_save: if not opt.skip_save:
for x_sample in x_samples_ddim: for x_sample in x_samples_ddim:
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')
Image.fromarray(x_sample.astype(np.uint8)).save(os.path.join(sample_path, f"{base_count:05}.png")) Image.fromarray(x_sample.astype(np.uint8)).save(
os.path.join(sample_path, f"{base_count:05}.png"))
base_count += 1 base_count += 1
all_samples.append(x_samples_ddim) all_samples.append(x_samples_ddim)
@ -256,7 +266,7 @@ def main():
toc = time.time() toc = time.time()
print(f"Your samples are ready and waiting for you here: \n{outpath} \n" print(f"Your samples are ready and waiting for you here: \n{outpath} \n"
f"Sampling took {toc-tic}s, i.e. produced {opt.n_iter * opt.n_samples / (toc - tic):.2f} samples/sec." f"Sampling took {toc - tic}s, i.e. produced {opt.n_iter * opt.n_samples / (toc - tic):.2f} samples/sec."
f" \nEnjoy.") f" \nEnjoy.")