ai-edge-torch-nightly 0.3.0.dev20240823__py3-none-any.whl → 0.3.0.dev20240828__py3-none-any.whl

ai_edge_torch/_convert/test/test_convert_multisig.py

@@ -17,11 +17,47 @@
 import ai_edge_torch
 from ai_edge_torch.testing import model_coverage
 import torch
-import torchvision
+from torch import nn
 
 from tensorflow.python.platform import googletest
 
 
+class FullyConnectedModel(nn.Module):
+  """A simple fully connected model with two fully connected layers."""
+
+  def __init__(self, input_size, hidden_size, output_size):
+    super(FullyConnectedModel, self).__init__()
+    self.fc = nn.Linear(input_size, hidden_size)  # Fully connected layer
+    self.relu = nn.ReLU()  # Activation function
+    self.output = nn.Linear(hidden_size, output_size)
+
+  def forward(self, x):
+    x = self.fc(x)
+    x = self.relu(x)
+    x = self.output(x)
+    return x
+
+
+class FullyConvModel(nn.Module):
+  """A simple fully convolutional model with two convolutions."""
+
+  def __init__(self):
+    super(FullyConvModel, self).__init__()
+    self.conv1 = nn.Conv2d(
+        3, 16, kernel_size=3, padding=1
+    )  # Input channels: 3 (RGB), Output channels: 16
+    self.relu = nn.ReLU(inplace=True)
+    self.conv2 = nn.Conv2d(
+        16, 1, kernel_size=1
+    )  # Output channels: 1 (single channel output)
+
+  def forward(self, x):
+    x = self.conv1(x)
+    x = self.relu(x)
+    x = self.conv2(x)
+    return x
+
+
 class TestConvertMultiSignature(googletest.TestCase):
   """Tests conversion of various modules through multi-signature conversion."""
 
@@ -29,12 +65,12 @@ class TestConvertMultiSignature(googletest.TestCase):
     super().setUp()
     torch.manual_seed(0)
 
-  def test_convert_mobilenet_v2_with_default(self):
+  def test_convert_with_default(self):
     """Tests conversion of a model with two signatures one of which is the default."""
-    torch_module = torchvision.models.mobilenet_v2().eval()
+    torch_module = FullyConvModel().eval()
 
-    args = (torch.randn(4, 3, 224, 224),)
-    large_args = (torch.randn(4, 3, 336, 336),)
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
 
     signature_name = "large_input"
 
@@ -51,12 +87,12 @@ class TestConvertMultiSignature(googletest.TestCase):
         )
     )
 
-  def test_convert_mobilenet_v2_no_default(self):
+  def test_convert_no_default(self):
     """Tests conversion of a model with two signatures none of which is the default."""
-    torch_module = torchvision.models.mobilenet_v2().eval()
+    torch_module = FullyConvModel().eval()
 
-    args = (torch.randn(4, 3, 224, 224),)
-    large_args = (torch.randn(4, 3, 336, 336),)
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
 
     signature_name_1 = "input"
     signature_name_2 = "large_input"
@@ -84,12 +120,12 @@ class TestConvertMultiSignature(googletest.TestCase):
         )
     )
 
-  def test_convert_mobilenet_v2_signature_helper(self):
+  def test_convert_signature_helper(self):
     """Tests the ai_edge_torch.signature helper function works."""
-    torch_module = torchvision.models.mobilenet_v2().eval()
+    torch_module = FullyConvModel().eval()
 
-    args = (torch.randn(4, 3, 224, 224),)
-    large_args = (torch.randn(4, 3, 336, 336),)
+    args = (torch.randn(4, 3, 12, 12),)
+    large_args = (torch.randn(4, 3, 24, 24),)
 
     signature_name = "large_input"
 
@@ -108,39 +144,43 @@ class TestConvertMultiSignature(googletest.TestCase):
 
   def test_convert_separate_modules(self):
     """Tests conversion of two completely different modules as separate signatures."""
-    mobilentv2 = torchvision.models.mobilenet_v2().eval()
-    resnet18 = torchvision.models.resnet18().eval()
+    fully_conv = FullyConvModel().eval()
+    fully_connected = FullyConnectedModel(10, 5, 10).eval()
 
-    mobilenet_args = (torch.randn(4, 3, 224, 224),)
-    resnet_args = (torch.randn(4, 3, 224, 224),)
+    fully_conv_args = (torch.randn(4, 3, 12, 12),)
+    fully_connected_args = (torch.randn(10),)
 
