### Install LitModels Python Package

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Installs the LitModels library using pip. This is the first step to using the library for managing AI model checkpoints.

```bash
pip install litmodels
```

--------------------------------

### Load Saved Model using litmodels

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Demonstrates how to download and load a previously saved model from cloud storage using litmodels.load_model. It then shows an example of running inference with the loaded model.

```python
from litmodels import load_model

# Download and load the model file from cloud storage
model = load_model(
    name="your_org/your_team/sklearn-svm-model", download_dir="my_models"
)

# Example: run inference with the loaded model
sample_input = [[5.1, 3.5, 1.4, 0.2]]
prediction = model.predict(sample_input)
print(f"Prediction: {prediction}")
```

--------------------------------

### Save and Load PyTorch Model using LitModels

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Provides examples for saving and loading standard PyTorch models. It utilizes `save_model` for persisting the model and `load_model` for retrieving it, using a specified name format for organization and identification.

```python
import torch
from litmodels import save_model

model = torch.nn.Module()
save_model(model=model, name="your_org/your_team/torch-model")
```

```python
from litmodels import load_model

model_ = load_model(name="your_org/your_team/torch-model")
```

--------------------------------

### Save and Load Custom PyTorch Model with PyTorchRegistryMixin

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Illustrates how to save and load custom PyTorch models using `PyTorchRegistryMixin`. This mixin centralizes serialization logic, ensuring consistent model persistence. The `download_model` method reconstructs the model directly from the registry without needing constructor arguments.

```python
import torch
from litmodels.integrations.mixins import PyTorchRegistryMixin


# Important: PyTorchRegistryMixin must be first in the inheritance order
class MyTorchModel(PyTorchRegistryMixin, torch.nn.Module):
    def __init__(self, input_size, hidden_size=128):
        super().__init__()
        self.linear = torch.nn.Linear(input_size, hidden_size)
        self.activation = torch.nn.ReLU()

    def forward(self, x):
        return self.activation(self.linear(x))


# Create and push the model
model = MyTorchModel(input_size=784)
model.upload_model(name="my-org/my-team/torch-model")


# Pull the model with the same architecture
loaded_model = MyTorchModel.download_model(name="my-org/my-team/torch-model")
```

--------------------------------

### Save and Load Custom Python Class with PickleRegistryMixin

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Demonstrates how to save and load any Python class using the `PickleRegistryMixin` from litmodels. This allows for modular model management and consistent code. The `upload_model` method pushes the model to a registry, and `download_model` reconstructs it.

```python
from litmodels.integrations.mixins import PickleRegistryMixin


class MyModel(PickleRegistryMixin):
    def __init__(self, param1, param2):
        self.param1 = param1
        self.param2 = param2
        # Your model initialization code
        ...


# Create and push a model instance
model = MyModel(param1=42, param2="hello")
model.upload_model(name="my-org/my-team/my-model")


# Load model:
loaded_model = MyModel.download_model(name="my-org/my-team/my-model")
```

--------------------------------

### PyTorch Lightning Callback for Model Checkpointing

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Shows how to integrate an automatic checkpointing callback into a PyTorch Lightning Trainer. This callback uploads the model at the end of each epoch to a specified model registry. Requires PyTorch, torchvision, and lightning libraries.

```python
import torch.utils.data as data
import torchvision as tv
from lightning import Trainer
from litmodels.integrations import LightningModelCheckpoint
from litmodels.demos import BoringModel

dataset = tv.datasets.MNIST(".", download=True, transform=tv.transforms.ToTensor())
train, val = data.random_split(dataset, [55000, 5000])

trainer = Trainer(
    max_epochs=2,
    callbacks=[
        LightningModelCheckpoint(
            # Define the model name - this should be unique to your model
            model_registry="<organization>/<teamspace>/<model-name>",
        )
    ],
)
trainer.fit(
    BoringModel(),
    data.DataLoader(train, batch_size=256),
    data.DataLoader(val, batch_size=256),
)
```

--------------------------------

### Download Checkpoint Files with LitModels

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Downloads model artifacts from Lightning Cloud to a local directory and returns the paths to the downloaded files. Supports downloading specific files or multiple files within a directory structure.

