ai-simple-engine-runtime-diffusers 0.0.2__tar.gz → 0.0.4__tar.gz

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-simple-engine-runtime-diffusers
-Version: 0.0.2
+Version: 0.0.4
 Summary: AI Simple Engine Runtime Diffusers Module
 License-File: LICENSE
 Author: danialcala94
@@ -9,11 +9,11 @@ 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.9,<1.0.0)
 Requires-Dist: ai-simple-engine-common (>=0.0.3,<1.0.0)
-Requires-Dist: ai-simple-engine-diffusion (>=0.0.2,<1.0.0)
+Requires-Dist: ai-simple-engine-diffusion (>=0.0.4,<1.0.0)
 Requires-Dist: diffusers (>=0.38.0,<9999.0.0)
-Requires-Dist: transformers (>=5.12.1,<9999.0.0)
+Requires-Dist: transformers (>=5.13.0,<9999.0.0)
 Description-Content-Type: text/markdown
 
 # AI Simple Engine Runtime Diffusers Module

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-simple-engine-runtime-diffusers"
-version = "0.0.2"
+version = "0.0.4"
 description = "AI Simple Engine Runtime Diffusers Module"
 authors = [
     {name = "danialcala94", email = "danielalcalavalera@gmail.com"}
@@ -10,10 +10,10 @@ 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.9,<1.0.0", optional = false }
 ai-simple-engine-common = { version = ">=0.0.3,<1.0.0", optional = false }
-ai-simple-engine-diffusion = { version = ">=0.0.2,<1.0.0", optional = false }
-transformers = { version = ">=5.12.1,<9999.0.0", optional = false }
+ai-simple-engine-diffusion = { version = ">=0.0.4,<1.0.0", optional = false }
+transformers = { version = ">=5.13.0,<9999.0.0", optional = false }
 diffusers = { version = ">=0.38.0,<9999.0.0", optional = false }
 # Optional
 

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/src/ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py

@@ -4,6 +4,7 @@ from ai_simple_engine.execution.execution_context import ExecutionContext
 from ai_simple_engine_runtime_diffusers.models.executor.prompt_embeddings.prompt_embeddings import PromptEmbeddings
 from ai_simple_engine_runtime_diffusers.models.scheduler.registry.diffusers_scheduler_registry import DiffusersSchedulerRegistry
 from ai_simple_engine.models.loaded_model import LoadedModel
+from ai_simple_engine.types.image import Image
 
 import torch
 
@@ -26,9 +27,9 @@ class DiffusersLatentDiffusionModelExecutor (
         context: ExecutionContext,
         model: LoadedModel,
     ):
-        scheduler_registry = context.services.resolve(DiffusersSchedulerRegistry)
+        scheduler_registry = context.services.get(DiffusersSchedulerRegistry)
 
-        scheduler_class = scheduler_registry.resolve(model.info.scheduler.name)
+        scheduler_class = scheduler_registry.resolve(model.info.scheduler.identifier)
 
         scheduler = scheduler_class.from_config(
             model.instance.scheduler_config,
@@ -109,15 +110,23 @@ class DiffusersLatentDiffusionModelExecutor (
     async def decode_latents(
         self,
         *,
-        model,
+        model: LoadedModel,
         latents
-    ):
-        runtime_model = model.model
+    ) -> Image:
+        runtime_model = model.instance
 
         vae = runtime_model.vae
 
         latents = latents / vae.config.scaling_factor
 
         image = vae.decode(latents).sample
+        """
+        TODO: I don't know if this must be done always
+        or depending or what, but I should pay attention
+        and refactor it.
+        """
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
+        image = Image(image)
 
         return image

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/src/ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py

@@ -9,12 +9,17 @@ from ai_simple_engine.models.loaded_model import LoadedModel
 from transformers import CLIPTokenizer, CLIPTextModel
 from diffusers import UNet2DConditionModel, AutoencoderKL
 
-import torch
-
 
 class LatentDiffusionModelLoader(
     ModelLoader
 ):
+    
+    @property
+    def family(
+        self
+    ) -> str:
+        # TODO: Transform into const
+        return 'latent_diffusion'
 
     async def load(
         self,
@@ -27,7 +32,7 @@ class LatentDiffusionModelLoader(
             path,
             subfolder = 'tokenizer',
             torch_dtype = get_torch_dtype_for(device)
-        ).to(str(device))
+        )#.to(str(device))
 
         text_encoder = CLIPTextModel.from_pretrained(
             path,
@@ -47,11 +52,25 @@ class LatentDiffusionModelLoader(
             torch_dtype = get_torch_dtype_for(device)
         ).to(str(device))
 
+        # TODO: This is if I want its own Scheduler
+        # DDIMScheduler.from_pretrained(
+        #     path,
+        #     subfolder = 'scheduler',
+        #     torch_dtype = get_torch_dtype_for(device)
+        # )#.to(str(device))
+
+        # TODO: We had this before but was failing
         scheduler_config = AutoencoderKL.load_config(
-            path,
+            path + '/scheduler/scheduler_config.json',
             subfolder = 'scheduler',
             torch_dtype = get_torch_dtype_for(device)
-        ).to(str(device))
+        )#.to(str(device))
+
+        # Free memory
+        # TODO: Explain why (?)
+        unet.eval()
+        text_encoder.eval()
+        vae.eval()
 
         runtime_model = DiffusersLatentDiffusionModel(
             tokenizer = tokenizer,

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/src/ai_simple_engine_runtime_diffusers/models/scheduler/diffusers_scheduler.py

@@ -1,9 +1,4 @@
-from ai_simple_engine_diffusion.scheduler.abstract import Scheduler
-
-
-class DiffusersScheduler(
-    Scheduler
-):
+class DiffusersScheduler:
 
     def __init__(
         self,
@@ -11,6 +6,16 @@ class DiffusersScheduler(
     ):
         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

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/src/ai_simple_engine_runtime_diffusers/models/scheduler/registry/diffusers_scheduler_registry.py

@@ -5,8 +5,10 @@ from diffusers.schedulers.scheduling_utils import SchedulerMixin
 class DiffusersSchedulerRegistry(
     ResolverRegistry[
         str,
-        type[SchedulerMixin]
-    ]
+        type[SchedulerMixin],
+        str
+    ],
+    # Generic[TModel, TExecutor]
 ):
     """
     Class that will register the different

ai_simple_engine_runtime_diffusers-0.0.4/src/ai_simple_engine_runtime_diffusers/models/scheduler/utils.py

@@ -0,0 +1,22 @@
+"""
+TODO: Maybe rename and move this module.
+"""
+from diffusers.schedulers import (
+    DDIMScheduler,
+    EulerDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    DPMSolverMultistepScheduler,
+    FlowMatchEulerDiscreteScheduler
+)
+
+SCHEDULERS = {
+    'ddim': DDIMScheduler,
+    'euler': EulerDiscreteScheduler,
+    'euler_a': EulerAncestralDiscreteScheduler,
+    'dpmpp_2m': DPMSolverMultistepScheduler,
+    'flow_match_euler': FlowMatchEulerDiscreteScheduler
+}
+"""
+The list including all the schedulers we are
+accepting.
+"""

{ai_simple_engine_runtime_diffusers-0.0.2 → ai_simple_engine_runtime_diffusers-0.0.4}/src/ai_simple_engine_runtime_diffusers/plugin.py

@@ -1,6 +1,10 @@
 from ai_simple_engine_runtime_diffusers.models.loader.latent_diffusion_model_loader import LatentDiffusionModelLoader
-from ai_simple_engine_runtime_diffusers.models.scheduler.utils import SCHEDULERS
+from ai_simple_engine_runtime_diffusers.models.executor.noise.pytorch_noise_generator import TorchNoiseGenerator
 from ai_simple_engine_runtime_diffusers.models.scheduler.registry.diffusers_scheduler_registry import DiffusersSchedulerRegistry
+from ai_simple_engine.models.executor.registry.family_model_executor_registry import FamilyModelExecutorRegistry
+from ai_simple_engine_runtime_diffusers.models.executor.diffusers_latent_diffusion_model_executor import DiffusersLatentDiffusionModelExecutor
+from ai_simple_engine_runtime_diffusers.models.scheduler.utils import SCHEDULERS
+from ai_simple_engine_diffusion.noise.noise_generator_abstract import NoiseGenerator
 from ai_simple_engine.engine_builder import EngineBuilder
 from ai_simple_engine.plugins.plugin import Plugin
 
@@ -20,7 +24,14 @@ class DiffusersRuntimePlugin(
         builder: EngineBuilder
     ):
         (
-            builder.add_model_loader(LatentDiffusionModelLoader)
+            builder.add_model_loader(LatentDiffusionModelLoader())
+        )
+
+        (
+            builder.add_service(
+                NoiseGenerator,
+                TorchNoiseGenerator()
+            )
         )
 
         # TODO: Register the SchedulerRegistry
@@ -34,6 +45,13 @@ class DiffusersRuntimePlugin(
             scheduler_spec.identifier
         )
         """
+        model_executor_registry = builder.get_or_add_service(FamilyModelExecutorRegistry)
+
+        model_executor_registry.register(
+            'latent_diffusion',
+            DiffusersLatentDiffusionModelExecutor()
+        )
+
         scheduler_registry = builder.get_or_add_service(DiffusersSchedulerRegistry)
 
         for identifier, cls in SCHEDULERS.items():

ai_simple_engine_runtime_diffusers-0.0.2/src/ai_simple_engine_runtime_diffusers/models/scheduler/utils.py

@@ -1,44 +0,0 @@
-"""
-TODO: Maybe rename and move this module.
-"""
-from diffusers.schedulers import (
-    DDIMScheduler,
-    EulerDiscreteScheduler,
-    EulerAncestralDiscreteScheduler,
-    DPMSolverMultistepScheduler,
-    FlowMatchEulerDiscreteScheduler
-)
-
-SCHEDULERS = {
-    'ddim': DDIMScheduler,
-    'euler': EulerDiscreteScheduler,
-    'euler_a': EulerAncestralDiscreteScheduler,
-    'dpmpp_2m': DPMSolverMultistepScheduler,
-    'flow_match_euler': FlowMatchEulerDiscreteScheduler
-}
-"""
-The list including all the schedulers we are
-accepting.
-"""
-
-
-def get_scheduler_class(
-    scheduler: str,
-    model: LoadedModel
-) -> Union[SchedulerMixin, None]:
-    scheduler_class = SCHEDULERS.get(scheduler, None)
-
-    if scheduler_class is None:
-        raise Exception(f'The scheduler "{scheduler}" is not accepted by our system.')
-    
-    return scheduler_class
-
-    scheduler_cls = registry.resolve(model.info.scheduler.identifier)
-
-    scheduler = scheduler_cls.from_config(...)
-
-    return scheduler
-
-    return scheduler_cls.from_config(
-        model.model.scheduler.config
-    )