Views
No views yet
1from huggingface_hub import snapshot_download
2from optimum.intel.openvino import OVStableDiffusionPipeline
3from optimum.intel.openvino.modeling_diffusion import OVModelVaeDecoder, OVModelVaeEncoder, OVBaseModel
4
5# Create class wrappers which allow us to specify model_dir of TAESD instead of original pipeline dir
6
7class CustomOVModelVaeDecoder(OVModelVaeDecoder):
8 def __init__(
9 self, model, parent_model, ov_config = None, model_dir = None,
10 ):
11 super(OVModelVaeDecoder, self).__init__(model, parent_model, ov_config, "vae_decoder", model_dir)
12
13class CustomOVModelVaeEncoder(OVModelVaeEncoder):
14 def __init__(
15 self, model, parent_model, ov_config = None, model_dir = None,
16 ):
17 super(OVModelVaeEncoder, self).__init__(model, parent_model, ov_config, "vae_encoder", model_dir)
18
19pipe = OVStableDiffusionPipeline.from_pretrained("OpenVINO/stable-diffusion-1-5-fp32", compile=False)
20
21# Inject TAESD
22
23taesd_dir = snapshot_download(repo_id="deinferno/taesd-openvino")
24pipe.vae_decoder = CustomOVModelVaeDecoder(model = OVBaseModel.load_model(f"{taesd_dir}/vae_decoder/openvino_model.xml"), parent_model = pipe, model_dir = taesd_dir)
25pipe.vae_encoder = CustomOVModelVaeEncoder(model = OVBaseModel.load_model(f"{taesd_dir}/vae_encoder/openvino_model.xml"), parent_model = pipe, model_dir = taesd_dir)
26
27pipe.reshape(batch_size=1, height=512, width=512, num_images_per_prompt=1)
28pipe.compile()
29
30prompt = "plant pokemon in jungle"
31output = pipe(prompt, num_inference_steps=50, output_type="pil")
32output.images[0].save("result.png")