```python
import litmodels as lm
import torch
from lightning.pytorch.demos.boring_classes import BoringModel

# Download checkpoint file
checkpoint_path = lm.download_model(
    name="my-org/my-team/boring-model:epoch-10",
    download_dir="./downloaded_models",
    progress_bar=True
)
print(f"Downloaded to: {checkpoint_path}")

# Load the checkpoint manually
model = BoringModel()
if isinstance(checkpoint_path, list):
    checkpoint_path = checkpoint_path[0]
model.load_state_dict(torch.load(checkpoint_path))
model.eval()

# Download specific version
paths = lm.download_model(
    name="my-org/my-team/my-model:v2.1",
    download_dir="./models"
)
```

--------------------------------

### Download and Load Models with LitModels

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Downloads a model from Lightning Cloud and loads it into memory based on file extension. Supports .ts (TorchScript), .keras (TensorFlow/Keras), and .pkl (pickled objects). Handles different model types for inference.

```python
import litmodels as lm

# Load PyTorch TorchScript model
torch_model = lm.load_model(
    name="my-org/my-team/pytorch-mnist:v1",
    download_dir="./models"
)

# Load scikit-learn model and run inference
sklearn_model = lm.load_model(
    name="my-org/my-team/sklearn-iris-svm",
    download_dir="./models"
)
sample_input = [[5.1, 3.5, 1.4, 0.2]]
prediction = sklearn_model.predict(sample_input)
print(f"Predicted class: {prediction[0]}")

# Load Keras model
keras_model = lm.load_model(
    name="my-org/my-team/keras-classifier:latest",
    download_dir="./models"
)
```

--------------------------------

### Save and Load PyTorch Lightning Model with LitModels

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Shows how to save the best checkpoint of a PyTorch Lightning model and then load it for further training. It uses `upload_model` to save the checkpoint to cloud storage and `download_model` to retrieve it, integrating with the `lightning` trainer.

```python
from lightning import Trainer
from litmodels import upload_model
from litmodels.demos import BoringModel

# Configure Lightning Trainer
trainer = Trainer(max_epochs=2)
# Define the model and train it
trainer.fit(BoringModel())

# Upload the best model to cloud storage
checkpoint_path = getattr(trainer.checkpoint_callback, "best_model_path")
# Define the model name - this should be unique to your model
upload_model(model=checkpoint_path, name="<organization>/<teamspace>/<model-name>")
```

```python
from lightning import Trainer
from litmodels import download_model
from litmodels.demos import BoringModel

# Load the model from cloud storage
checkpoint_path = download_model(
    # Define the model name and version - this needs to be unique to your model
    name="<organization>/<teamspace>/<model-name>:<model-version>",
    download_dir="my_models",
)
print(f"model: {checkpoint_path}")

# Train the model with extended training period
trainer = Trainer(max_epochs=4)
trainer.fit(BoringModel(), ckpt_path=checkpoint_path)
```

--------------------------------

### Save and Load PyTorch Model with LitModels

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Demonstrates basic saving and loading of a PyTorch model using the litmodels library. It requires PyTorch and the litmodels package. The `save_model` function takes a model object and a name, while `load_model` retrieves it by name.

```python
import litmodels as lm
import torch

# save a model
model = torch.nn.Module()
lm.save_model(model=model, name="model-name")

# load a model
model = lm.load_model(name="model-name")
```

--------------------------------

### Download Model Files - Python

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Downloads all associated files for a specified model version from cloud storage. It supports progress bar display and can process downloaded files, such as JSON configurations.

```python
from pathlib import Path

files = download_model_files(
    name="my-org/my-team/custom-artifacts:v1.0",
    download_dir="./artifacts",
    progress_bar=True
)

for file_path in files:
    print(f"Downloaded: {file_path}")
    if Path(file_path).suffix == ".json":
        import json
        with open(file_path) as f:
            config = json.load(f)
            print(f"Config: {config}")
```

--------------------------------

### download_model_files: Low-level Cloud File Download

Source: https://context7.com/lightning-ai/litmodels/llms.txt

A low-level function for downloading model artifacts from cloud storage without automatic loading. It returns a list of the paths to the downloaded files, allowing for manual processing or integration into custom workflows.

