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

{ai_simple_engine_diffusion-0.0.1 → 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.1
+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.0,<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.1 → ai_simple_engine_diffusion-0.0.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-simple-engine-diffusion"
-version = "0.0.1"
+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.0,<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.3/src/ai_simple_engine_diffusion/graph/operation/generate_image.py

@@ -0,0 +1,67 @@
+from ai_simple_engine_diffusion.graph.operation.sample_latents import SampleLatents
+from ai_simple_engine_diffusion.graph.operation.generate_image_from_latents import GenerateImageFromLatents
+from ai_simple_engine.graph.operation.abstract.composite_operation import CompositeOperation
+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 INT, LOADED_MODEL, STRING, FLOAT, IMAGE
+
+
+class GenerateImage(
+    CompositeOperation
+):
+    """
+    *Composite Operation*
+
+    Generate an image from the given `prompt`,
+    `negative_prompt` and using the `model`.
+
+    Inputs:
+    - `model` (`LOADED_MODEL`)
+    - `prompt` (`STRING`)
+    - `negative_prompt` (`STRING`)
+    - `width` (`INT`)
+    - `height` (`INT`)
+    - `steps` (`INT`)
+    - `guidance_scale` (`FLOAT`)
+    - `seed` (`INT`)
+
+    Outputs:
+    - `image` (`IMAGE`)
+    """
+
+    model = Input(LOADED_MODEL)
+    prompt = Input(STRING)
+    negative_prompt = Input(
+        STRING,
+        default = ''
+    )
+    width = Input(INT)
+    height = Input(INT)
+    steps = Input(INT)
+    guidance_scale = Input(FLOAT)
+    seed = Input(INT)
+
+    image = Output(IMAGE)
+
+    def expand(
+        self
+    ):
+        latents = SampleLatents(
+            model = self.model,
+            width = self.width,
+            height = self.height,
+            seed = self.seed
+        )
+
+        image = GenerateImageFromLatents(
+            model = self.model,
+            latents = latents.latents,
+            prompt = self.prompt,
+            negative_prompt = self.negative_prompt,
+            steps = self.steps,
+            guidance_scale = self.guidance_scale
+        )
+
+        return {
+            'image': image.image
+        }

ai_simple_engine_diffusion-0.0.3/src/ai_simple_engine_diffusion/graph/operation/generate_image_from_latents.py

@@ -0,0 +1,63 @@
+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
+
+
+class GenerateImageFromLatents(
+    AtomicOperation
+):
+    """
+    *Atomic Operation*
+
+    Generate an image from an initial latent
+    tensor.
+
+    Inputs:
+    - `model` (`LOADED_MODEL`)
+    - `latents` (`LATENTS`)
+    - `prompt` (`STRING`)
+    - `negative_prompt` (`STRING`)
+    - `steps` (`INT`)
+    - `guidance_scale` (`FLOAT`)
+
+    Outputs:
+    - `image` (`IMAGE`)
+    """
+
+    model = Input(LOADED_MODEL)
+    latents = Input(LATENTS)
+    prompt = Input(STRING)
+    negative_prompt = Input(
+        STRING,
+        default = ''
+    )
+    steps = Input(INT)
+    guidance_scale = Input(FLOAT)
+
+    image = Output(IMAGE)
+
+    async def execute(
+        self,
+        context
+    ):
+        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.instance,
+            prompt = self.prompt,
+            negative_prompt = self.negative_prompt,
+            steps = self.steps,
+            guidance_scale = self.guidance_scale
+        )
+
+        return {
+            'image': image
+        }

ai_simple_engine_diffusion-0.0.3/src/ai_simple_engine_diffusion/graph/operation/sample_latents.py

