enh(TextToImage): Add getExpectedRuntime to IProvider and run tasks during request lifetime if possible

Signed-off-by: Marcel Klehr <mklehr@gmx.net>

4 files changed, 55 insertions, 3 deletions

diff --git a/core/Controller/TextToImageApiController.php b/core/Controller/TextToImageApiController.php
index 402d2983732..c7878d7cdc3 100644
--- a/core/Controller/TextToImageApiController.php
+++ b/core/Controller/TextToImageApiController.php
@@ -36,6 +36,7 @@ use OCP\AppFramework\Http\Attribute\PublicPage;
 use OCP\AppFramework\Http\Attribute\UserRateLimit;
 use OCP\AppFramework\Http\DataResponse;
 use OCP\AppFramework\Http\FileDisplayResponse;
+use OCP\DB\Exception;
 use OCP\Files\NotFoundException;
 use OCP\IL10N;
 use OCP\IRequest;
@@ -91,7 +92,11 @@ class TextToImageApiController extends \OCP\AppFramework\OCSController {
 	public function schedule(string $input, string $appId, string $identifier = ''): DataResponse {
 		$task = new Task($input, $appId, $this->userId, $identifier);
 		try {
-			$this->textToImageManager->scheduleTask($task);
+			try {
+				$this->textToImageManager->runOrScheduleTask($task);
+			} catch (\RuntimeException) {
+				// noop
+			}
 
 			$json = $task->jsonSerialize();
 
@@ -100,6 +105,8 @@ class TextToImageApiController extends \OCP\AppFramework\OCSController {
 			]);
 		} catch (PreConditionNotMetException) {
 			return new DataResponse(['message' => $this->l->t('No text to image provider is available')], Http::STATUS_PRECONDITION_FAILED);
+		} catch (Exception) {
+			return new DataResponse(['message' => $this->l->t('Internal error')], Http::STATUS_INTERNAL_SERVER_ERROR);
 		}
 	}
 
diff --git a/lib/private/TextToImage/Manager.php b/lib/private/TextToImage/Manager.php
index 181748caf99..c309b7264e1 100644
--- a/lib/private/TextToImage/Manager.php
+++ b/lib/private/TextToImage/Manager.php
@@ -187,12 +187,12 @@ class Manager implements IManager {
 
 	/**
 	 * @inheritDoc
-	 * @throws Exception
 	 */
 	public function scheduleTask(Task $task): void {
 		if (!$this->hasProviders()) {
 			throw new PreConditionNotMetException('No text to image provider is installed that can handle this task');
 		}
+		$this->logger->warning('Scheduling Text2Image Task');
 		$task->setStatus(Task::STATUS_SCHEDULED);
 		$taskEntity = DbTask::fromPublicTask($task);
 		$this->taskMapper->insert($taskEntity);
@@ -205,6 +205,36 @@ class Manager implements IManager {
 	/**
 	 * @inheritDoc
 	 */
+	public function runOrScheduleTask(Task $task) : void {
+		if (!$this->hasProviders()) {
+			throw new PreConditionNotMetException('No text to image provider is installed that can handle this task');
+		}
+		$providers = $this->getProviders();
+
+		$json = $this->config->getAppValue('core', 'ai.text2image_provider', '');
+		if ($json !== '') {
+			try {
+				$className = json_decode($json, true, 512, JSON_THROW_ON_ERROR);
+				$provider = current(array_filter($providers, fn ($provider) => $provider::class === $className));
+				if ($provider !== false) {
+					$providers = [$provider];
+				}
+			} catch (\JsonException $e) {
+				$this->logger->warning('Failed to decode Text2Image setting `ai.text2image_provider`', ['exception' => $e]);
+			}
+		}
+		$maxExecutionTime = (int) ini_get('max_execution_time');
+		// Offload the tttttttask to a background job if the expected runtime of the likely provider is longer than 80% of our max execution time
+		if ($providers[0]->getExpectedRuntime() > $maxExecutionTime * 0.8) {
+			$this->scheduleTask($task);
+			return;
+		}
+		$this->runTask($task);
+	}
+
+	/**
+	 * @inheritDoc
+	 */
 	public function deleteTask(Task $task): void {
 		$taskEntity = DbTask::fromPublicTask($task);
 		$this->taskMapper->delete($taskEntity);
diff --git a/lib/public/TextToImage/IManager.php b/lib/public/TextToImage/IManager.php
index c5c8a0a22a6..cd97312779c 100644
--- a/lib/public/TextToImage/IManager.php
+++ b/lib/public/TextToImage/IManager.php
@@ -26,6 +26,7 @@ declare(strict_types=1);
 
 namespace OCP\TextToImage;
 
+use OCP\DB\Exception;
 use OCP\PreConditionNotMetException;
 use OCP\TextToImage\Exception\TaskNotFoundException;
 use RuntimeException;
@@ -55,12 +56,20 @@ interface IManager {
 	 * If inference fails a \OCP\TextToImage\Events\TaskFailedEvent will be dispatched instead
 	 *
 	 * @param Task $task The task to schedule
-	 * @throws PreConditionNotMetException If no or not the requested provider was registered but this method was still called
+	 * @throws PreConditionNotMetException If no provider was registered but this method was still called
+	 * @throws Exception If there was a problem inserting the task into the database
 	 * @since 28.0.0
 	 */
 	public function scheduleTask(Task $task) : void;
 
 	/**
+	 * @throws Exception if there was a problem inserting the task into the database
+	 * @throws PreConditionNotMetException if no provider is registered
+	 * @throws RuntimeException If the task run fail
+	 */
+	public function runOrScheduleTask(Task $task) : void;
+
+	/**
 	 * Delete a task that has been scheduled before
 	 *
 	 * @param Task $task The task to delete
diff --git a/lib/public/TextToImage/IProvider.php b/lib/public/TextToImage/IProvider.php
index 4fc73089243..12cf39bb713 100644
--- a/lib/public/TextToImage/IProvider.php
+++ b/lib/public/TextToImage/IProvider.php
@@ -49,4 +49,10 @@ interface IProvider {
 	 * @throws RuntimeException If the text could not be processed
 	 */
 	public function generate(string $prompt, $resource): void;
+
+	/**
+	 * The expected runtime for one task with this provider in seconds
+	 * @since 28.0.0
+	 */
+	public function getExpectedRuntime(): int;
 }