```python
from litmodels.io.cloud import download_model_files
from pathlib import Path


```

--------------------------------

### Core API Functions

Source: https://context7.com/lightning-ai/litmodels/llms.txt

This section details the core functions for managing AI model checkpoints with LitModels.

```APIDOC
## save_model - Serialize and upload models to Lightning Cloud

### Description
Serializes an in-memory model (PyTorch, Keras, scikit-learn, or any picklable object) and uploads it to Lightning Cloud. Automatically detects the model type and uses the appropriate serialization format (.pth for PyTorch state_dict, .keras for TensorFlow, .pkl for others).

### Method
`save_model`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **name** (str) - Required - The name of the model to save, including organization and team (e.g., "my-org/my-team/model-name").
- **model** (object) - Required - The model object to serialize and upload (e.g., PyTorch model, scikit-learn estimator, Keras model).
- **progress_bar** (bool) - Optional - Whether to display a progress bar during upload. Defaults to False.
- **verbose** (int) - Optional - Verbosity level for the upload process. Defaults to 0.
- **metadata** (dict) - Optional - A dictionary of metadata to associate with the model checkpoint.

### Request Example
```python
import torch
import litmodels as lm
from sklearn import datasets, svm, model_selection

# PyTorch model - saves state_dict as .pth
model = torch.nn.Sequential(
    torch.nn.Linear(784, 128),
    torch.nn.ReLU(),
    torch.nn.Linear(128, 10)
)
lm.save_model(
    name="my-org/my-team/pytorch-mnist",
    model=model,
    progress_bar=True,
    metadata={"accuracy": "0.95", "epoch": "10"}
)

# Scikit-learn model - saves as .pkl
iris = datasets.load_iris()
X_train, X_test, y_train, y_test = model_selection.train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)
clf = svm.SVC(kernel='rbf', gamma='auto')
clf.fit(X_train, y_train)
lm.save_model(
    name="my-org/my-team/sklearn-iris-svm",
    model=clf,
    verbose=1
)

# TensorFlow/Keras model - saves as .keras
from tensorflow import keras
keras_model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(784,)),
    keras.layers.Dense(10, activation='softmax')
])
keras_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
lm.save_model(
    name="my-org/my-team/keras-classifier",
    model=keras_model
)
```

### Response
#### Success Response (200)
- **info** (object) - Information about the uploaded model, including its version.

#### Response Example
```json
{
  "version": "v1.2.0",
  "name": "my-org/my-team/pytorch-mnist"
}
```

## load_model - Download and load models into memory

### Description
Downloads a model from Lightning Cloud and automatically loads it into memory based on file extension. Supports .ts (TorchScript), .keras (TensorFlow/Keras), and .pkl (pickled objects).

### Method
`load_model`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **name** (str) - Required - The name of the model to download, including organization, team, and optionally version (e.g., "my-org/my-team/model-name:v1" or "my-org/my-team/model-name:latest").
- **download_dir** (str) - Optional - The directory to download the model artifacts to. If not provided, a temporary directory will be used.
- **progress_bar** (bool) - Optional - Whether to display a progress bar during download. Defaults to False.
- **verbose** (int) - Optional - Verbosity level for the download process. Defaults to 0.

### Request Example
```python
import litmodels as lm

# Load PyTorch TorchScript model
torch_model = lm.load_model(
    name="my-org/my-team/pytorch-mnist:v1",
    download_dir="./models"
)

# Load scikit-learn model and run inference
sklearn_model = lm.load_model(
    name="my-org/my-team/sklearn-iris-svm",
    download_dir="./models"
)
sample_input = [[5.1, 3.5, 1.4, 0.2]]
prediction = sklearn_model.predict(sample_input)
print(f"Predicted class: {prediction[0]}")

