neuralnode 2.1.12__tar.gz → 2.1.13__tar.gz

{neuralnode-2.1.12 → neuralnode-2.1.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neuralnode
-Version: 2.1.12
+Version: 2.1.13
 Summary: Comprehensive AI Framework with 50+ LLM Providers, Advanced Agents, Chains, Memory, RAG, and 100+ Tools
 Project-URL: Homepage, https://assem.cloud/
 Project-URL: Documentation, https://neuralnode.readthedocs.io

{neuralnode-2.1.12 → neuralnode-2.1.13}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "neuralnode"
-version = "2.1.12"
+version = "2.1.13"
 description = "Comprehensive AI Framework with 50+ LLM Providers, Advanced Agents, Chains, Memory, RAG, and 100+ Tools"
 readme = "README.md"
 requires-python = ">=3.9"

{neuralnode-2.1.12 → neuralnode-2.1.13}/src/neuralnode/providers/horus.py

@@ -504,6 +504,8 @@ class HorusProvider(BaseLLMProvider):
         if torch is None:
             return None
         if torch_dtype is None:
+            if self.device == "cuda" and torch.cuda.is_available() and torch.cuda.is_bf16_supported():
+                return torch.bfloat16
             return torch.float16 if self.device == "cuda" else torch.float32
         if isinstance(torch_dtype, str):
             mapping = {
@@ -701,6 +703,9 @@ class HorusProvider(BaseLLMProvider):
 
         if hasattr(self.model, "to"):
             self.model = self.model.to(self.device)
+        if hasattr(self.model, "vae") and self.model.vae is not None:
+            if self.device == "cuda" and self.torch_dtype == torch.float16:
+                self.model.vae.to(dtype=torch.float32)
         if hasattr(self.model, "set_progress_bar_config") and self.suppress_native_output:
             self.model.set_progress_bar_config(disable=True)
         return self
@@ -925,6 +930,9 @@ class HorusProvider(BaseLLMProvider):
 
         if hasattr(self.model, "to"):
             self.model = self.model.to(self.device)
+        if hasattr(self.model, "vae") and self.model.vae is not None:
+            if self.device == "cuda" and self.torch_dtype == torch.float16:
+                self.model.vae.to(dtype=torch.float32)
         if hasattr(self.model, "set_progress_bar_config") and self.suppress_native_output:
             self.model.set_progress_bar_config(disable=True)
         return self