(model, controller, latents, context, t, guidance_scale, low_resource=False)
| 62 | |
| 63 | |
| 64 | def diffusion_step(model, controller, latents, context, t, guidance_scale, low_resource=False): |
| 65 | if low_resource: |
| 66 | noise_pred_uncond = model.unet(latents, t, encoder_hidden_states=context[0])["sample"] |
| 67 | noise_prediction_text = model.unet(latents, t, encoder_hidden_states=context[1])["sample"] |
| 68 | else: |
| 69 | latents_input = torch.cat([latents] * 2) |
| 70 | noise_pred = model.unet(latents_input, t, encoder_hidden_states=context)["sample"] |
| 71 | noise_pred_uncond, noise_prediction_text = noise_pred.chunk(2) |
| 72 | noise_pred = noise_pred_uncond + guidance_scale * (noise_prediction_text - noise_pred_uncond) |
| 73 | latents = model.scheduler.step(noise_pred, t, latents)["prev_sample"] |
| 74 | latents = controller.step_callback(latents) |
| 75 | return latents |
| 76 | |
| 77 | |
| 78 | def latent2image(vae, latents): |
no test coverage detected