# Load Keras model
keras_model = lm.load_model(
    name="my-org/my-team/keras-classifier:latest",
    download_dir="./models"
)
```

### Response
#### Success Response (200)
- **model** (object) - The loaded model object.

#### Response Example
```python
# Output depends on the model type loaded
# For scikit-learn, it would be the fitted estimator object
print(prediction[0]) # Example output for scikit-learn prediction
```

## upload_model - Upload existing checkpoint files

### Description
Uploads a local checkpoint file or directory to Lightning Cloud without serialization. Use this when you already have saved checkpoint files on disk.

### Method
`upload_model`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **name** (str) - Required - The name of the model to upload, including organization and team (e.g., "my-org/my-team/model-name"). Can include a version tag (e.g., "my-org/my-team/model-name:v2").
- **model** (str or list[str]) - Required - The path to the local checkpoint file or directory to upload.
- **progress_bar** (bool) - Optional - Whether to display a progress bar during upload. Defaults to False.
- **verbose** (int) - Optional - Verbosity level for the upload process. Defaults to 0.
- **metadata** (dict) - Optional - A dictionary of metadata to associate with the model checkpoint.

### Request Example
```python
import litmodels as lm
import torch
from lightning.pytorch.demos.boring_classes import BoringModel

# Save checkpoint locally first
model = BoringModel()
torch.save(model.state_dict(), "./checkpoints/model-epoch-10.pt")

# Upload the checkpoint file
info = lm.upload_model(
    name="my-org/my-team/boring-model:epoch-10",
    model="./checkpoints/model-epoch-10.pt",
    progress_bar=True,
    metadata={"epoch": "10", "loss": "0.032"}
)
print(f"Uploaded version: {info.version}")

# Upload entire checkpoint directory
lm.upload_model(
    name="my-org/my-team/complete-experiment",
    model="./checkpoints",
    verbose=2
)
```

### Response
#### Success Response (200)
- **info** (object) - Information about the uploaded model, including its version.

#### Response Example
```json
{
  "version": "v1.2.1",
  "name": "my-org/my-team/boring-model"
}
```

## download_model - Download checkpoint files

### Description
Downloads model artifacts from Lightning Cloud to a local directory. Returns the path(s) to downloaded files.

### Method
`download_model`

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **name** (str) - Required - The name of the model to download, including organization, team, and optionally version (e.g., "my-org/my-team/model-name:v1" or "my-org/my-team/model-name:latest").
- **download_dir** (str) - Optional - The directory to download the model artifacts to. If not provided, a temporary directory will be used.
- **progress_bar** (bool) - Optional - Whether to display a progress bar during download. Defaults to False.
- **verbose** (int) - Optional - Verbosity level for the download process. Defaults to 0.

### Request Example
```python
import litmodels as lm
import torch
from lightning.pytorch.demos.boring_classes import BoringModel

# Download checkpoint file
checkpoint_path = lm.download_model(
    name="my-org/my-team/boring-model:epoch-10",
    download_dir="./downloaded_models",
    progress_bar=True
)
print(f"Downloaded to: {checkpoint_path}")

# Load the checkpoint manually
model = BoringModel()
if isinstance(checkpoint_path, list):
    checkpoint_path = checkpoint_path[0]
model.load_state_dict(torch.load(checkpoint_path))
model.eval()

# Download specific version
paths = lm.download_model(
    name="my-org/my-team/my-model:v2.1",
    download_dir="./models"
)
```

### Response
#### Success Response (200)
- **paths** (str or list[str]) - The path(s) to the downloaded file(s) or directory.

#### Response Example
```json
[
  "./downloaded_models/model-epoch-10.pt"
]
```
```

--------------------------------

### LightningModelCheckpoint: Automatic Cloud Checkpoint Uploads

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Drop-in replacement for PyTorch Lightning's ModelCheckpoint callback. Automatically uploads checkpoints to Lightning Cloud in background threads after each save, preventing training blockage. Configurable with options like model registry name, monitor metric, save count, and file naming.

```python
import torch.utils.data as data
import torchvision as tv
from lightning import Trainer
from lightning.pytorch.demos.boring_classes import BoringModel
from litmodels.integrations import LightningModelCheckpoint

# Prepare dataset
dataset = tv.datasets.MNIST(
    "./data",
    download=True,
    transform=tv.transforms.ToTensor()
)
train_data, val_data = data.random_split(dataset, [55000, 5000])

# Train with automatic cloud uploads
trainer = Trainer(
    max_epochs=10,
    callbacks=[
        LightningModelCheckpoint(
            model_registry="my-org/my-team/mnist-classifier",
            monitor="val_loss",
            save_top_k=3,
            keep_all_uploaded=False,  # Remove old versions from cloud
            clear_all_local=True,     # Delete local files after upload
            filename="epoch-{epoch:02d}-loss-{val_loss:.2f}"
        )
    ],
)
trainer.fit(
    BoringModel(),
    data.DataLoader(train_data, batch_size=256),
    data.DataLoader(val_data, batch_size=256)
)
# All checkpoints are uploaded in background without blocking training
```

