optimum-rbln 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

optimum/rbln/transformers/models/t5/t5_architecture.py

@@ -43,6 +43,12 @@ if TYPE_CHECKING:
     from transformers import T5ForConditionalGeneration
 
 
+class T5Wrapper:
+    def __init__(self, model):
+        self.encoder = T5EncoderWrapper(model)
+        self.decoder = T5DecoderWrapper(model)
+
+
 class T5Encoder(T5Stack):
     def forward(
         self,
@@ -122,19 +128,26 @@ class T5EncoderWrapper(torch.nn.Module):
         )
         self.encoder_max_length = None
         self.decoder_max_length = None
-        self.decoder_batch_size = 1
 
-    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
-                cross_key_value: torch.Tensor = None, batch_idx: torch.Tensor = None,
-                ) -> torch.Tensor:
-        encoder_batch_size = input_ids.shape[0]
-        decoder_batch_size = self.decoder_batch_size
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        cross_key_value: torch.Tensor = None,
+        batch_idx: torch.Tensor = None,
+    ) -> torch.Tensor:
         decoder_max_length = self.decoder_max_length or self.default_max_length
         encoder_max_length = self.encoder_max_length or self.default_max_length
 
         attn_layer = self.encoder.block[0].layer[0].SelfAttention
         encoder_position_bias = T5Attention.compute_bias(attn_layer, encoder_max_length, encoder_max_length)
-        encoder_outputs = T5Encoder.forward(self.encoder, input_ids, attention_mask, encoder_position_bias, batch_ids=torch.tensor(0, dtype=torch.int32))
+        encoder_outputs = T5Encoder.forward(
+            self.encoder,
+            input_ids,
+            attention_mask,
+            encoder_position_bias,
+            batch_ids=torch.tensor(0, dtype=torch.int32),
+        )
 
         attn_layer = self.decoder.block[0].layer[0].SelfAttention
         decoder_position_bias = T5Attention.compute_bias(attn_layer, decoder_max_length, decoder_max_length)
@@ -145,22 +158,14 @@ class T5EncoderWrapper(torch.nn.Module):
 
         dummy_past_key_value = []
         for i in range(self.config.num_layers):
-            pkv_self_attn_key = torch.zeros(
-                decoder_batch_size, self.config.num_heads, decoder_max_length, self.config.d_kv
-            )
-            pkv_self_attn_value = torch.zeros(
-                decoder_batch_size, self.config.num_heads, decoder_max_length, self.config.d_kv
-            )
-            pkv_cross_attn_key = torch.zeros(
-                encoder_batch_size, self.config.num_heads, encoder_max_length, self.config.d_kv
-            )
-            pkv_cross_attn_value = torch.zeros(
-                encoder_batch_size, self.config.num_heads, encoder_max_length, self.config.d_kv
-            )
+            pkv_self_attn_key = torch.zeros(1, self.config.num_heads, decoder_max_length, self.config.d_kv)
+            pkv_self_attn_value = torch.zeros(1, self.config.num_heads, decoder_max_length, self.config.d_kv)
+            pkv_cross_attn_key = torch.zeros(1, self.config.num_heads, encoder_max_length, self.config.d_kv)
+            pkv_cross_attn_value = torch.zeros(1, self.config.num_heads, encoder_max_length, self.config.d_kv)
             layer_pkv = (pkv_self_attn_key, pkv_self_attn_value, pkv_cross_attn_key, pkv_cross_attn_value)
             dummy_past_key_value.append(layer_pkv)
 
-        decoder_attention_mask = torch.zeros(decoder_batch_size, decoder_max_length, dtype=torch.float32)
+        decoder_attention_mask = torch.zeros(1, decoder_max_length, dtype=torch.float32)
         decoder_attention_mask[:, :1] = 1
 
         # Since first step of decoder has different graph to further step of it,
