Skip to content

Commit f3f60fc

Browse files
correct merg
2 parents eef5da9 + 536eb1a commit f3f60fc

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

scripts/txt2img.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -239,9 +239,7 @@ def main():
239239
if opt.fixed_code:
240240
start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
241241

242-
print("start code", start_code.abs().sum())
243242
precision_scope = autocast if opt.precision=="autocast" else nullcontext
244-
precision_scope = nullcontext
245243
with torch.no_grad():
246244
with precision_scope("cuda"):
247245
with model.ema_scope():
@@ -297,6 +295,13 @@ def main():
297295
Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
298296
grid_count += 1
299297

298+
image = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy()
299+
300+
# run safety checker
301+
safety_checker_input = pipe.feature_extractor(numpy_to_pil(image), return_tensors="pt")
302+
image, has_nsfw_concept = pipe.safety_checker(images=image, clip_input=safety_checker_input.pixel_values)
303+
toc = time.time()
304+
300305
print(f"Your samples are ready and waiting for you here: \n{outpath} \n"
301306
f" \nEnjoy.")
302307

0 commit comments

Comments
 (0)