ai-simple-engine-diffusion 0.0.1__py3-none-any.whl

ai_simple_engine_diffusion/__init__.py

@@ -0,0 +1,4 @@
+"""
+The plugin to add classes and utilities related
+to Diffusion.
+"""

ai_simple_engine_diffusion/consts.py

@@ -0,0 +1,8 @@
+from ai_simple_engine_common.consts.model_spec.families import TEXT_ENCODER_MODEL_FAMILY, VAE_MODEL_FAMILY, UNET_MODEL_FAMILY
+
+
+__all__ = [
+    'TEXT_ENCODER_MODEL_FAMILY',
+    'VAE_MODEL_FAMILY',
+    'UNET_MODEL_FAMILY'
+]

ai_simple_engine_diffusion/graph/operation/decode_latents.py

@@ -0,0 +1,46 @@
+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/graph/operation/diffusion_loop.py

@@ -0,0 +1,57 @@
+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/graph/operation/encode_image.py

@@ -0,0 +1,46 @@
+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 EncodeImage(
+    AtomicOperation
+):
+    """
+    *Atomic Operation*
+    
+    Encode the `image` provided with the `model`
+    given, obtaining the corresponding `latents`.
+
+    Inputs:
+    - `model` (`LOADED_MODEL`)
+    - `image` (`IMAGE`)
+
+    Outputs:
+    - `latents` (`LATENTS`)
+    """
+
+    model = Input(LOADED_MODEL)
+    image = Input(IMAGE)
+
+    latents = Output(LATENTS)
+
+    async def execute(
+        self,
+        context
+    ):
+        registry = context.services.get(FamilyModelExecutorRegistry)
+
+        executor = registry.resolve(self.model)
+
+        latents = await executor.encode(
+            model = self.model,
+            image = self.image
+        )
+
+        return {
+            'latents': latents
+        }

ai_simple_engine_diffusion/graph/operation/encode_prompt.py

@@ -0,0 +1,46 @@
+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/graph/operation/generate_latents.py

@@ -0,0 +1,74 @@
+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/graph/operation/predict_noise.py

@@ -0,0 +1,53 @@
+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/graph/operation/sample_latents.py

@@ -0,0 +1,53 @@
+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/model/executor/abstract/diffusion_model_executor_abstract.py

@@ -0,0 +1,66 @@
+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/model/executor/diffusers_stable_diffusion_model_executor.py

@@ -0,0 +1,43 @@
+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/model/executor/text_encoder_model_executor_abstract.py

@@ -0,0 +1,27 @@
+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/model/executor/unet_model_executor_abstract.py

@@ -0,0 +1,31 @@
+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/model/executor/vae_model_executor_abstract.py

@@ -0,0 +1,39 @@
+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/model/info/__init__.py

@@ -0,0 +1,4 @@
+
+
+# TODO: Create a 'ModelInfo' abstract class to
+# inherit for this (?)

ai_simple_engine_diffusion/model/info/diffusion_model_info.py

@@ -0,0 +1,36 @@
+from ai_simple_engine_diffusion.scheduler.spec.base import SchedulerSpec
+from ai_simple_engine_diffusion.types.latent_space import LatentSpace
+from ai_simple_engine.models.info.abstract import ModelInfo
+from dataclasses import dataclass
+# from typing import Union
+
+
+@dataclass(frozen = True)
+class DiffusionModelInfo(
+    ModelInfo
+):
+
+    latent_channels: int
+    vae_scale_factor: int
+    # latent_shape: tuple[int, ...] instead (?)
+    scheduler: SchedulerSpec
+    """
+    The scheduler that must be used for this
+    specific model: 'euler', 'flow_match', etc.
+    """
+
+    def latent_space(
+        self,
+        width: int,
+        height: int
+    ) -> LatentSpace:
+        return LatentSpace(
+            width = width // self.vae_scale_factor,
+            height = height // self.vae_scale_factor,
+            channels = self.latent_channels
+        )
+
+    # TODO: Near in the future
+    # supports_negative_prompt: bool
+    # max_resolution: Union[tuple[int, int], None]
+    # scheduler_defaults: ...

ai_simple_engine_diffusion/noise/noise_generator_abstract.py

@@ -0,0 +1,37 @@
+from ai_simple_engine_diffusion.types.latents import Latents
+from ai_simple_engine_diffusion.types.latent_space import LatentSpace
+from abc import ABC, abstractmethod
+
+
+class NoiseGenerator(
+    ABC
+):
+
+    @abstractmethod
+    async def generate(
+        self,
+        *,
+        space: LatentSpace,
+        seed: int
+    ) -> Latents:
+        """
+        Generate the initial latent noise.
+        """
+        ...
+
+
+"""
+This will be implemented as:
+class TorchNoiseGenerator(
+    NoiseGenerator
+):
+    ...
+
+in 'ai_simple_engine_runtime_diffusers'
+
+and will be registered like this:
+builder.add_service(
+    NoiseGenerator,
+    TorchNoiseGenerator()
+)
+"""

ai_simple_engine_diffusion/prompt_embeddings/abstract.py

