(pipe, trans_vae, args)
| 8 | import numpy as np |
| 9 | |
| 10 | def generate_img(pipe, trans_vae, args): |
| 11 | |
| 12 | latents = pipe( |
| 13 | prompt=args.prompt, |
| 14 | height=args.height, |
| 15 | width=args.width, |
| 16 | num_inference_steps=args.steps, |
| 17 | output_type="latent", |
| 18 | generator=torch.Generator("cuda").manual_seed(args.seed), |
| 19 | guidance_scale=args.guidance, |
| 20 | |
| 21 | ).images |
| 22 | |
| 23 | latents = pipe._unpack_latents(latents, args.height, args.width, pipe.vae_scale_factor) |
| 24 | latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor |
| 25 | |
| 26 | with torch.no_grad(): |
| 27 | original_x, x = trans_vae.decode(latents) |
| 28 | |
| 29 | x = x.clamp(0, 1) |
| 30 | x = x.permute(0, 2, 3, 1) |
| 31 | img = Image.fromarray((x*255).float().cpu().numpy().astype(np.uint8)[0]) |
| 32 | |
| 33 | return img |
| 34 | |
| 35 | if __name__ == "__main__": |
| 36 | parser = argparse.ArgumentParser() |
no test coverage detected