(
img,
model,
sampler,
value_dict,
num_samples,
force_uc_zero_embeddings=[],
additional_kwargs={},
offset_noise_level: int = 0.0,
return_latents=False,
skip_encode=False,
filter=None,
add_noise=True,
)
| 656 | |
| 657 | @torch.no_grad() |
| 658 | def do_img2img( |
| 659 | img, |
| 660 | model, |
| 661 | sampler, |
| 662 | value_dict, |
| 663 | num_samples, |
| 664 | force_uc_zero_embeddings=[], |
| 665 | additional_kwargs={}, |
| 666 | offset_noise_level: int = 0.0, |
| 667 | return_latents=False, |
| 668 | skip_encode=False, |
| 669 | filter=None, |
| 670 | add_noise=True, |
| 671 | ): |
| 672 | st.text("Sampling") |
| 673 | |
| 674 | outputs = st.empty() |
| 675 | precision_scope = autocast |
| 676 | with torch.no_grad(): |
| 677 | with precision_scope("cuda"): |
| 678 | with model.ema_scope(): |
| 679 | load_model(model.conditioner) |
| 680 | batch, batch_uc = get_batch( |
| 681 | get_unique_embedder_keys_from_conditioner(model.conditioner), |
| 682 | value_dict, |
| 683 | [num_samples], |
| 684 | ) |
| 685 | c, uc = model.conditioner.get_unconditional_conditioning( |
| 686 | batch, |
| 687 | batch_uc=batch_uc, |
| 688 | force_uc_zero_embeddings=force_uc_zero_embeddings, |
| 689 | ) |
| 690 | unload_model(model.conditioner) |
| 691 | for k in c: |
| 692 | c[k], uc[k] = map(lambda y: y[k][:num_samples].to("cuda"), (c, uc)) |
| 693 | |
| 694 | for k in additional_kwargs: |
| 695 | c[k] = uc[k] = additional_kwargs[k] |
| 696 | if skip_encode: |
| 697 | z = img |
| 698 | else: |
| 699 | load_model(model.first_stage_model) |
| 700 | z = model.encode_first_stage(img) |
| 701 | unload_model(model.first_stage_model) |
| 702 | |
| 703 | noise = torch.randn_like(z) |
| 704 | |
| 705 | sigmas = sampler.discretization(sampler.num_steps).cuda() |
| 706 | sigma = sigmas[0] |
| 707 | |
| 708 | st.info(f"all sigmas: {sigmas}") |
| 709 | st.info(f"noising sigma: {sigma}") |
| 710 | if offset_noise_level > 0.0: |
| 711 | noise = noise + offset_noise_level * append_dims( |
| 712 | torch.randn(z.shape[0], device=z.device), z.ndim |
| 713 | ) |
| 714 | if add_noise: |
| 715 | noised_z = z + noise * append_dims(sigma, z.ndim).cuda() |
no test coverage detected