@@ -168,7 +173,7 @@ class T5EncoderWrapper(torch.nn.Module):
         # TODO(jongho): Separate first-step-decoder.
         decoder_outputs = T5Decoder.forward(
             self.decoder,
-            input_ids=torch.zeros(decoder_batch_size, 1, dtype=torch.int64),
+            input_ids=torch.zeros(1, 1, dtype=torch.int64),
             attention_mask=decoder_attention_mask,
             position_bias=decoder_position_bias,
             encoder_decoder_position_bias=encoder_decoder_position_bias,
@@ -187,7 +192,7 @@ class T5EncoderWrapper(torch.nn.Module):
             cross_kv_cache.append(past_key_values[i][3])
         cross_kv_cache = torch.stack(cross_kv_cache, dim=0)
 
-        cross_key_value = cross_key_value.slice_scatter(cross_kv_cache, dim=1, start=batch_idx, end=batch_idx+1)
+        cross_key_value = cross_key_value.slice_scatter(cross_kv_cache, dim=1, start=batch_idx, end=batch_idx + 1)
 
         return cross_key_value
 
@@ -240,6 +245,7 @@ class T5DecoderWrapper(torch.nn.Module):
         attn_layer = self.model.decoder.block[0].layer[0].SelfAttention
         _decoder_position_bias = T5Attention.compute_bias(attn_layer, decoder_max_length, decoder_max_length)
 
+        # position_bias need to compute with batch (for cb)
         batch_decoder_position_bias = []
         for i in range(input_ids.shape[0]):
             batch_position_bias = _decoder_position_bias[:, :, cache_position[i][0]].unsqueeze(2)
@@ -259,7 +265,7 @@ class T5DecoderWrapper(torch.nn.Module):
             encoder_decoder_position_bias=encoder_decoder_position_bias,
             past_key_values=kv_cache,
             cache_position=cache_position,
-            batch_ids=rbln_batch_position
+            batch_ids=rbln_batch_position,
         )
 
         past_key_values = decoder_outputs.past_key_values
@@ -312,7 +318,7 @@ class _T5Attention(T5Attention):
             value_states = shape(self.v(hidden_states), batch_size)
         else:
             # cross-attn
-            if cache_position.dim() == 0 :
+            if cache_position.dim() == 0:
                 key_states = shape(self.k(key_value_states), key_value_states.shape[0])
                 value_states = shape(self.v(key_value_states), key_value_states.shape[0])
                 past_key_value = key_states, value_states
@@ -331,18 +337,24 @@ class _T5Attention(T5Attention):
                 batch_value_states = value_states[b].unsqueeze(0)
 
                 if is_self_attn and past_key_value is not None:
