From 2b3435a0ec9e79509508f6543e6c0129bc093dd0 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Sun, 23 Feb 2025 17:48:32 +0300
Subject: [PATCH] OpenVINO fix Img2Img

---
 CHANGELOG.md                       |  1 +
 modules/intel/openvino/__init__.py | 13 +++++++++++++
 2 files changed, 14 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1c955f2b2..570d5a03a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -21,6 +21,7 @@
   - log full config path, full log path, system name, extensions path
   - zluda hotfixes  
   - zluda force sync  
+  - openvino with img2img  
 
 ## Update for 2025-02-18
 
diff --git a/modules/intel/openvino/__init__.py b/modules/intel/openvino/__init__.py
index d6dd73567..0a73bc614 100644
--- a/modules/intel/openvino/__init__.py
+++ b/modules/intel/openvino/__init__.py
@@ -34,6 +34,19 @@ except Exception:
     pass
 
 
+# Fix "torch is not defined" error on sd 1.5 img2img when torch.compile for vae.encode is enabled
+from dataclasses import dataclass
+from torch.compiler import disable as disable_compile
+import diffusers.models.autoencoders.autoencoder_kl
+
+@dataclass
+@disable_compile
+class AutoencoderKLOutput(diffusers.utils.BaseOutput):
+    latent_dist: "DiagonalGaussianDistribution"  # noqa: F821
+diffusers.models.autoencoders.autoencoder_kl.AutoencoderKLOutput = AutoencoderKLOutput
+
+
+# Set default params
 torch._dynamo.config.cache_size_limit = 64 # pylint: disable=protected-access
 torch._dynamo.eval_frame.check_if_dynamo_supported = lambda: True # pylint: disable=protected-access
 if hasattr(torch._dynamo.config, "inline_inbuilt_nn_modules"):