@@ -0,0 +1,3 @@
+class PromptEmbeddings:
+    
+    pass

ai_simple_engine_diffusion/scheduler/abstract.py

@@ -0,0 +1,46 @@
+"""
+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/scheduler/spec/base.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen = True)
+class SchedulerSpec:
+    """
+    The specifications of the scheduler that must
+    be used, which means the arguments that must
+    be provided to the scheduler when creating it.
+    """
+
+    name: str
+    kwargs: dict[str, object] = field(default_factory = dict)
+
+
+"""
+Example:
+
+SchedulerSpec(
+    name="euler",
+    kwargs={
+        "use_karras_sigmas": True
+    }
+)
+"""

ai_simple_engine_diffusion/scheduler/utils.py

@@ -0,0 +1,20 @@
+"""
+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
+    )

ai_simple_engine_diffusion/types/data_type.py

@@ -0,0 +1,28 @@
+from ai_simple_engine_diffusion.types.latents import Latents
+from ai_simple_engine_diffusion.types.embeddings import Embeddings
+from ai_simple_engine_diffusion.types.noise_prediction import NoisePrediction
+from ai_simple_engine.types.data_type.base import DataType
+
+
+LATENTS = DataType(
+    name = 'Latents',
+    runtime_type = Latents,
+    parent = None
+)
+EMBEDDINGS = DataType(
+    name = 'Embeddings',
+    runtime_type = Embeddings,
+    parent = None
+)
+NOISE_PREDICTION = DataType(
+    name = 'NoisePrediction',
+    runtime_type = NoisePrediction,
+    parent = None
+)
+
+
+__all__ = [
+    'LATENTS',
+    'EMBEDDINGS',
+    'NOISE_PREDICTION'
+]

ai_simple_engine_diffusion/types/embeddings.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen = True)
+class Embeddings:
+
+    instance: object
+
+
+"""
+I don't use 'torch.Tensor' here because this
+one is general and the implementation could
+be 'torch.Tensor' or could be not.
+
+Having a specific 'Latents' class allows us
+to avoid receiving 'Embeddings' in a method
+that is not expected, even if the 'Latents'
+and the 'Embeddings' have a 'torch' value
+inside. They are distinct concepts. same with
+'NoisePrediction'.
+"""

ai_simple_engine_diffusion/types/latent_space.py

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

ai_simple_engine_diffusion/types/latents.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen = True)
+class Latents:
+    
+    instance: object
+
+
+"""
+I don't use 'torch.Tensor' here because this
+one is general and the implementation could
+be 'torch.Tensor' or could be not.
+
+Having a specific 'Latents' class allows us
+to avoid receiving 'Embeddings' in a method
+that is not expected, even if the 'Latents'
+and the 'Embeddings' have a 'torch' value
+inside. They are distinct concepts. same with
+'NoisePrediction'.
+"""

ai_simple_engine_diffusion/types/noise_prediction.py

@@ -0,0 +1,21 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen = True)
+class NoisePrediction:
+
+    instance: object
+
+
+"""
+I don't use 'torch.Tensor' here because this
+one is general and the implementation could
+be 'torch.Tensor' or could be not.
+
+Having a specific 'Latents' class allows us
+to avoid receiving 'Embeddings' in a method
+that is not expected, even if the 'Latents'
+and the 'Embeddings' have a 'torch' value
+inside. They are distinct concepts. same with
+'NoisePrediction'.
+"""

ai_simple_engine_diffusion-0.0.1.dist-info/METADATA

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: ai-simple-engine-diffusion
+Version: 0.0.1
+Summary: AI Simple Engine Diffusion Module
+License-File: LICENSE
+Author: danialcala94
+Author-email: danielalcalavalera@gmail.com
+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-common (>=0.0.3,<1.0.0)
+Description-Content-Type: text/markdown
+
+# AI Simple Engine Diffusion Module
+
+The module that implements the abstract classes and utilities to be able to do Diffusion (image creation), for the AI Simple Engine `ai-simple-engine`.

ai_simple_engine_diffusion-0.0.1.dist-info/RECORD

