diff --git a/custom_nodes/openvino_unet.py b/custom_nodes/openvino_unet.py new file mode 100644 index 000000000..f16dfb53f --- /dev/null +++ b/custom_nodes/openvino_unet.py @@ -0,0 +1,77 @@ +import torch + +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) + +# Test OpenVINO Execution +from openvino.runtime import Core +from optimum.intel import OVStableDiffusionPipeline +from PIL import Image, ImageOps +from io import BytesIO +import numpy as np +import struct +import comfy.utils +import time +from pathlib import Path + +MODEL_ID = "helenai/stabilityai-stable-diffusion-2-1-ov" +MODEL_DIR = Path("diffusion_pipeline") +DEVICE="CPU" + +batch_size = 1 +num_images_per_prompt = 1 +height =512 +width = 512 + +class OpenVINOUNetInference: + def __init__(self): + if not MODEL_DIR.exists(): + self.pipe = OVStableDiffusionPipeline.from_pretrained( + MODEL_ID, + compile=False, + device=DEVICE, + ) + self.pipe.save_pretrained(MODEL_DIR) + else: + self.pipe = OVStableDiffusionPipeline.from_pretrained( + MODEL_ID, + compile=False, + device=DEVICE, + ) + self.pipe.reshape(batch_size=batch_size, height=height, width=width, num_images_per_prompt=num_images_per_prompt) + self.pipe.compile() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "prompt": ("STRING", {"tooltip": "The text prompt for generation."}), + "num_inference_steps": ("INT", {"tooltip": "The number of inference steps."}), + } + } + + RETURN_TYPES = ("IMAGE",) + OUTPUT_TOOLTIPS = ("The generated image from the prompt.",) + FUNCTION = "generate_image" + CATEGORY = "generation" + + def generate_image(self, prompt, num_inference_steps=50, guidance_scale=7.5): + """ + Generates an image from the given prompt using the OpenVINO optimized Stable Diffusion model. + """ + # Generate the image from the prompt using the pipeline + #image = self.pipe(prompt, num_inference_steps=num_inference_steps).images[0] + image = self.pipe(prompt, num_inference_steps=num_inference_steps) + final_image = image["images"][0] + return (final_image,) # Return the generated image as output + + +NODE_CLASS_MAPPINGS = { + "OpenVINOUNetInference": OpenVINOUNetInference, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "OpenVINOUnetInference": "OpenVINO Inference" +}