@@ -0,0 +1,83 @@
+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
+from ai_simple_engine.graph.output import Output
+from ai_simple_engine.types.data_type.base import INT, LOADED_MODEL
+
+
+class SampleLatents(
+    AtomicOperation
+):
+    """
+    *Atomic Operation*
+    
+    Get `latents` sampled using the given
+    `model` and the instructions provided.
+
+    Inputs:
+    - `model` (`LOADED_MODEL`)
+    - `width` (`INT`)
+    - `height` (`INT`)
+    - `seed` (`INT`)
+
+    Outputs:
+    - `latents` (`LATENTS`)
+    """
+
+    model = Input(LOADED_MODEL)
+    width = Input(INT)
+    height = Input(INT)
+    seed = Input(INT)
+
+    latents = Output(LATENTS)
+
+    async def execute(
+        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(
+            width = self.width,
+            height = self.height
+        )
+
+        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.3/src/ai_simple_engine_diffusion/model/executor/abstract/latent_diffusion_model_executor_abstract.py

@@ -0,0 +1,172 @@
+from ai_simple_engine.models.executor.abstract import ModelExecutor
+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,
+    ABC
+):
+    """
+    *Abstract class*
+
+    Abstract executor for latent diffusion models.
+
+    Defines the common latent diffusion pipeline:
+
+    1. Encode the prompt.
+    2. Generate an initial latent tensor.
+    3. Perform the diffusion process.
+    4. Decode the final latents into an image.
+
+    Concrete implementations only need to provide the
+    model-specific operations, while the diffusion loop
+    is shared.
+
+    Models that should use this executor:
+    - Stable Diffusion 1.x
+    - SD 2.x
+    - SDXL
+    - FLUX
+    - SD3
+    - AuraFlow
+    - HunyuanDiT
+    - PixArt-α
+    - Sana
+    - Wan (imagen)
+    """
+
+    @abstractmethod
+    def create_scheduler(
+        self,
+        context: ExecutionContext,
+        model: LoadedModel
+    ):
+        ...
+
+    @abstractmethod
+    async def predict_noise(
+        self,
+        *,
+        model: LoadedModel,
+        latents,
+        embeddings,
+        timestep,
+        guidance_scale
+    ):
+        ...
+
+    @abstractmethod
+    async def encode_prompt(
+        self,
+        *,
+        model: LoadedModel,
+        prompt: str,
+        negative_prompt: str
+    ):
+        ...
+
+    @abstractmethod
+    async def decode_latents(
+        self,
+        *,
+        model: LoadedModel,
+        latents
+    ):
+        ...
+
+    async def generate(
+        self,
+        *,
+        context = ExecutionContext,
+        model: LoadedModel,
+        latents,
+        prompt: str,
+        negative_prompt: str,
+        steps: int,
+        guidance_scale: float
+    ):
+        """
+        Generate an image with the given `prompt`,
+        `negative_prompt` and `model`.
+        """
+        
+        """
+        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
+            )
+
+            image = await self.decode_latents(
+                model = model,
+                latents = latents
+            )
+
+        return image
+
+    async def _diffuse(
+        self,
+        *,
+        context: ExecutionContext,
+        model: LoadedModel,
+        latents,
+        embeddings,
+        steps: int,
+        guidance_scale: float
+    ):
+        scheduler = self.create_scheduler(
+            context = context,
+            model = model
+        )
+
+        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,
+                timestep
+            )
+
+            noise = await self.predict_noise(
+                model = model,
+                latents = latent_input,
+                embeddings = embeddings,
+                timestep = timestep,
+                guidance_scale = guidance_scale
+            )
+
+            step = scheduler.step(
+                model_output = noise,
+                timestep = timestep,
+                sample = latents
+            )
+
+            latents = step.prev_sample
+
+        return latents

{ai_simple_engine_diffusion-0.0.1 → ai_simple_engine_diffusion-0.0.3}/src/ai_simple_engine_diffusion/model/info/diffusion_model_info.py

@@ -25,6 +25,7 @@ class DiffusionModelInfo(
         height: int
     ) -> LatentSpace:
         return LatentSpace(
+            batch_size = 1,
             width = width // self.vae_scale_factor,
             height = height // self.vae_scale_factor,
             channels = self.latent_channels

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.1 → ai_simple_engine_diffusion-0.0.3}/src/ai_simple_engine_diffusion/scheduler/spec/base.py

@@ -9,7 +9,7 @@ class SchedulerSpec:
     be provided to the scheduler when creating it.
     """
 
-    name: str
+    identifier: str
     kwargs: dict[str, object] = field(default_factory = dict)
 
 