-    mobilenet_signature_name = "mobilentv2"
-    resnet_signature_name = "resnet18"
+    fully_conv_signature_name = "fully_conv"
+    fully_connected_signature_name = "fully_connected"
 
     edge_model = (
         ai_edge_torch.signature(
-            mobilenet_signature_name, mobilentv2, mobilenet_args
+            fully_conv_signature_name, fully_conv, fully_conv_args
+        )
+        .signature(
+            fully_connected_signature_name,
+            fully_connected,
+            fully_connected_args,
         )
-        .signature(resnet_signature_name, resnet18, resnet_args)
-        .convert(resnet18, resnet_args)
+        .convert(fully_connected, fully_connected_args)
     )
 
-    mobilenet_inference_args = (torch.randn(4, 3, 224, 224),)
-    resnet_inference_args = (torch.randn(4, 3, 224, 224),)
+    fully_conv_inference_args = (torch.randn(4, 3, 12, 12),)
+    fully_connected_inference_args = (torch.randn(10),)
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             edge_model,
-            mobilentv2,
-            mobilenet_inference_args,
-            signature_name=mobilenet_signature_name,
+            fully_conv,
+            fully_conv_inference_args,
+            signature_name=fully_conv_signature_name,
         )
     )
     self.assertTrue(
         model_coverage.compare_tflite_torch(
             edge_model,
-            resnet18,
-            resnet_inference_args,
-            signature_name=resnet_signature_name,
+            fully_connected,
+            fully_connected_inference_args,
+            signature_name=fully_connected_signature_name,
         )
     )
 

ai_edge_torch/generative/examples/stable_diffusion/pipeline.py

@@ -15,16 +15,16 @@
 
 import argparse
 import os
-from pathlib import Path
-from typing import Dict, Optional
+import pathlib
+from typing import Optional
 
-import ai_edge_torch.generative.examples.stable_diffusion.samplers as samplers
-from ai_edge_torch.generative.examples.stable_diffusion.tokenizer import Tokenizer  # NOQA
-import ai_edge_torch.generative.examples.stable_diffusion.util as util
-from ai_edge_torch.model import TfLiteModel
+import ai_edge_torch
+from ai_edge_torch.generative.examples.stable_diffusion import samplers
+from ai_edge_torch.generative.examples.stable_diffusion import tokenizer
+from ai_edge_torch.generative.examples.stable_diffusion import util
 import numpy as np
 from PIL import Image
-from tqdm import tqdm
+import tqdm
 
 arg_parser = argparse.ArgumentParser()
 arg_parser.add_argument(
@@ -104,12 +104,12 @@ class StableDiffusion:
       diffusion_ckpt: str,
       decoder_ckpt: str
   ):
-    self.tokenizer = Tokenizer(tokenizer_vocab_dir)
-    self.clip = TfLiteModel.load(clip_ckpt)
-    self.decoder = TfLiteModel.load(decoder_ckpt)
-    self.diffusion = TfLiteModel.load(diffusion_ckpt)
+    self.tokenizer = tokenizer.Tokenizer(tokenizer_vocab_dir)
+    self.clip = ai_edge_torch.model.TfLiteModel.load(clip_ckpt)
+    self.decoder = ai_edge_torch.model.TfLiteModel.load(decoder_ckpt)
+    self.diffusion = ai_edge_torch.model.TfLiteModel.load(diffusion_ckpt)
     if encoder_ckpt is not None:
