ai-simple-engine-runtime-diffusers 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py

@@ -26,7 +26,7 @@ class DiffusersLatentDiffusionModelExecutor (
         self,
         context: ExecutionContext,
         model: LoadedModel,
-    ):
+    ) -> DiffusersScheduler:
         scheduler_registry = context.services.get(DiffusersSchedulerRegistry)
 
         scheduler_class = scheduler_registry.resolve(model.info.scheduler.identifier)
@@ -47,7 +47,7 @@ class DiffusersLatentDiffusionModelExecutor (
         embeddings: PromptEmbeddings,
         timestep,
         guidance_scale
-    ):
+    ) -> torch.Tensor:
         runtime_model = model.instance
         unet = runtime_model.unet
         latent_input = torch.cat([latents, latents])
@@ -112,21 +112,66 @@ class DiffusersLatentDiffusionModelExecutor (
         *,
         model: LoadedModel,
         latents
-    ) -> Image:
+    ) -> torch.Tensor:
         runtime_model = model.instance
 
         vae = runtime_model.vae
 
+        # Scale based on the vae (mandatory)
         latents = latents / vae.config.scaling_factor
 
+        # Decode the image (mandatory)
         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)
+
+        """
+        The VAE of Stable Diffusion return in the [-1, 1]
+        so we need to transform to [0, 1]. This is
+        specific from the Stable Diffusion.
+        """
+        image = self._normalize_image(
+            image = image,
+            model = model
+        )
+        # image = (image / 2 + 0.5).clamp(0, 1)
+
+        return image
+
+        """
+        TODO: This is to make the image ready be read as a
+        real image and should be done here.
+        """
         image = image.cpu().permute(0, 2, 3, 1).float().detach().numpy()
         image = Image(image)
 
-        return image
+        return image
+    
+    # TODO: Maybe rename to '_normalize_tensor' (?)
+    def _normalize_image(
+        self,
+        image: torch.Tensor,
+        model: LoadedModel
+    ) -> torch.Tensor:
+        """
+        *For internal use only*
+
+        Normalize the `image` provided according to
+        the specifications in the information of the
+        `model` given.
+        """
+        minimum, maximum = model.info.tensor_output_spec.value_range
+
+        if (minimum, maximum) == (-1.0, 1.0):
+            image = (image + 1.0) / 2.0
+            # Previously:
+            # image = (image / 2 + 0.5)
+        elif (minimum, maximum) == (0.0, 1.0):
+            pass
+        else:
+            image = (image - minimum) / (maximum - minimum)
+
+        return image.clamp(0.0, 1.0)

ai_simple_engine_runtime_diffusers/models/executor/noise/pytorch_noise_generator.py

@@ -6,8 +6,11 @@ import torch
 class TorchNoiseGenerator(
     NoiseGenerator
 ):
+    """
+    Class to generate random noise using `torch`.
+    """
 