--------------------------------

### Upload Existing Checkpoint Files with LitModels

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Uploads a local checkpoint file or directory to Lightning Cloud without requiring serialization. This is useful when models are already saved on disk. Supports uploading individual files or entire directories.

```python
import litmodels as lm
import torch
from lightning.pytorch.demos.boring_classes import BoringModel

# Save checkpoint locally first
model = BoringModel()
torch.save(model.state_dict(), "./checkpoints/model-epoch-10.pt")

# Upload the checkpoint file
info = lm.upload_model(
    name="my-org/my-team/boring-model:epoch-10",
    model="./checkpoints/model-epoch-10.pt",
    progress_bar=True,
    metadata={"epoch": "10", "loss": "0.032"}
)
print(f"Uploaded version: {info.version}")

# Upload entire checkpoint directory
lm.upload_model(
    name="my-org/my-team/complete-experiment",
    model="./checkpoints",
    verbose=2
)
```

--------------------------------

### Save and Load Scikit-learn Model with LitModels

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Demonstrates saving a trained scikit-learn SVC model using LitModels. It involves loading an Iris dataset, splitting it, training an SVC model, and then using `save_model` to persist it.

```python
from sklearn import datasets, model_selection, svm
from litmodels import save_model

# Load example dataset
iris = datasets.load_iris()
X, y = iris.data, iris.target

# Split dataset into training and test sets
X_train, X_test, y_train, y_test = model_selection.train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Train a simple SVC model
model = svm.SVC()
model.fit(X_train, y_train)

# Save the model
save_model(model=model, name="sklearn-iris-svm")

```

--------------------------------

### PickleRegistryMixin: Add Pickle-based Save/Load to Python Classes

Source: https://context7.com/lightning-ai/litmodels/llms.txt

A mixin class that adds `upload_model()` and `download_model()` methods to any Python class using pickle serialization. This allows for easy saving and loading of custom Python objects, including their state, without complex serialization logic.

```python
from litmodels.integrations.mixins import PickleRegistryMixin

class CustomModel(PickleRegistryMixin):
    def __init__(self, num_layers, hidden_size):
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.weights = [i * 0.1 for i in range(num_layers)]

    def predict(self, x):
        return sum(w * x for w in self.weights)

# Create and upload model
model = CustomModel(num_layers=5, hidden_size=128)
model.upload_model(
    name="my-org/my-team/custom-model",
    version="v1.0",
    metadata={"architecture": "custom", "layers": "5"}
)

# Download and use model (no __init__ args needed)
loaded_model = CustomModel.download_model(
    name="my-org/my-team/custom-model",
    version="v1.0"
)
result = loaded_model.predict(2.5)
print(f"Prediction: {result}")
```

--------------------------------

### PyTorchRegistryMixin: Add PyTorch Save/Load to nn.Module

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Mixin class for PyTorch models that saves architecture parameters as JSON and state_dict as .pth. It must be the first in the inheritance order. The `download_model()` method reconstructs the model without requiring constructor arguments, simplifying model loading.

```python
import torch
from litmodels.integrations.mixins import PyTorchRegistryMixin

# IMPORTANT: PyTorchRegistryMixin must be first in inheritance
class ConvNet(PyTorchRegistryMixin, torch.nn.Module):
    def __init__(self, input_channels=3, num_classes=10, hidden_size=128):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(input_channels, 32, kernel_size=3)
        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=3)
        self.fc1 = torch.nn.Linear(64 * 6 * 6, hidden_size)
        self.fc2 = torch.nn.Linear(hidden_size, num_classes)
        self.relu = torch.nn.ReLU()

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = self.relu(self.fc1(x))
        return self.fc2(x)

# Train model
model = ConvNet(input_channels=3, num_classes=10, hidden_size=256)
# ... training code ...

# Upload model (saves __init__ args + state_dict)
model.upload_model(
    name="my-org/my-team/convnet",
    version="v2.0",
    metadata={"dataset": "cifar10", "accuracy": "0.89"}
)

# Download model (automatically reconstructs with correct args)
loaded_model = ConvNet.download_model(
    name="my-org/my-team/convnet",
    version="v2.0",
    torch_load_kwargs={"map_location": "cpu"}
)
loaded_model.eval()
```