-      self.encoder = TfLiteModel.load(encoder_ckpt)
+      self.encoder = ai_edge_torch.model.TfLiteModel.load(encoder_ckpt)
 
 
 def run_tflite_pipeline(
@@ -128,48 +128,32 @@ def run_tflite_pipeline(
 ):
   """Run stable diffusion pipeline with tflite model.
 
-  model:
-
-    StableDiffsuion model.
-  prompt:
-    The prompt to guide the image generation.
-  output_path:
-    The path to the generated output image.
-  uncond_prompt:
-    The prompt not to guide the image generation.
-  cfg_scale:
-    Guidance scale of classifier-free guidance. Higher guidance scale encourages
-    to generate
-    images that are closely linked to the text `prompt`, usually at the expense
-    of lower
-    image quality.
-  height:
-    The height in pixels of the generated image.
-  width:
-    The width in pixels of the generated image.
-  sampler:
-    A sampler to be used to denoise the encoded image latents. Can be one of
-    `k_lms, `k_euler`,
-    or `k_euler_ancestral`.
-  n_inference_steps:
-    The number of denoising steps. More denoising steps usually lead to a higher
-    quality image at the
-    expense of slower inference. This parameter will be modulated by `strength`.
-  seed:
-    A seed to make generation deterministic.
-  strength:
-    Conceptually, indicates how much to transform the reference `input_image`.
-    Must be between 0 and 1.
-    `input_image` will be used as a starting point, adding more noise to it the
-    larger the `strength`.
-    The number of denoising steps depends on the amount of noise initially
-    added. When `strength` is 1,
-    added noise will be maximum and the denoising process will run for the full
-    number of iterations
-    specified in `n_inference_steps`. A value of 1, therefore, essentially
-    ignores `input_image`.
-  input_image:
-    Image which is served as the starting point for the image generation.
+  Args:
+    model: StableDiffsuion model.
+    prompt: The prompt to guide the image generation.
+    output_path: The path to the generated output image.
+    uncond_prompt: The prompt not to guide the image generation.
+    cfg_scale: Guidance scale of classifier-free guidance. Higher guidance scale
+      encourages to generate images that are closely linked to the text
+      `prompt`, usually at the expense of lower image quality.
+    height: The height in pixels of the generated image.
+    width: The width in pixels of the generated image.
+    sampler: A sampler to be used to denoise the encoded image latents. Can be
+      one of `k_lms, `k_euler`, or `k_euler_ancestral`.
+    n_inference_steps: The number of denoising steps. More denoising steps
+      usually lead to a higher quality image at the expense of slower inference.
+      This parameter will be modulated by `strength`.
+    seed: A seed to make generation deterministic.
+    strength: Conceptually, indicates how much to transform the reference
+      `input_image`. Must be between 0 and 1. `input_image` will be used as a
+      starting point, adding more noise to it the larger the `strength`. The
+      number of denoising steps depends on the amount of noise initially added.
+      When `strength` is 1, added noise will be maximum and the denoising
+      process will run for the full number of iterations specified in
+      `n_inference_steps`. A value of 1, therefore, essentially ignores
+      `input_image`.
+    input_image: Image which is served as the starting point for the image
+      generation.
   """
   if not 0 < strength < 1:
     raise ValueError('strength must be between 0 and 1')
@@ -202,7 +186,8 @@ def run_tflite_pipeline(
   context = np.concatenate([cond_context, uncond_context], axis=0)
   noise_shape = (1, 4, height // 8, width // 8)
 
-  # Initialization starts from input_image if any, otherwise, starts from a random sampling.
+  # Initialization starts from input_image if any, otherwise, starts from a
+  # random sampling.
   if input_image:
     if not hasattr(model, 'encoder'):
       raise AttributeError(
@@ -210,7 +195,6 @@ def run_tflite_pipeline(
           ' input_image.'
       )
     input_image = input_image.resize((width, height))
-    input_image_np = np.array(input_image).astype(np.float32)
     input_image_np = util.rescale(input_image, (0, 255), (-1, 1))
     input_image_np = util.move_channel(input_image_np, to='first')
     encoder_noise = np.random.normal(size=noise_shape).astype(np.float32)
@@ -223,8 +207,8 @@ def run_tflite_pipeline(
     latents *= sampler.initial_scale
 
   # Diffusion process.
-  timesteps = tqdm(sampler.timesteps)
-  for i, timestep in enumerate(timesteps):
+  timesteps = tqdm.tqdm(sampler.timesteps)
+  for _, timestep in enumerate(timesteps):
     time_embedding = util.get_time_embedding(timestep)
 
     input_latents = latents * sampler.get_input_scale()
@@ -242,7 +226,7 @@ def run_tflite_pipeline(
   images = util.rescale(images, (-1, 1), (0, 255), clamp=True)
   images = util.move_channel(images, to='last')
   if not os.path.exists(output_path):
-    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    pathlib.Path(output_path).parent.mkdir(parents=True, exist_ok=True)
   Image.fromarray(images[0].astype(np.uint8)).save(output_path)
 
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20240823"
+__version__ = "0.3.0.dev20240828"

{ai_edge_torch_nightly-0.3.0.dev20240823.dist-info → ai_edge_torch_nightly-0.3.0.dev20240828.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20240823
+Version: 0.3.0.dev20240828
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI
@@ -20,7 +20,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.9, <3.12
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy

{ai_edge_torch_nightly-0.3.0.dev20240823.dist-info → ai_edge_torch_nightly-0.3.0.dev20240828.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=48qP37uHT90YPs4eIUQxCiWVwqGEX3idCUs6mQKvX1U,116
 ai_edge_torch/config.py,sha256=PCd9PVrbUNeVIUDFUCnW4goDWU4bjouK28yMYU6VOi0,877
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=7tox6sdFIlCYPLDYpjFcD8cPTSivURCL_VV6-Dt5Sfc,4910
-ai_edge_torch/version.py,sha256=a63GrjqX4sRjk0WbC_0gGhT-ax_TLEi4iCEw0Iys7bw,706
+ai_edge_torch/version.py,sha256=FYYHMZbGfkHhNzKfScFWouShgQ9DOVNXZ7WWrrsKjPY,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=kcv_QgNgeyDmrqwdzHicGNP68w6zF7GJg7YkMEIXp4Q,3759
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -28,7 +28,7 @@ ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/layout_partitio
 ai_edge_torch/_convert/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/test/test_convert.py,sha256=tvj7fWHHmA9ddtcu-Fp3lJ6emaAQMrtK9wCG0cjgRAo,14413
 ai_edge_torch/_convert/test/test_convert_composites.py,sha256=CBiOqq-m7QT2ggBI1jBl9MkTIT5d0nK1tA0BUga0LGs,7994
-ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=0xIkahEU26Qx9GGn6Dm05ObIqJvsCdh692dREcaHEdE,4725
+ai_edge_torch/_convert/test/test_convert_multisig.py,sha256=ShZeakqVN9jg9mgFvvMWP0BoPF-u4BTM2hoEkDmLwj0,5780
 ai_edge_torch/_convert/test/test_to_channel_last_io.py,sha256=jLAmyHw5llT2ff8qA8mem3eVN57e_o5EpBnW72ZtP2I,3026
 ai_edge_torch/debug/__init__.py,sha256=N05Mmvi41KgSuK0JhuMejERESgP8QekiGdp9_PEyuKU,742
 ai_edge_torch/debug/culprit.py,sha256=7UYVpVWpiCXbMAyThVtHt_kc_poT7sCTh5UUPvcycgk,14832
@@ -64,7 +64,7 @@ ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=7
 ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=slieF2-QcDCwd4DRZ7snsZIphT97IXpp4plRRsRSwL8,13983
 ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=7oUIJ6HO0vmlhFdkXpqGm9KTB-eM4Ob9VrHSDlIGFOg,30926
 ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
-ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=Wc94X_kEnbInTAXFgf-VuCvv1A0HxxWrFZ7Tsq3Z8n0,8662
+ai_edge_torch/generative/examples/stable_diffusion/pipeline.py,sha256=x9lEEENGNbpx6VTf_LTVudd9d6bs9tLvFUKTl252zEY,8623
 ai_edge_torch/generative/examples/stable_diffusion/tokenizer.py,sha256=xychak9hdLd6ieXBYEwrK2BkF8NRZWZSSCijIsESpBA,3420
 ai_edge_torch/generative/examples/stable_diffusion/util.py,sha256=XIXIB0vCvQKOGyIyiZeiIA5DLeSXjkudywvJS4FK7AM,2431
 ai_edge_torch/generative/examples/stable_diffusion/samplers/__init__.py,sha256=uQWKzCD_49ackNFrt50H04dkDXxfAwUCtMWWQre5SVE,830
@@ -137,8 +137,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20240823.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20240823.dist-info/METADATA,sha256=OPYTq5RQCL2lvIFeBKOwbFusi4rq_Qo2ytgn_JQTVb0,1885
-ai_edge_torch_nightly-0.3.0.dev20240823.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
-ai_edge_torch_nightly-0.3.0.dev20240823.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20240823.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20240828.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20240828.dist-info/METADATA,sha256=Q5exnNzCT4z8p2mh5MFKl7MgfezYvW0ebW9k4PW8jJo,1878
+ai_edge_torch_nightly-0.3.0.dev20240828.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
+ai_edge_torch_nightly-0.3.0.dev20240828.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20240828.dist-info/RECORD,,