@@ -17,7 +17,7 @@ class SchedulerSpec:
 Example:
 
 SchedulerSpec(
-    name="euler",
+    identifier="euler",
     kwargs={
         "use_karras_sigmas": True
     }

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

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

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/graph/operation/decode_latents.py

@@ -1,46 +0,0 @@
-from ai_simple_engine_diffusion.types.data_type import LATENTS
-from ai_simple_engine.models.executor.registry.family_model_executor_registry import FamilyModelExecutorRegistry
-from ai_simple_engine.graph.operation.abstract.atomic_operation import AtomicOperation
-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 IMAGE, LOADED_MODEL
-
-
-class DecodeLatents(
-    AtomicOperation
-):
-    """
-    *Atomic Operation*
-    
-    Decode the `latents` provided with the `model`
-    given, obtaining an `image`.
-
-    Inputs:
-    - `model` (`LOADED_MODEL`)
-    - `latents` (`LATENTS`)
-
-    Outputs:
-    - `image` (`IMAGE`)
-    """
-
-    model = Input(LOADED_MODEL)
-    latents = Input(LATENTS)
-
-    image = Output(IMAGE)
-
-    async def execute(
-        self,
-        context
-    ):
-        registry = context.services.get(FamilyModelExecutorRegistry)
-
-        executor = registry.resolve(self.model)
-
-        image = await executor.decode(
-            model = self.model,
-            latents = self.latents
-        )
-
-        return {
-            'image': image
-        }

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/graph/operation/diffusion_loop.py

@@ -1,57 +0,0 @@
-from ai_simple_engine_diffusion.types.data_type import EMBEDDINGS, LATENTS
-from ai_simple_engine.graph.operation.abstract.atomic_operation import AtomicOperation
-from ai_simple_engine.types.data_type.base import INT, FLOAT, LOADED_MODEL
-from ai_simple_engine.graph.input import Input
-from ai_simple_engine.graph.output import Output
-
-
-class DiffusionLoop(
-    AtomicOperation
-):
-    """
-    *Atomic Operation*
-
-    Performs the complete denoising process,
-    iteratively updating the latent tensor.
-
-    This operation encapsulates the whole diffusion
-    loop. Internally it repeatedly asks the model
-    to predict the noise and delegates the latent
-    update to the scheduler.
-
-    Inputs:
-    - `model` (`LOADED_MODEL`)
-    - `latents` (`LATENTS`)
-    - `embeddings` (`EMBEDDINGS`)
-    - `steps` (`INT)
-    - `guidance_scale` (`FLOAT`)
-
-    Outputs:
-    - `latents_output` (`LATENTS`)
-    """
-
-    model = Input(LOADED_MODEL)
-    latents = Input(LATENTS)
-    embeddings = Input(EMBEDDINGS)
-    steps = Input(INT)
-    guidance_scale = Input(FLOAT)
-
-    latents_output = Output(LATENTS)
-
-    async def execute(
-        self,
-        context
-    ):
-        executor = context.model_executors.resolve(self.model)
-
-        latents = await executor.diffuse(
-            model = self.model,
-            latents = self.latents,
-            embeddings = self.embeddings,
-            steps = self.steps,
-            guidance_scale = self.guidance_scale
-        )
-
-        return {
-            'latents': latents
-        }

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/graph/operation/encode_prompt.py

@@ -1,46 +0,0 @@
-from ai_simple_engine_diffusion.types.data_type import EMBEDDINGS
-from ai_simple_engine.types.data_type.base import LOADED_MODEL, STRING
-from ai_simple_engine.models.executor.registry.family_model_executor_registry import FamilyModelExecutorRegistry
-from ai_simple_engine.graph.operation.abstract.atomic_operation import AtomicOperation
-from ai_simple_engine.graph.input import Input
-from ai_simple_engine.graph.output import Output
-
-
-class EncodePrompt(
-    AtomicOperation
-):
-    """
-    *Atomic Operation*
-    
-    Encode the `prompt` given with the `text_encoder`
-    also provided.
-
-    Inputs:
-    - `text_encoder` (`LOADED_MODEL`)
-    - `prompt` (`STRING`)
-
-    Outputs:
-    - `embeddings` (`EMBEDDINGS`)
-    """
-
-    text_encoder = Input(LOADED_MODEL)
-    prompt = Input(STRING)
-
-    embeddings = Output(EMBEDDINGS)
-
-    async def execute(
-        self,
-        context
-    ):
-        registry = context.services.get(FamilyModelExecutorRegistry)
-
-        runner = registry.resolve(self.text_encoder)
-
-        embeddings = await runner.encode(
-            model = self.text_encoder,
-            prompt = self.prompt
-        )
-
-        return {
-            'embeddings': embeddings
-        }

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/graph/operation/generate_latents.py

@@ -1,74 +0,0 @@
-from ai_simple_engine_diffusion.types.data_type import LATENTS
-from ai_simple_engine_diffusion.graph.operation.encode_prompt import EncodePrompt
-from ai_simple_engine_diffusion.graph.operation.sample_latents import SampleLatents
-from ai_simple_engine_diffusion.graph.operation.diffusion_loop import DiffusionLoop
-from ai_simple_engine.graph.operation.abstract.composite_operation import CompositeOperation
-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 INTEGER, LOADED_MODEL, STRING, FLOAT
-
-
-class GenerateLatents(
-    CompositeOperation
-):
-    """
-    *Composite Operation*
-    
-    Generate `latents` for the given `prompt`
-    and instructions.
-
-    Inputs:
-    - `model` (`LOADED_MODEL`)
-    - `prompt` (`STRING`)
-    - `negative_prompt` (`STRING`)
-    - `width` (`INTEGER`)
-    - `height` (`INTEGER`)
-    - `steps` (`INTEGER`)
-    - `guidance_scale` (`FLOAT`)
-    - `seed` (`INTEGER`)
-
-    Outputs:
-    - `latents` (`LATENTS`)
-    """
-
-    model = Input(LOADED_MODEL)
-    prompt = Input(STRING)
-    negative_prompt = Input(
-        STRING,
-        default = ''
-    )
-    width = Input(INTEGER)
-    height = Input(INTEGER)
-    steps = Input(INTEGER)
-    guidance_scale = Input(FLOAT)
-    seed = Input(INTEGER)
-
-    latents = Output(LATENTS)
-
-    def expand(
-        self
-    ):
-        embeddings = EncodePrompt(
-            model = self.model,
-            prompt = self.prompt,
-            negative_prompt = self.negative_prompt
-        )
-
-        sampled_latents = SampleLatents(
-            model = self.model,
-            width = self.width,
-            height = self.height,
-            seed = self.seed
-        )
-
-        latents = DiffusionLoop(
-            model = self.model,
-            latents = sampled_latents.latents,
-            embeddings = embeddings.embeddings,
-            steps = self.steps,
-            guidance_scale = self.guidance_scale
-        )
-
-        return {
-            'latents': latents.latents
-        }

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/graph/operation/predict_noise.py

@@ -1,53 +0,0 @@
-from ai_simple_engine_diffusion.types.data_type import EMBEDDINGS, LATENTS, NOISE_PREDICTION
-from ai_simple_engine.models.executor.registry.family_model_executor_registry import FamilyModelExecutorRegistry
-from ai_simple_engine.types.data_type.base import LOADED_MODEL, INT
-from ai_simple_engine.graph.operation.abstract.atomic_operation import AtomicOperation
-from ai_simple_engine.graph.input import Input
-from ai_simple_engine.graph.output import Output
-
-
-class PredictNoise(
-    AtomicOperation
-):
-    """
-    *Atomic Operation*
-    
-    Predict the `noise` from the given
-    `latents` and `embeddings`, using the
-    `model` provided.
-
-    Inputs:
-    - `model` (`LOADED_MODEL`)
-    - `latents` (`LATENTS`)
-    - `embeddings` (`EMBEDDINGS`)
-    - `timestep` (`INT`)
-
-    Outputs:
-    - `noise` (`NOISE_PREDICTION`)
-    """
-    
-    model = Input(LOADED_MODEL)
-    latents = Input(LATENTS)
-    embeddings = Input(EMBEDDINGS)
-    timestep = Input(INT)
-
-    noise = Output(NOISE_PREDICTION)
-
-    async def execute(
-        self,
-        context
-    ):
-        registry = context.services.get(FamilyModelExecutorRegistry)
-
-        executor = registry.resolve(self.model)
-
-        noise = await executor.predict_noise(
-            model = self.model,
-            latents = self.latents,
-            embeddings = self.embeddings,
-            timestep = self.timestep
-        )
-
-        return {
-            'noise': noise
-        }

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/graph/operation/sample_latents.py

@@ -1,53 +0,0 @@
-from ai_simple_engine_diffusion.types.data_type 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
-from ai_simple_engine.graph.output import Output
-from ai_simple_engine.types.data_type.base import INTEGER, LOADED_MODEL
-
-
-class SampleLatents(
-    AtomicOperation
-):
-    """
-    *Atomic Operation*
-    
-    Get `latents` sampled using the given
-    `model` and the instructions provided.
-
-    Inputs:
-    - `model` (`LOADED_MODEL`)
-    - `width` (`INTEGER`)
-    - `height` (`INTEGER`)
-    - `seed` (`INTEGER`)
-
-    Outputs:
-    - `latents` (`LATENTS`)
-    """
-
-    model = Input(LOADED_MODEL)
-    width = Input(INTEGER)
-    height = Input(INTEGER)
-    seed = Input(INTEGER)
-
-    latents = Output(LATENTS)
-
-    async def execute(
-        self,
-        context
-    ):
-        generator = context.services.get(NoiseGenerator)
-
-        space = self.model.info.latent_space(
-            width = self.width,
-            height = self.height
-        )
-
-        latents = await generator.generate(
-            space = space,
-            seed = self.seed
-        )
-
-        return {
-            'latents': latents
-        }

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/model/executor/abstract/diffusion_model_executor_abstract.py

@@ -1,66 +0,0 @@
-from ai_simple_engine_diffusion.model.info.diffusion_model_info import DiffusionModelInfo
-from ai_simple_engine.models.executor.abstract import ModelExecutor
-from ai_simple_engine.models.loaded_model import LoadedModel
-from abc import ABC, abstractmethod
-
-
-class DiffusionModelExecutorAbstract(
-    ModelExecutor,
-    ABC
-):
-
-    @abstractmethod
-    def info(
-        self,
-        model: LoadedModel
-    ) -> DiffusionModelInfo:
-        ...
-
-    @abstractmethod
-    def create_scheduler(
-        self,
-        model: LoadedModel
-    ):
-        ...
-
-    @abstractmethod
-    async def predict_noise(
-        self,
-        *,
-        model: LoadedModel,
-        latents,
-        embeddings,
-        timestep,
-        guidance_scale
-    ):
-        ...
-
-    async def diffuse(
-        self,
-        *,
-        model: LoadedModel,
-        latents,
-        embeddings,
-        steps: int,
-        guidance_scale: float
-    ):
-        scheduler = self.create_scheduler(model)
-
-        scheduler.set_timesteps(steps)
-
-        for timestep in scheduler.timesteps:
-            noise = await self.predict_noise(
-                model = model,
-                latents = latents,
-                embeddings = embeddings,
-                timestep = timestep,
-                guidance_scale = guidance_scale
-            )
-
-            latents = scheduler.step(
-                noise_prediction = noise,
-                latents = latents,
-                timestep = timestep
-            )
-
-        return latents

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/model/executor/diffusers_stable_diffusion_model_executor.py

@@ -1,43 +0,0 @@
-from ai_simple_engine_diffusion.model.executor.abstract.diffusion_model_executor_abstract import DiffusionModelExecutorAbstract
-from ai_simple_engine.models.loaded_model import LoadedModel
-
-
-class DiffusersStableDiffusionModelExecutor(
-    DiffusionModelExecutorAbstract
-):
-
-    def create_scheduler(
-        self,
-        model: LoadedModel
-    ):
-        # TODO: Move this
-        from ai_simple_engine_diffusion.scheduler.utils import SCHEDULERS
-
-        scheduler_cls = SCHEDULERS[model.info.scheduler]
-
-        return scheduler_cls.from_config(
-            model.model.scheduler.config
-        )
-
-    async def predict_noise(
-        self,
-        *,
-        model: LoadedModel,
-        latents,
-        embeddings,
-        timestep,
-        guidance_scale
-    ):
-        pipeline = model.model
-
-        noise_prediction = pipeline.unet(
-            latents,
-            timestep,
-            encoder_hidden_states = embeddings
-        ).sample
-
-        if guidance_scale > 1:
-            # TODO: Apply CFG (?)
-            pass
-
-        return noise_prediction

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/model/executor/text_encoder_model_executor_abstract.py

@@ -1,27 +0,0 @@
-from ai_simple_engine.models.executor.abstract import ModelExecutor
-from ai_simple_engine.models.loaded_model import LoadedModel
-from ai_simple_engine_diffusion.types.data_type import Embeddings
-from abc import ABC, abstractmethod
-
-
-# This will be to perform 'EncodePrompt()'
-class TextEncoderModelExecutorAbstract(
-    ModelExecutor,
-    ABC
-):
-    """
-    *Abstract class*
-
-    To convert text into embeddings
-    """
-
-    @abstractmethod
-    async def encode(
-        self,
-        model: LoadedModel,
-        prompt: str
-    ) -> Embeddings:
-        """
-        Encode the given prompt into embeddings.
-        """
-        ...

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/model/executor/unet_model_executor_abstract.py

@@ -1,31 +0,0 @@
-from ai_simple_engine_diffusion.types.latents import Latents
-from ai_simple_engine_diffusion.types.noise_prediction import NoisePrediction
-from ai_simple_engine_diffusion.types.embeddings import Embeddings
-from ai_simple_engine.models.executor.abstract import ModelExecutor
-from ai_simple_engine.models.loaded_model import LoadedModel
-from abc import ABC, abstractmethod
-
-
-class UNetModelExecutorAbstract(
-    ModelExecutor,
-    ABC
-):
-    """
-    *Abstract class*
-
-    To predict the noise.
-    """
-
-    @abstractmethod
-    async def predict_noise(
-        self,
-        model: LoadedModel,
-        latents: Latents,
-        embeddings: Embeddings,
-        timestep: int
-    ) -> NoisePrediction:
-        """
-        Predict the noise residual for the given
-        latents at the specified timestep.
-        """
-        ...

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/model/executor/vae_model_executor_abstract.py

@@ -1,39 +0,0 @@
-from ai_simple_engine_diffusion.types.latents import Latents
-from ai_simple_engine.models.executor.abstract import ModelExecutor
-from ai_simple_engine.models.loaded_model import LoadedModel
-from ai_simple_engine.types.image import Image
-from abc import ABC, abstractmethod
-
-
-class VAEModelExecutorAbstract(
-    ModelExecutor,
-    ABC
-):
-    """
-    *Abstract class*
-
-    To transform images into latents and
-    latents into images.
-    """
-
-    @abstractmethod
-    async def encode(
-        self,
-        model: LoadedModel,
-        image: Image
-    ) -> Latents:
-        """
-        Encode an image into latent space.
-        """
-        ...
-
-    @abstractmethod
-    async def decode(
-        self,
-        model: LoadedModel,
-        latents: Latents
-    ) -> Image:
-        """
-        Decode latent representations into an image.
-        """
-        ...

ai_simple_engine_diffusion-0.0.1/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.
-        """
-        ...

ai_simple_engine_diffusion-0.0.1/src/ai_simple_engine_diffusion/scheduler/utils.py

@@ -1,20 +0,0 @@
-"""
-TODO: Maybe rename and move this module.
-"""
-SCHEDULERS = {
-    'ddim': DDIMScheduler,
-    'euler': EulerDiscreteScheduler,
-    'euler_a': EulerAncestralDiscreteScheduler,
-    'dpmpp_2m': DPMSolverMultistepScheduler,
-}
-
-
-def get_scheduler_class(
-    scheduler: str,
-    model: LoadedModel
-) -> object:
-    scheduler_cls = SCHEDULERS[model.info.scheduler]
-
-    return scheduler_cls.from_config(
-        model.model.scheduler.config
-    )