--------------------------------

### upload_model_files: Low-level Cloud File Upload

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Provides direct access to upload files or directories to cloud storage with custom metadata. This function is used internally by higher-level upload functions and offers granular control over the upload process, including progress bar and verbosity settings.

```python
from litmodels.io.cloud import upload_model_files

info = upload_model_files(
    name="my-org/my-team/custom-artifacts:v1.0",
    path=["model.pth", "config.json", "vocab.txt"],
    progress_bar=True,
    verbose=2,
    metadata={
        "framework": "pytorch",
        "task": "text-classification",
        "trained_on": "2025-01-15"
    }
)
print(f"Created version: {info.version}")
print(f"Model URL: {info.url}")
```

--------------------------------

### Save and Upload PyTorch, Scikit-learn, and Keras Models with LitModels

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Serializes in-memory models (PyTorch, Keras, scikit-learn, or any picklable object) and uploads them to Lightning Cloud. It automatically detects the model type and uses appropriate serialization formats (.pth, .keras, .pkl).

```python
import torch
import litmodels as lm
from sklearn import datasets, svm, model_selection

# PyTorch model - saves state_dict as .pth
model = torch.nn.Sequential(
    torch.nn.Linear(784, 128),
    torch.nn.ReLU(),
    torch.nn.Linear(128, 10)
)
lm.save_model(
    name="my-org/my-team/pytorch-mnist",
    model=model,
    progress_bar=True,
    metadata={"accuracy": "0.95", "epoch": "10"}
)

# Scikit-learn model - saves as .pkl
iris = datasets.load_iris()
X_train, X_test, y_train, y_test = model_selection.train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=42
)
clf = svm.SVC(kernel='rbf', gamma='auto')
clf.fit(X_train, y_train)
lm.save_model(
    name="my-org/my-team/sklearn-iris-svm",
    model=clf,
    verbose=1
)

# TensorFlow/Keras model - saves as .keras
from tensorflow import keras
keras_model = keras.Sequential([
    keras.layers.Dense(64, activation='relu', input_shape=(784,)),
    keras.layers.Dense(10, activation='softmax')
])
keras_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
lm.save_model(
    name="my-org/my-team/keras-classifier",
    model=keras_model
)
```

--------------------------------

### Save and Load TensorFlow/Keras Model with LitModels

Source: https://github.com/lightning-ai/litmodels/blob/main/README.md

Illustrates saving and loading TensorFlow/Keras models using LitModels. The `save_model` function persists the compiled Keras model, and `load_model` retrieves it, supporting specified download directories.

```python
from tensorflow import keras

from litmodels import save_model

# Define the model
model = keras.Sequential(
    [
        keras.layers.Dense(10, input_shape=(784,), name="dense_1"),
        keras.layers.Dense(10, name="dense_2"),
    ]
)

# Compile the model
model.compile(optimizer="adam", loss="categorical_crossentropy")

# Save the model
save_model("lightning-ai/jirka/sample-tf-keras-model", model=model)
```

```python
from litmodels import load_model

model_ = load_model(
    "lightning-ai/jirka/sample-tf-keras-model", download_dir="./my-model"
)
```

--------------------------------

### Delete Model Version - Python

Source: https://context7.com/lightning-ai/litmodels/llms.txt

Removes a specific model version from Lightning Cloud, aiding in storage management and cleanup. It can delete individual versions or iterate through a list for bulk removal.

```python
from litmodels.io.cloud import delete_model_version

# Delete specific version
delete_model_version(
    name="my-org/my-team/experimental-model",
    version="v0.1-alpha"
)

# Clean up old experimental versions
for version in ["v0.1", "v0.2", "v0.3"]:
    delete_model_version(
        name="my-org/my-team/mnist-model",
        version=version
    )
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.