Add files via upload

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"
+}