-                    batch_key_states = past_key_value[0][b].unsqueeze(0).slice_scatter(
-                        batch_key_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                    batch_key_states = (
+                        past_key_value[0][b]
+                        .unsqueeze(0)
+                        .slice_scatter(
+                            batch_key_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                        )
                     )
-                    batch_value_states = past_key_value[1][b].unsqueeze(0).slice_scatter(
-                        batch_value_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                    batch_value_states = (
+                        past_key_value[1][b]
+                        .unsqueeze(0)
+                        .slice_scatter(
+                            batch_value_states, dim=-2, start=cache_position[b][0], end=cache_position[b][0] + 1
+                        )
                     )
 
                 scores = torch.matmul(batch_query_states, batch_key_states.transpose(3, 2))
                 scores += position_bias[b]
-                attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(
-                    scores
-                )
+                attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
                 attn_output = unshape(torch.matmul(attn_weights, batch_value_states), 1)
                 all_key_states.append(batch_key_states)
                 all_value_states.append(batch_value_states)
@@ -371,7 +383,9 @@ class _T5Attention(T5Attention):
                 scores
             )  # (batch_size, n_heads, seq_length, key_length)
 
-            attn_output = unshape(torch.matmul(attn_weights, value_states), batch_size)  # (batch_size, seq_length, dim)
+            attn_output = unshape(
+                torch.matmul(attn_weights, value_states), batch_size
+            )  # (batch_size, seq_length, dim)
 
         attn_output = self.o(attn_output)
         present_key_value = (key_states, value_states)

optimum/rbln/transformers/models/wav2vec2/modeling_wav2vec2.py

@@ -65,7 +65,6 @@ class RBLNWav2Vec2ForCTC(RBLNModel):
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    model_type = "rbln_model"
     main_input_name = "input_values"
     auto_model_class = AutoModelForMaskedLM
 

optimum/rbln/transformers/models/whisper/modeling_whisper.py

@@ -59,14 +59,16 @@ if TYPE_CHECKING:
 class RBLNRuntimeEncoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
+    def forward(self, input_features: torch.Tensor = None):
         # backward compatibility transformers==4.40.2
         # https://github.com/huggingface/transformers/blob/4fdf58afb72b0754da30037fc800b6044e7d9c99/src/transformers/pipelines/automatic_speech_recognition.py#L494
-        input_features = kwargs.get("input_features", None)
-        if input_features is None:
-            input_features = args[0]
+
+        n_pad_to_batch = self.batch_size - input_features.shape[0]
+        if n_pad_to_batch > 0:
+            input_features = torch.nn.functional.pad(input_features, (0, 0, 0, 0, 0, n_pad_to_batch))
 
         _ = super().forward(input_features=input_features)
+
         # dummy output for generation
         return BaseModelOutput(last_hidden_state=torch.tensor([[-1.0]]))
 
@@ -74,18 +76,33 @@ class RBLNRuntimeEncoder(RBLNPytorchRuntime):
 class RBLNRuntimeDecoder(RBLNPytorchRuntime):
     mandatory_members = ["main_input_name"]
 
-    def forward(self, *args: List[torch.Tensor], **kwargs: Dict[str, torch.Tensor]):
-        outputs = super().forward(*args, **kwargs)
+    def forward(
+        self,
+        decoder_input_ids: torch.Tensor = None,
+        decoder_attention_mask: torch.Tensor = None,
+        cache_position: torch.Tensor = None,
+    ):
+        inputs_bsz = decoder_input_ids.shape[0]
+        padded_bsz = self.batch_size - inputs_bsz
+        if padded_bsz > 0:
+            decoder_input_ids = torch.nn.functional.pad(decoder_input_ids, (0, 0, 0, padded_bsz))
+
+        outputs = super().forward(
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            cache_position=cache_position,
+        )
+
         if isinstance(outputs, torch.Tensor):
-            return Seq2SeqLMOutput(logits=outputs, cross_attentions=None)
+            return Seq2SeqLMOutput(logits=outputs[:inputs_bsz], cross_attentions=None)
         else:
-            return Seq2SeqLMOutput(logits=outputs[0], cross_attentions=outputs[1])
+            return Seq2SeqLMOutput(logits=outputs[0][:inputs_bsz], cross_attentions=outputs[1][:, :inputs_bsz])
 
 
 class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin):
     """
     The Whisper Model with a language modeling head. Can be used for automatic speech recognition.
-    This model inherits from [`RBLNMultiModel`]. Check the superclass documentation for the generic methods the library implements for all its models.
+    This model inherits from [`RBLNDecoderOnlyModelForCausalLM`]. Check the superclass documentation for the generic methods the library implements for all its models.
 
     A class to convert and run pre-trained transformers based LlamaForCausalLM model on RBLN devices.
     It implements the methods to convert a pre-trained transformers LlamaForCausalLM model into a RBLN transformer model by:
@@ -93,7 +110,6 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
     - compiling the resulting graph using the RBLN compiler.
     """
 
-    model_type = "rbln_model"
     auto_model_class = AutoModelForSpeechSeq2Seq
     main_input_name = "input_ids"
 
@@ -104,8 +120,12 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         self.dec_max_seq_len = self.rbln_config.model_cfg["dec_max_seq_len"]
         self.rbln_token_timestamps = self.rbln_config.model_cfg["token_timestamps"]
 