-    def generate(
+    async def generate(
         self,
         *,
         shape,

ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py

@@ -1,6 +1,7 @@
 from ai_simple_engine_runtime_diffusers.models.runtime.diffusers_latent_diffusion_model import DiffusersLatentDiffusionModel
 from ai_simple_engine_runtime_diffusers.models.runtime.utils import get_torch_dtype_for
 from ai_simple_engine_diffusion.model.info.diffusion_model_info import DiffusionModelInfo
+from ai_simple_engine_diffusion.model.info.tensor_output_spec import TensorOutputSpec
 from ai_simple_engine_diffusion.scheduler.spec.base import SchedulerSpec
 from ai_simple_engine.device.base import Device
 from ai_simple_engine.models.loaders.abstract import ModelLoader
@@ -84,7 +85,11 @@ class LatentDiffusionModelLoader(
             latent_channels = unet.config.in_channels,
             vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1),
             # TODO: Make dynamic depending on the 'installed_model'
-            scheduler = self._default_scheduler(installed_model)
+            scheduler = self._default_scheduler(installed_model),
+            # TODO: Is this ok (?)
+            tensor_output_spec = TensorOutputSpec(
+                value_range = (-1.0, 1.0)
+            )
         )
 
         return LoadedModel(

{ai_simple_engine_runtime_diffusers-0.0.4.dist-info → ai_simple_engine_runtime_diffusers-0.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ai-simple-engine-runtime-diffusers
-Version: 0.0.4
+Version: 0.0.6
 Summary: AI Simple Engine Runtime Diffusers Module
 License-File: LICENSE
 Author: danialcala94
@@ -11,7 +11,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 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.4,<1.0.0)
+Requires-Dist: ai-simple-engine-diffusion (>=0.0.8,<1.0.0)
 Requires-Dist: diffusers (>=0.38.0,<9999.0.0)
 Requires-Dist: transformers (>=5.13.0,<9999.0.0)
 Description-Content-Type: text/markdown

{ai_simple_engine_runtime_diffusers-0.0.4.dist-info → ai_simple_engine_runtime_diffusers-0.0.6.dist-info}/RECORD

@@ -1,16 +1,16 @@
 ai_simple_engine_runtime_diffusers/__init__.py,sha256=oeCwDniRRNaeu1i8oSoIGl24B_le57R-tjJfbQIcaTA,101
 ai_simple_engine_runtime_diffusers/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_simple_engine_runtime_diffusers/models/executor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py,sha256=mBKcpDmnSo-AroNNaMlQq64aqpynbIm_ULY2aaXi8KE,4093
+ai_simple_engine_runtime_diffusers/models/executor/diffusers_latent_diffusion_model_executor.py,sha256=ublq2CEj20YLjHZVIAhrPcCNOwGNJl26Tv2Y9lgse_Y,5458
 ai_simple_engine_runtime_diffusers/models/executor/noise/__init__.py,sha256=gtBak1TwMesgIsuc56HdYmCwrTjdysO68FO6XzAs-ME,69
-ai_simple_engine_runtime_diffusers/models/executor/noise/pytorch_noise_generator.py,sha256=kFL42BiG2odZp7hAA-7FZakihakOzc8BpYorysDvWx0,508
+ai_simple_engine_runtime_diffusers/models/executor/noise/pytorch_noise_generator.py,sha256=hwEdnPua6jMWnFiNlrnodsMJhAOihN3A-_62RGWUdhM,583
 ai_simple_engine_runtime_diffusers/models/executor/prompt_embeddings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_simple_engine_runtime_diffusers/models/executor/prompt_embeddings/prompt_embeddings.py,sha256=Y-Krd8UkaaZofSTZoY3Iq1lYetNBBfa_x8ivkplrpBA,321
 ai_simple_engine_runtime_diffusers/models/executor/text_encoder_model_executor_abstract.py,sha256=DdGE73KVvPZOjD1ulk3ffBC3RjAx4MIf9vcsyatgUyE,805
 ai_simple_engine_runtime_diffusers/models/executor/unet_model_executor_abstract.py,sha256=xj5Fw8J1jAwJZRoE5dEnH4Ewf92jIyEh6BJEOHOY1n0,1004
 ai_simple_engine_runtime_diffusers/models/executor/vae_model_executor_abstract.py,sha256=t3j4jYHKbo7bSjMhLLOJYsY2IZm7uQGXXK-c7QDCTv8,1045
 ai_simple_engine_runtime_diffusers/models/loader/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py,sha256=C-vENcRTMk9__Es_TdP38j3avsqK2olIFtXXmC_Tcao,3645
+ai_simple_engine_runtime_diffusers/models/loader/latent_diffusion_model_loader.py,sha256=euhlOFfivrkPFkLmjX8uwHvZ2v6L5S2TAayrCykWRoY,3879
 ai_simple_engine_runtime_diffusers/models/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_simple_engine_runtime_diffusers/models/runtime/diffusers_latent_diffusion_model.py,sha256=UTgR2mWW3rDtJXyRATpR_lfcWS6wMugeKHlvKneXTbw,478
 ai_simple_engine_runtime_diffusers/models/runtime/utils.py,sha256=niSrKW6RdZwtRPQVEFCArTyYdGgWWk-rlFSd2ledSPo,347
@@ -20,7 +20,7 @@ ai_simple_engine_runtime_diffusers/models/scheduler/registry/__init__.py,sha256=
 ai_simple_engine_runtime_diffusers/models/scheduler/registry/diffusers_scheduler_registry.py,sha256=FOuIAAwNmi3_jB25ICOlCPbD7U1WlOcGFZT3rhql5cY,565
 ai_simple_engine_runtime_diffusers/models/scheduler/utils.py,sha256=tIzwp7fy0Y1pE7XcLJU0dEp4XH2TMxBzY7SnvyfwITM,554
 ai_simple_engine_runtime_diffusers/plugin.py,sha256=o_5OJSwHj_4Qy1LW-Y_JRIP3-LotSOEUVgB3LdnK3mU,2205
-ai_simple_engine_runtime_diffusers-0.0.4.dist-info/licenses/LICENSE,sha256=6kbiFSfobTZ7beWiKnHpN902HgBx-Jzgcme0SvKqhKY,1091
-ai_simple_engine_runtime_diffusers-0.0.4.dist-info/METADATA,sha256=BI4tzPK8BpyBX-IFB-ZZZ-sSnSVxXgGiVwhSYhGQxbs,841
-ai_simple_engine_runtime_diffusers-0.0.4.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
-ai_simple_engine_runtime_diffusers-0.0.4.dist-info/RECORD,,
+ai_simple_engine_runtime_diffusers-0.0.6.dist-info/licenses/LICENSE,sha256=6kbiFSfobTZ7beWiKnHpN902HgBx-Jzgcme0SvKqhKY,1091
+ai_simple_engine_runtime_diffusers-0.0.6.dist-info/METADATA,sha256=1IycPuH0vBg3I95R0i9MDKFWUCrlcTUAoI3Jr1_tdLg,841
+ai_simple_engine_runtime_diffusers-0.0.6.dist-info/WHEEL,sha256=M5asmiAlL6HEcOq52Yi5mmk9KmTVjY2RDPtO4p9DMrc,88
+ai_simple_engine_runtime_diffusers-0.0.6.dist-info/RECORD,,