neuralnode 2.1.13__tar.gz → 2.1.14__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: neuralnode
-Version: 2.1.13
+Version: 2.1.14
 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.13 → neuralnode-2.1.14}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "neuralnode"
-version = "2.1.13"
+version = "2.1.14"
 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.13 → neuralnode-2.1.14}/src/neuralnode/providers/horus.py

@@ -385,6 +385,7 @@ class HorusProvider(BaseLLMProvider):
         self.model_id = model_id
         self.device = device or self._resolve_device()
         self.n_gpu_layers = n_gpu_layers
+        self._user_specified_dtype = torch_dtype  # track explicit user choice
         self.torch_dtype = self._resolve_torch_dtype(torch_dtype)
         self.load_in_8bit = load_in_8bit
         self.load_in_4bit = load_in_4bit
@@ -520,6 +521,24 @@ class HorusProvider(BaseLLMProvider):
             return mapping.get(torch_dtype.lower(), torch.float16)
         return torch_dtype
 
+    def _resolve_t2i_dtype(self) -> Any:
+        """Determine a safe dtype for text-to-image (Lens) models.
+
+        Flux/ZImage-class transformers overflow in float16 attention layers on GPUs
+        that lack native bfloat16 (e.g. Tesla T4, GTX 16xx).  This helper picks the
+        safest dtype that will actually produce a visible image:
+          * bfloat16 if hardware-supported (Ampere+)  – fast & numerically safe
+          * float32 otherwise (T4 / Turing)            – slower but correct
+        """
+        if torch is None:
+            return None
+        if self.device != "cuda" or not torch.cuda.is_available():
+            return torch.float32
+        if torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+        # T4 / Turing – float16 overflows in transformer attention, use float32
+        return torch.float32
+
     @staticmethod
     def _is_gguf_model_id(model_id: str) -> bool:
         return model_id.lower().endswith(".gguf")
@@ -640,8 +659,12 @@ class HorusProvider(BaseLLMProvider):
             pipeline_kwargs["revision"] = self.revision
         if self.variant:
             pipeline_kwargs["variant"] = self.variant
-        if self.torch_dtype is not None:
-            pipeline_kwargs["torch_dtype"] = self.torch_dtype
+        # Use a safe dtype for text-to-image; user-explicit overrides are honoured.
+        t2i_dtype = self.torch_dtype
+        if self._user_specified_dtype is None:
+            t2i_dtype = self._resolve_t2i_dtype()
+        if t2i_dtype is not None:
+            pipeline_kwargs["torch_dtype"] = t2i_dtype
 
         try:
             # Check if this model is Horus-Lens-1.0 which requires custom component mapping
@@ -701,13 +724,7 @@ class HorusProvider(BaseLLMProvider):
                 "Only public Hugging Face repositories are supported."
             ) from exc
 
-        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)
+        self._finalize_t2i_pipeline()
         return self
 
     def _load_gguf_text_to_image(self) -> "HorusProvider":
@@ -818,7 +835,9 @@ class HorusProvider(BaseLLMProvider):
             except ImportError:
                 GGUFQuantizationConfig = None
 
-        compute_dtype = self.torch_dtype or (torch.bfloat16 if self.device == "cuda" else torch.float32)
+        compute_dtype = self._resolve_t2i_dtype() if self._user_specified_dtype is None else self.torch_dtype
+        if compute_dtype is None:
+            compute_dtype = torch.float32
         
         transformer_kwargs = {
             "torch_dtype": compute_dtype,
@@ -875,8 +894,9 @@ class HorusProvider(BaseLLMProvider):
             pipeline_kwargs["revision"] = self.revision
         if self.variant:
             pipeline_kwargs["variant"] = self.variant
-        if self.torch_dtype is not None:
-            pipeline_kwargs["torch_dtype"] = self.torch_dtype
+        t2i_dtype = self._resolve_t2i_dtype() if self._user_specified_dtype is None else self.torch_dtype
+        if t2i_dtype is not None:
+            pipeline_kwargs["torch_dtype"] = t2i_dtype
 
         try:
             try:
@@ -928,14 +948,35 @@ class HorusProvider(BaseLLMProvider):
             with self._quiet_native_output():
                 self.model = DiffusionPipeline.from_pretrained(base_repo_id, **pipeline_kwargs)
 
+        self._finalize_t2i_pipeline()
+        return self
+
+    def _finalize_t2i_pipeline(self) -> None:
+        """Post-load hardening shared by both standard and GGUF Lens paths.
+
+        1. Move pipeline to self.device
+        2. Upcast VAE to float32 to prevent decoding NaN (essential on float16 pipelines)
+        3. Enable attention slicing for lower peak VRAM
+        4. Suppress progress bars if requested
+        """
+        if self.model is None:
+            return
         if hasattr(self.model, "to"):
             self.model = self.model.to(self.device)
+        # Always keep VAE in float32 — ZImage VAE overflows in fp16/bf16 decoding
         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)
+            self.model.vae.to(dtype=torch.float32)
+            # Tell the pipeline to upcast latents before passing them to the VAE
+            if hasattr(self.model.vae, "config"):
+                self.model.vae.config.force_upcast = True
+        # Reduce peak VRAM by slicing attention — negligible speed impact
+        if hasattr(self.model, "enable_attention_slicing"):
+            try:
+                self.model.enable_attention_slicing()
+            except Exception:
+                pass
         if hasattr(self.model, "set_progress_bar_config") and self.suppress_native_output:
             self.model.set_progress_bar_config(disable=True)
-        return self
 
     def _load_gguf(self) -> "HorusProvider":
         if not HF_HUB_AVAILABLE and self.auto_install_deps:
@@ -1567,6 +1608,21 @@ class HorusProvider(BaseLLMProvider):
         if not images:
             raise RuntimeError("Horus text-to-image generation did not return an image.")
         image = images[0]
+
+        # ── Black-image / NaN guard ─────────────────────────────────────
+        try:
+            import numpy as np
+            arr = np.asarray(image)
+            if arr.max() == 0:
+                logger.warning(
+                    "Horus Lens generated an all-black image. This usually means the "
+                    "transformer produced NaN values due to float16 overflow. "
+                    "Try loading with torch_dtype='bfloat16' (Ampere+ GPUs) "
+                    "or torch_dtype='float32' (Tesla T4 / GTX 16xx)."
+                )
+        except Exception:
+            pass
+
         if output_path:
             path = Path(output_path)
             path.parent.mkdir(parents=True, exist_ok=True)