-        self.encoder = RBLNRuntimeEncoder(runtime=self.model[0], main_input_name="input_features")
-        self.decoder = RBLNRuntimeDecoder(runtime=self.model[1], main_input_name="input_ids")
+        self.encoder = RBLNRuntimeEncoder(
+            runtime=self.model[0], main_input_name="input_features", batch_size=self.batch_size
+        )
+        self.decoder = RBLNRuntimeDecoder(
+            runtime=self.model[1], main_input_name="input_ids", batch_size=self.batch_size
+        )
 
         # skip encoder &  first decoder when language detected
         self.is_language_detected = False
@@ -200,7 +220,11 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         expected_seq_len = model_config.max_source_positions * 2
         num_mel_bins = model_config.num_mel_bins
         enc_max_seq_len = model_config.max_source_positions
-        rbln_dec_max_seq_len = model_config.max_length
+
+        # 'whisper-large-v3-turbo' doesn't have 'max_length', but PretrainedConfig have default value for the key 'max_length'
+        rbln_dec_max_seq_len = getattr(model_config, "max_target_positions", None)
+        if rbln_dec_max_seq_len is None:
+            rbln_dec_max_seq_len = model_config.max_length
 
         # model input info
         enc_input_info = [("input_features", [rbln_batch_size, num_mel_bins, expected_seq_len], "float32")]
