Writing a fine-tuning plugin¶

A plugin lives in src/Plugin/AiFinetune/ of your module, carries the #[AiFinetune] attribute and implements AiFinetuneInterface. Extend AiFinetuneBase to get sensible defaults for the optional methods.

The attribute¶

#[AiFinetune(
  id: 'krea_lora',
  label: new TranslatableMarkup('Krea style (LoRA)'),
  description: new TranslatableMarkup('Train a custom Krea style from 5-30 images.'),
  ai_provider: 'krea',
  dataset_types: ['image'],
)]

Property	Required	Meaning
`id`	yes	Plugin id, stored on the job's `plugin` field.
`label`	yes	Shown in the provider select and in Drush.
`description`	no	Free text.
`ai_provider`	no	The id of the AI provider plugin this belongs to. Informational.
`dataset_types`	no	`['image']`, `['text']` or both. Defaults to `['image']`. The job form only offers these types and the runner refuses others.

The interface¶

Method	Purpose
`isUsable(): bool`	Whether the provider is configured (API key present). Unusable plugins are listed but cannot start.
`getBaseModels(): array`	`id => label` of the models that can be fine-tuned. Required, the form needs at least one.
`getTrainingTypes(): array`	Optional `id => label`, for example Style or Character. Return `[]` to hide the field.
`getMinimumItems(): int`	Smallest dataset the provider accepts. The runner refuses smaller datasets. Default 1.
`getSupportedDatasetTypes(): array`	Defaults to the attribute's `dataset_types`.
`buildSettingsForm(array $form, FormStateInterface $form_state, array $settings): array`	Provider specific form elements. Whatever the user enters is stored as an array on the job and comes back through `$job->getSettings()`. Return `$form` unchanged for no settings.
`validateSettingsForm(array &$form, FormStateInterface $form_state): void`	Validate those elements.
`startTraining(FinetuneJobInterface $job, FinetuneDataset $dataset): FinetuneStatus`	Upload the dataset, start the remote training, return the remote job id in a status.
`checkStatus(FinetuneJobInterface $job): FinetuneStatus`	Poll the provider using `$job->getRemoteJobId()`. Return the result id once the training is done.
`cancelTraining(FinetuneJobInterface $job): FinetuneStatus`	Cancel remotely. The base class marks the job cancelled locally and says the provider does not support cancelling.
`getUsageInstructions(FinetuneJobInterface $job): array`	A render array shown on the edit form once a result id exists. The base class prints the result id.

Example¶

<?php

declare(strict_types=1);

namespace Drupal\my_provider\Plugin\AiFinetune;

use Drupal\Core\Form\FormStateInterface;
use Drupal\Core\StringTranslation\TranslatableMarkup;
use Drupal\ai_finetuner\AiFinetuneBase;
use Drupal\ai_finetuner\Attribute\AiFinetune;
use Drupal\ai_finetuner\Entity\FinetuneJobInterface;
use Drupal\ai_finetuner\FinetuneDataset;
use Drupal\ai_finetuner\FinetuneStatus;
use Drupal\my_provider\ApiClientFactory;
use Symfony\Component\DependencyInjection\ContainerInterface;

#[AiFinetune(
  id: 'my_provider_lora',
  label: new TranslatableMarkup('My provider LoRA'),
  ai_provider: 'my_provider',
  dataset_types: ['image'],
)]
class MyProviderLora extends AiFinetuneBase {

  public function __construct(array $configuration, $plugin_id, $plugin_definition, protected ApiClientFactory $clientFactory) {
    parent::__construct($configuration, $plugin_id, $plugin_definition);
  }

  public static function create(ContainerInterface $container, array $configuration, $plugin_id, $plugin_definition) {
    return new static($configuration, $plugin_id, $plugin_definition, $container->get('my_provider.client_factory'));
  }

  public function isUsable(): bool {
    return $this->clientFactory->hasApiKey();
  }