@@ -0,0 +1,40 @@
+ai_simple_engine_diffusion/__init__.py,sha256=Ornd59rNW32VowUSiT_7QnuryUp9NMREBiw34ms7qUU,72
+ai_simple_engine_diffusion/consts.py,sha256=nxSeOtv3Y_ezmMoPAdjCjzZg533iqaJVu8jTerU2QRE,229
+ai_simple_engine_diffusion/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/graph/operation/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/graph/operation/decode_latents.py,sha256=y69ivvQnqy8VigNyYJYZyJpwyN7vq0KgjAqkpAP7BZ8,1197
+ai_simple_engine_diffusion/graph/operation/diffusion_loop.py,sha256=81FxdgzROaN7_i9T6N08IgVlhC1ayypVifB1PxW0CLc,1573
+ai_simple_engine_diffusion/graph/operation/encode_image.py,sha256=0C46darLjm3erFgR3HNPsXeoPH3xKMxJbYEtwJmCc7E,1212
+ai_simple_engine_diffusion/graph/operation/encode_prompt.py,sha256=d5QctkPhIkbzmlzHLFbHIKPwAgwWvrTPH1vQ58_HKLc,1242
+ai_simple_engine_diffusion/graph/operation/generate_latents.py,sha256=dl-3MeDHct6QrXQL91xg67-RElnMoayBCTLrnYoLQa4,2113
+ai_simple_engine_diffusion/graph/operation/predict_noise.py,sha256=cfykfp9GFHFSpkY-bIl7BGD5tDN7z-rsVNvJdQ7SdKA,1483
+ai_simple_engine_diffusion/graph/operation/sample_latents.py,sha256=pmeWCz7Jn7T8w77uChXUoDcOV4VQpXHSjo1eXd0st7I,1348
+ai_simple_engine_diffusion/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/model/executor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/model/executor/abstract/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/model/executor/abstract/diffusion_model_executor_abstract.py,sha256=lb50F6hyKXLMpIlWoQBNUX-LwMCSPyRhsQQy4QIAKrY,1561
+ai_simple_engine_diffusion/model/executor/diffusers_stable_diffusion_model_executor.py,sha256=zqoOlVapev-hvVo_zuHwHp9G3zB6G1MoQK2uMAqKDuw,1107
+ai_simple_engine_diffusion/model/executor/text_encoder_model_executor_abstract.py,sha256=78SLwFMIQvnYhSpe9bNS5eVMaNGYHVimCiphaNSoPYU,656
+ai_simple_engine_diffusion/model/executor/unet_model_executor_abstract.py,sha256=7pIgudVQHpsKUSWRmlcDtAbYM_zpSeIA2rw4Q4Niz5E,855
+ai_simple_engine_diffusion/model/executor/vae_model_executor_abstract.py,sha256=VmAmEiSfH1Mk9vwp2ZcQtWYA3-NQRshRtnNRCCGHEr8,896
+ai_simple_engine_diffusion/model/info/__init__.py,sha256=OwlQ20tWdB11fWKzDnPlSXr1Es2LJ-JQ_VdBq10xCrA,74
+ai_simple_engine_diffusion/model/info/diffusion_model_info.py,sha256=ITx8dbRCWqTDARKBKldCO9XwwwR0eJKaU5suidALTsc,1044
+ai_simple_engine_diffusion/noise/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/noise/noise_generator_abstract.py,sha256=WlSvbjgNUUtueYUjQ9pAQ4bP2VrUS5ZAmAF7U4VxRE0,691
+ai_simple_engine_diffusion/prompt_embeddings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/prompt_embeddings/abstract.py,sha256=BqhDxBpioPa2ZwIdCapmgqQeOozk-1Msf5sOlp90MP8,39
+ai_simple_engine_diffusion/scheduler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/scheduler/abstract.py,sha256=BkE46-rRKVa3bpEhSpV3hwiA_0m4pp5Aw5JUPeiMV-M,1007
+ai_simple_engine_diffusion/scheduler/spec/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/scheduler/spec/base.py,sha256=8IswyXklrmmUrjfGkQKTtz4HFQ27NMoA9mnZrp3hZNs,464
+ai_simple_engine_diffusion/scheduler/utils.py,sha256=d0ToobCjOxjfP8V555WRW8ctQz8N7M3naXTl4fNhcfU,459
+ai_simple_engine_diffusion/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ai_simple_engine_diffusion/types/data_type.py,sha256=_30TMYvyOXlzFAzV4RyGKKTv10_pg8LmBxMoiRzCgGg,669
+ai_simple_engine_diffusion/types/embeddings.py,sha256=dSr_o3h_owAAkjao8fHujQg4rF5_3W_BUCB_TB5R0lY,495
+ai_simple_engine_diffusion/types/latent_space.py,sha256=tRtT1afGxc0H8_J_irW8NLcvBYjgQT12JPV-hKfKwiI,138
+ai_simple_engine_diffusion/types/latents.py,sha256=KobTwk-LlJ1ftFuiuy1T-LgwWMGSY5hqU0vhHQ5uzt8,496
+ai_simple_engine_diffusion/types/noise_prediction.py,sha256=10FxCykTxPbzV5CcJj9LvaIEBiyLFCH2tR3nRP0xP60,500
+ai_simple_engine_diffusion-0.0.1.dist-info/licenses/LICENSE,sha256=6kbiFSfobTZ7beWiKnHpN902HgBx-Jzgcme0SvKqhKY,1091
+ai_simple_engine_diffusion-0.0.1.dist-info/METADATA,sha256=3xjsj58ohnrVjYmzI_M2gKIhqw_hngX-xtR8EAWfcnc,714
+ai_simple_engine_diffusion-0.0.1.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+ai_simple_engine_diffusion-0.0.1.dist-info/RECORD,,

ai_simple_engine_diffusion-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.2.0
+Root-Is-Purelib: true
+Tag: py3-none-any

ai_simple_engine_diffusion-0.0.1.dist-info/licenses/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2018 The Python Packaging Authority
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.