@@ -273,6 +297,7 @@ class RBLNWhisperForConditionalGeneration(RBLNModel, RBLNWhisperGenerationMixin)
         self,
         input_ids,
         cache_position: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,  # need for support transformers>=4.45.0
         **kwargs,
     ):
         """

optimum/rbln/transformers/models/xlm_roberta/modeling_xlm_roberta.py

@@ -25,7 +25,7 @@ import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import torch
-from transformers import AutoModel, PretrainedConfig, PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
+from transformers import PretrainedConfig, PreTrainedModel, XLMRobertaConfig, XLMRobertaModel
 
 from ....modeling_base import RBLNModel
 from ....modeling_config import RBLNCompileConfig, RBLNConfig
@@ -38,7 +38,6 @@ if TYPE_CHECKING:
 
 
 class RBLNXLMRobertaModel(RBLNModel):
-    auto_model_class = AutoModel  # feature extraction
     original_model_class = XLMRobertaModel
     original_config_class = XLMRobertaConfig
 

optimum/rbln/transformers/utils/rbln_quantization.py

@@ -31,8 +31,13 @@ from torch.nn import functional as F
 
 # Constants
 QUANTIZED_WEIGHTS = {
-    "q_proj", "k_proj", "v_proj", "o_proj",
-    "gate_proj", "up_proj", "down_proj",
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj",
+    "gate_proj",
+    "up_proj",
+    "down_proj",
 }
 
 
@@ -81,6 +86,7 @@ def create_qlinear(layer: Linear) -> Linear:
     """
     Converts a standard linear layer to a quantized linear (qlinear) layer with a custom forward pass.
     """
+
     def qlinear_forward(self, inputs: torch.Tensor) -> torch.Tensor:
         if inputs.dtype != self.scales.dtype:
             raise TypeError(f"Expected input dtype {self.scales.dtype}, but got {inputs.dtype}")

optimum/rbln/utils/context.py

@@ -0,0 +1,58 @@
+# Copyright 2024 Rebellions Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Portions of this software are licensed under the Apache License,
+# Version 2.0. See the NOTICE file distributed with this work for
+# additional information regarding copyright ownership.
+
+# All other portions of this software, including proprietary code,
+# are the intellectual property of Rebellions Inc. and may not be
+# copied, modified, or distributed without prior written permission
+# from Rebellions Inc.
+
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Union
+
+from optimum.exporters import TasksManager
+from transformers import AutoConfig, AutoModel
+
+
+@contextmanager
+def override_auto_classes(config_func=None, model_func=None, skip_taskmanager=True):
+    """Temporarily override Auto classes with original model classes"""
+    original_config = AutoConfig.from_pretrained
+    original_model = AutoModel.from_pretrained
+    original_get_model_from_task = TasksManager.get_model_from_task
+
+    def get_model_from_task(
+        task: str,
+        model_name_or_path: Union[str, Path],
+        **kwargs,
+    ):
+        return model_func(model_name_or_path, **kwargs)
+
+    def none_func(*args, **kwargs):
+        return None
+
+    try:
+        AutoConfig.from_pretrained = config_func or none_func
+        AutoModel.from_pretrained = model_func or none_func
+        if skip_taskmanager:
+            TasksManager.get_model_from_task = none_func if model_func is None else get_model_from_task
+        yield
+    finally:
+        AutoConfig.from_pretrained = original_config
+        AutoModel.from_pretrained = original_model
+        TasksManager.get_model_from_task = original_get_model_from_task

optimum/rbln/utils/decorator_utils.py

@@ -0,0 +1,55 @@
+from functools import wraps
+
+from .logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def remove_compile_time_kwargs(func):
+    """
+    Decorator to handle compile-time parameters during inference.
+
+    For RBLN-optimized pipelines, several parameters must be determined during compilation
+    and cannot be modified during inference. This decorator:
+    1. Removes and warns about LoRA scale in cross_attention_kwargs
+    2. Removes and warns about image dimension parameters (height, width)
+
+    Args:
+        func: The pipeline's __call__ method to be wrapped
+    """
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        height_exists = "height" in kwargs and kwargs["height"] is not None
+        width_exists = "width" in kwargs and kwargs["width"] is not None
+        if height_exists or width_exists:
+            logger.warning(
+                "Image dimension parameters (`height`, `width`) will be ignored during inference. "
+                "Image dimensions must be specified during model compilation using from_pretrained()."
+            )
+            kwargs.pop("width", None)
+            kwargs.pop("height", None)
+
+        if "cross_attention_kwargs" in kwargs:
+            cross_attention_kwargs = kwargs.get("cross_attention_kwargs")
+            if not cross_attention_kwargs:
+                return func(self, *args, **kwargs)
+
+            has_scale = "scale" in cross_attention_kwargs
+            if has_scale:
+                logger.warning(
+                    "LoRA scale in cross_attention_kwargs will be ignored during inference. "
+                    "To adjust LoRA scale, specify it during model compilation using from_pretrained()."
+                )
+
+                # If scale is the only key, set to None
+                # Otherwise, remove scale and preserve other settings
+                if len(cross_attention_kwargs) == 1:
+                    kwargs["cross_attention_kwargs"] = None
+                else:
+                    kwargs["cross_attention_kwargs"].pop("scale")
+
+        return func(self, *args, **kwargs)
+
+    return wrapper

optimum/rbln/utils/import_utils.py

@@ -37,6 +37,27 @@ class VersionCompat:
 
 
 RBLN_VERSION_COMPATS = {
+    "0.1.13": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.6.0",
+            max_version="0.6.1",
+        ),
+    ],
+    "0.1.12": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.5.12",
+            max_version="0.5.13",
+        ),
+    ],
+    "0.1.11": [
+        VersionCompat(
+            package_name="rebel-compiler",
+            min_version="0.5.10",
+            max_version="0.5.11",
+        ),
+    ],
     "0.1.10": [
         VersionCompat(
             package_name="rebel-compiler",

optimum/rbln/utils/logging.py

@@ -22,7 +22,7 @@ log_levels = {
     "critical": logging.CRITICAL,
 }
 
-_default_log_level = logging.WARNING
+_default_log_level = logging.INFO
 
 
 def _get_default_logging_level():

optimum/rbln/utils/runtime_utils.py

@@ -67,7 +67,7 @@ class UnavailableRuntime:
         return iter([self])
 
     def forward(self, *args: List["torch.Tensor"], **kwargs: Dict[str, "torch.Tensor"]):
-        raise RuntimeError("RBLN-Runtime is not created, So it is not available.")
+        raise RuntimeError("The model can't run because the runtime hasn't been created.")
 
     def __repr__(self) -> str:
         return "UnavailableRuntime"
@@ -76,17 +76,17 @@ class UnavailableRuntime:
 class ContextRblnConfig:
     _local = threading.local()
 
-    def __init__(self, device, device_map, create_runtimes, optimze_host_mem):
+    def __init__(self, device=None, device_map=None, create_runtimes=None, optimize_host_mem=None):
         self.device = device
         self.device_map = device_map
         self.create_runtimes = create_runtimes
-        self.optimze_host_mem = optimze_host_mem
+        self.optimize_host_mem = optimize_host_mem
 
     def __enter__(self):
         self._local.device = self.device
         self._local.device_map = self.device_map
         self._local.create_runtimes = self.create_runtimes
-        self._local.optimize_host_memory = self.optimze_host_mem
+        self._local.optimize_host_memory = self.optimize_host_mem
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):

optimum/rbln/utils/timer_utils.py

@@ -1,5 +1,8 @@
+import os
 from datetime import datetime
 
+from halo import Halo
+
 from .logging import get_logger
 
 
@@ -9,11 +12,32 @@ logger = get_logger()
 def rbln_timer(print_name):
     def decorator(function):
         def wrapper(*args, **kwargs):
+            disable = os.getenv("OPTIMUM_RBLN_DISABLE_SPIN", "False").lower() in ("true", "1", "t")
+            if disable:
+                logger.info(f"{print_name} ...")
+
+            spinner = Halo(text=f"{print_name} ...", spinner="dots", color="green", enabled=(not disable))
+            spinner.start()
+
+            # Start timer
             tick = datetime.now()
-            result = function(*args, **kwargs)
-            logger.debug(f"{print_name}. Elasped time: {str(datetime.now() - tick)[:7]}")
+            try:
+                result = function(*args, **kwargs)
+            except Exception as e:
+                spinner.fail(f"{print_name} failed.")
+                raise e
+
+            # Print elapsed time.
+            if disable:
+                logger.info(f"{print_name} done. Elasped time: {format_elapsed_time(tick)}")
+
+            spinner.stop()
+            spinner.succeed(text=f"{print_name} done. Elasped time: {format_elapsed_time(tick)}")
             return result
 
         return wrapper
 
+    def format_elapsed_time(start_time: datetime) -> str:
+        return str(datetime.now() - start_time)[:7]
+
     return decorator

{optimum_rbln-0.1.11.dist-info → optimum_rbln-0.1.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: optimum-rbln
-Version: 0.1.11
+Version: 0.1.13
 Summary: Optimum RBLN is the interface between the Hugging Face Transformers and Diffusers libraries and RBLN accelerators.
          It provides a set of tools enabling easy model loading and inference on single and multiple rbln device settings for different downstream tasks.
          
@@ -14,22 +14,24 @@ Classifier: Intended Audience :: Education
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: POSIX :: Linux
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Project-URL: Homepage, https://rebellions.ai
 Project-URL: Documentation, https://docs.rbln.ai
-Requires-Python: <3.11,>=3.8
-Requires-Dist: torch<=2.2.1
-Requires-Dist: torchvision<=0.17.1
-Requires-Dist: torchaudio<=2.2.1
-Requires-Dist: optimum<=1.22.0
+Requires-Python: <3.13,>=3.9
+Requires-Dist: torch<=2.5.1
+Requires-Dist: torchvision<=0.20.1
+Requires-Dist: torchaudio<=2.5.1
+Requires-Dist: optimum==1.23.1
 Requires-Dist: accelerate>=0.28.0
-Requires-Dist: transformers<=4.44.2,>=4.43.2
-Requires-Dist: diffusers<=0.30.3
+Requires-Dist: transformers==4.45.2
+Requires-Dist: diffusers<=0.31.0
 Requires-Dist: einops>=0.8.0
 Requires-Dist: packaging>=24.1
+Requires-Dist: halo
 Provides-Extra: tests
 Requires-Dist: pytest>=8.1.1; extra == "tests"
 Requires-Dist: psutil>=5.9.8; extra == "tests"