  public function getBaseModels(): array {
    return ['flux_dev' => $this->t('Flux Dev')];
  }

  public function getMinimumItems(): int {
    return 5;
  }

  public function buildSettingsForm(array $form, FormStateInterface $form_state, array $settings): array {
    $form['steps'] = [
      '#type' => 'number',
      '#title' => $this->t('Training steps'),
      '#default_value' => $settings['steps'] ?? 500,
      '#min' => 100,
    ];
    return $form;
  }

  public function startTraining(FinetuneJobInterface $job, FinetuneDataset $dataset): FinetuneStatus {
    $client = $this->clientFactory->create();
    $asset_ids = [];
    foreach ($dataset->images as $image) {
      $asset_ids[] = $client->upload($image->getFilename(), $image->getMimeType(), $image->getBinary());
    }
    $response = $client->train([
      'name' => $job->label(),
      'base_model' => $job->getBaseModel(),
      'trigger_word' => $job->getTriggerWord() ?: $job->label(),
      'steps' => $job->getSettings()['steps'] ?? 500,
      'assets' => $asset_ids,
    ]);
    return new FinetuneStatus(FinetuneStatus::PENDING, $response['job_id'], NULL, 'Submitted', $response);
  }

  public function checkStatus(FinetuneJobInterface $job): FinetuneStatus {
    $remote = $this->clientFactory->create()->job($job->getRemoteJobId());
    $status = match ($remote['status']) {
      'queued' => FinetuneStatus::PENDING,
      'training' => FinetuneStatus::RUNNING,
      'done' => FinetuneStatus::COMPLETED,
      'cancelled' => FinetuneStatus::CANCELLED,
      default => FinetuneStatus::FAILED,
    };
    return new FinetuneStatus($status, $job->getRemoteJobId(), $remote['model_id'] ?? NULL, $remote['message'] ?? NULL, $remote);
  }

  public function cancelTraining(FinetuneJobInterface $job): FinetuneStatus {
    $this->clientFactory->create()->cancel($job->getRemoteJobId());
    return new FinetuneStatus(FinetuneStatus::CANCELLED, $job->getRemoteJobId(), NULL, 'Cancelled');
  }

  public function getUsageInstructions(FinetuneJobInterface $job): array {
    return ['#markup' => $this->t('Set model id @id on the provider configuration.', ['@id' => $job->getResultId()])];
  }

}

What the plugin receives¶

$job->getBaseModel(), getTrainingType(), getTriggerWord(), getSettings(): what the user chose in the form.
$job->getRemoteJobId(), getResultId(): what earlier calls returned.
$dataset->type: image or text.
$dataset->images: Drupal\ai\OperationType\GenericType\ImageFile[] with getBinary(), getMimeType(), getFilename().
$dataset->texts: arrays with id, title, text (and entity_type, entity_id when they came from entities). $dataset->toJsonl() gives one JSON object per line.

Do not build datasets in the plugin. The runner already handles uploads, entity displays and Views and hands you a finished FinetuneDataset.

Status mapping¶

Map the provider's statuses onto the FinetuneStatus constants. Anything still in progress must be PENDING or RUNNING, otherwise cron stops polling. COMPLETED should carry the result id; the runner stores it on the job and the UI shows your usage instructions from then on. The optional raw array is stored on the job's result field for debugging.

Throw a \RuntimeException (or one of the AI module's exceptions) on failures. The UI shows the message as an error and Drush prints it.

startTraining() should return as soon as the provider has accepted the job. Long polling belongs in checkStatus(), which cron, the Refresh status operation and drush ai-finetuner:status call.

Verifying¶

drush cr
drush ai-finetuner:providers
drush ai-finetuner:create "Test" --plugin=my_provider_lora --images=a.png,b.png --start --wait

For automated tests, look at tests/modules/ai_finetuner_test, which implements the interface with State instead of an API.