ai-simple-engine-diffusion 0.0.2__tar.gz → 0.0.3__tar.gz

{ai_simple_engine_diffusion-0.0.2 → ai_simple_engine_diffusion-0.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-simple-engine-diffusion
-Version: 0.0.2
+Version: 0.0.3
 Summary: AI Simple Engine Diffusion Module
 License-File: LICENSE
 Author: danialcala94
@@ -9,7 +9,7 @@ Requires-Python: >=3.10,<3.12
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: ai-simple-engine (>=0.2.6,<1.0.0)
+Requires-Dist: ai-simple-engine (>=0.2.8,<1.0.0)
 Requires-Dist: ai-simple-engine-common (>=0.0.3,<1.0.0)
 Description-Content-Type: text/markdown
 

{ai_simple_engine_diffusion-0.0.2 → ai_simple_engine_diffusion-0.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-simple-engine-diffusion"
-version = "0.0.2"
+version = "0.0.3"
 description = "AI Simple Engine Diffusion Module"
 authors = [
     {name = "danialcala94", email = "danielalcalavalera@gmail.com"}
@@ -10,7 +10,7 @@ requires-python = ">=3.10,<3.12"
 
 [tool.poetry.dependencies]
 # Mandatory
-ai-simple-engine = { version = ">=0.2.6,<1.0.0", optional = false }
+ai-simple-engine = { version = ">=0.2.8,<1.0.0", optional = false }
 ai-simple-engine-common = { version = ">=0.0.3,<1.0.0", optional = false }
 # Optional
 

{ai_simple_engine_diffusion-0.0.2 → ai_simple_engine_diffusion-0.0.3}/src/ai_simple_engine_diffusion/graph/operation/generate_image_from_latents.py

@@ -1,5 +1,6 @@
 from ai_simple_engine_diffusion.types.data_type import LATENTS
 from ai_simple_engine.graph.operation.abstract.atomic_operation import AtomicOperation
+from ai_simple_engine.models.executor.registry.family_model_executor_registry import FamilyModelExecutorRegistry
 from ai_simple_engine.graph.input import Input
 from ai_simple_engine.graph.output import Output
 from ai_simple_engine.types.data_type.base import LOADED_MODEL, STRING, INT, FLOAT, IMAGE
@@ -42,11 +43,15 @@ class GenerateImageFromLatents(
         self,
         context
     ):
-        executor = context.model_executors.resolve(self.model)
+        model_executors_registry = context.services.get(FamilyModelExecutorRegistry)
+        executor = model_executors_registry.resolve(self.model)
+
+        # executor = context.model_executors.resolve(self.model)
 
         image = await executor.generate(
+            context = context,
             model = self.model,
-            latents = self.latents,
+            latents = self.latents.instance,
             prompt = self.prompt,
             negative_prompt = self.negative_prompt,
             steps = self.steps,

{ai_simple_engine_diffusion-0.0.2 → ai_simple_engine_diffusion-0.0.3}/src/ai_simple_engine_diffusion/graph/operation/sample_latents.py

@@ -1,4 +1,6 @@
 from ai_simple_engine_diffusion.types.data_type import LATENTS
+# from ai_simple_engine_runtime_diffusers.models.executor.noise.pytorch_noise_generator import TorchNoiseGenerator
+from ai_simple_engine_diffusion.types.latents import Latents
 from ai_simple_engine_diffusion.noise.noise_generator_abstract import NoiseGenerator
 from ai_simple_engine.graph.operation.abstract.atomic_operation import AtomicOperation
 from ai_simple_engine.graph.input import Input
@@ -36,6 +38,13 @@ class SampleLatents(
         self,
         context
     ):
+        # noise_generator = context.services.get(TorchNoiseGenerator)
+        """
+        This plugin is registering it with the abstract
+        as the key, but including an instance of the
+        Torch version as the actual service registered.
+        """
+        # TODO: Is this ok like this (?)
         noise_generator = context.services.get(NoiseGenerator)
 
         space = self.model.info.latent_space(
@@ -43,11 +52,32 @@ class SampleLatents(
             height = self.height
         )
 
-        latents = await noise_generator.generate(
-            space = space,
+        shape = (
+            space.batch_size,
+            space.channels,
+            space.height,
+            space.width
+        )
+
+        runtime = self.model.instance
+
+        # Dynamically detect 'device' and 'dtype'
+        device = next(runtime.unet.parameters()).device
+        dtype = next(runtime.unet.parameters()).dtype
+
+        latents_torch = noise_generator.generate(
+            # space = space,
+            shape = shape,
+            device = device,
+            dtype = dtype,
             seed = self.seed
         )
 
+        latents = Latents(
+            latents_torch
+            # TODO: Maybe include the 'space' (?)
+        )
+
         return {
             'latents': latents
         }

{ai_simple_engine_diffusion-0.0.2 → ai_simple_engine_diffusion-0.0.3}/src/ai_simple_engine_diffusion/model/executor/abstract/latent_diffusion_model_executor_abstract.py

@@ -3,6 +3,8 @@ from ai_simple_engine.models.loaded_model import LoadedModel
 from ai_simple_engine.execution.execution_context import ExecutionContext
 from abc import ABC, abstractmethod
 
+import torch
+
 
 class LatentDiffusionModelExecutorAbstract(
     ModelExecutor,
@@ -91,25 +93,32 @@ class LatentDiffusionModelExecutorAbstract(
         Generate an image with the given `prompt`,
         `negative_prompt` and `model`.
         """
-        embeddings = await self.encode_prompt(
-            model = model,
-            prompt = prompt,
-            negative_prompt = negative_prompt
-        )
+        
+        """
+        Disable view tracking and version counter
+        bumps to improve the performance as it is
+        inference process.
+        """
+        with torch.inference_mode():
+            embeddings = await self.encode_prompt(
+                model = model,
+                prompt = prompt,
+                negative_prompt = negative_prompt
+            )
 
-        latents = await self._diffuse(
-            context = context,
-            model = model,
-            latents = latents,
-            embeddings = embeddings,
-            steps = steps,
-            guidance_scale = guidance_scale
-        )
+            latents = await self._diffuse(
+                context = context,
+                model = model,
+                latents = latents,
+                embeddings = embeddings,
+                steps = steps,
+                guidance_scale = guidance_scale
+            )
 
-        image = await self.decode_latents(
-            model = model,
-            latents = latents
-        )
+            image = await self.decode_latents(
+                model = model,
+                latents = latents
+            )
 
         return image
 
@@ -130,6 +139,14 @@ class LatentDiffusionModelExecutorAbstract(
 
         scheduler.set_timesteps(steps)
 
+        """
+        The noise that the `torch.randn(...)` generates
+        has a standard deviation of ~ 1, but each
+        scheduler expects to start with a different
+        amplitude, calculated with this formula below:
+        """
+        latents = latents * scheduler.init_noise_sigma
+
         for timestep in scheduler.timesteps:
             latent_input = scheduler.scale_model_input(
                 latents,

ai_simple_engine_diffusion-0.0.3/src/ai_simple_engine_diffusion/scheduler/abstract.py

@@ -0,0 +1,57 @@
+from ai_simple_engine_diffusion.scheduler.abstract import Scheduler
+
+
+class DiffusersScheduler(
+    Scheduler
+):
+
+    def __init__(
+        self,
+        scheduler
+    ):
+        self._scheduler = scheduler
+
+    @property
+    def init_noise_sigma(
+        self
+    ):
+        """
+        The `gamma` that must be used when starting
+        to create the noise.
+        """
+        return self._scheduler.init_noise_sigma
+
+    @property
+    def timesteps(
+        self
+    ):
+        return self._scheduler.timesteps
+
+    def set_timesteps(
+        self,
+        steps: int
+    ):
+        self._scheduler.set_timesteps(steps)
+
+    def scale_model_input(
+        self,
+        latents,
+        timestep
+    ):
+        return self._scheduler.scale_model_input(
+            latents,
+            timestep
+        )
+    
+    def step(
+        self,
+        *,
+        model_output,
+        timestep,
+        sample
+    ):
+        return self._scheduler.step(
+            model_output = model_output,
+            timestep = timestep,
+            sample = sample
+        )

{ai_simple_engine_diffusion-0.0.2 → ai_simple_engine_diffusion-0.0.3}/src/ai_simple_engine_diffusion/types/latent_space.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 @dataclass(frozen = True)
 class LatentSpace:
 
-    batch_size: int = 1
     channels: int
     width: int
-    height: int
+    height: int
+    batch_size: int = 1

ai_simple_engine_diffusion-0.0.2/src/ai_simple_engine_diffusion/scheduler/abstract.py

@@ -1,46 +0,0 @@
-"""
-TODO: Remove this class, it is no longer used.
-"""
-from ai_simple_engine_diffusion.types.latents import Latents
-from ai_simple_engine_diffusion.types.noise_prediction import NoisePrediction
-from abc import ABC, abstractmethod
-
-
-class Scheduler(
-    ABC
-):
-
-    @abstractmethod
-    def set_timesteps(
-        self,
-        *,
-        steps: int
-    ) -> None:
-        """
-        Prepare the scheduler for a diffusion run.
-        """
-        ...
-
-    @abstractmethod
-    def timesteps(
-        self
-    ) -> list[int]:
-        """
-        Return the timesteps used during inference.
-        """
-        ...
-
-    @abstractmethod
-    def step(
-        self,
-        *,
-        latents: Latents,
-        noise_prediction: NoisePrediction,
-        timestep: int
-    ) -> Latents:
-        """
-        Update the latent tensor using the predicted
-        noise for the current timestep, which means
-        computing the next latent sample.
-        """
-        ...