From 8968573d9fa0b9abac1dd5c7684dd94258a080cf Mon Sep 17 00:00:00 2001
From: Marcel Klehr <mklehr@gmx.net>
Date: Fri, 20 Oct 2023 12:17:17 +0200
Subject: enh(TextToImage): Add getExpectedRuntime to IProvider and run tasks
 during request lifetime if possible

Signed-off-by: Marcel Klehr <mklehr@gmx.net>
---
 lib/private/TextToImage/Manager.php  | 32 +++++++++++++++++++++++++++++++-
 lib/public/TextToImage/IManager.php  | 11 ++++++++++-
 lib/public/TextToImage/IProvider.php |  6 ++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

(limited to 'lib')

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);
@@ -202,6 +202,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
 	 */
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,11 +56,19 @@ 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
 	 *
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;
 }
-- 
cgit v1.2.3