enh(TextToImage): Allow generating multiple images with one task

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

2 files changed, 27 insertions, 27 deletions

diff --git a/lib/private/TextToImage/Db/Task.php b/lib/private/TextToImage/Db/Task.php
index bff827533a5..12818d956c2 100644
--- a/lib/private/TextToImage/Db/Task.php
+++ b/lib/private/TextToImage/Db/Task.php
@@ -27,10 +27,6 @@ namespace OC\TextToImage\Db;
 
 use DateTime;
 use OCP\AppFramework\Db\Entity;
-use OCP\Files\AppData\IAppDataFactory;
-use OCP\Files\NotFoundException;
-use OCP\Files\NotPermittedException;
-use OCP\Image;
 use OCP\TextToImage\Task as OCPTask;
 
 /**
@@ -48,6 +44,8 @@ use OCP\TextToImage\Task as OCPTask;
  * @method string getAppId()
  * @method setIdentifier(string $identifier)
  * @method string|null getIdentifier()
+ * @method setNumberOfImages(int $numberOfImages)
+ * @method int getNumberOfImages()
  */
 class Task extends Entity {
 	protected $lastUpdated;
@@ -57,16 +55,17 @@ class Task extends Entity {
 	protected $userId;
 	protected $appId;
 	protected $identifier;
+	protected $numberOfImages;
 
 	/**
 	 * @var string[]
 	 */
-	public static array $columns = ['id', 'last_updated', 'input', 'status', 'user_id', 'app_id', 'identifier'];
+	public static array $columns = ['id', 'last_updated', 'input', 'status', 'user_id', 'app_id', 'identifier', 'number_of_images'];
 
 	/**
 	 * @var string[]
 	 */
-	public static array $fields = ['id', 'lastUpdated', 'input', 'status', 'userId', 'appId', 'identifier'];
+	public static array $fields = ['id', 'lastUpdated', 'input', 'status', 'userId', 'appId', 'identifier', 'numberOfImages'];
 
 
 	public function __construct() {
@@ -78,6 +77,7 @@ class Task extends Entity {
 		$this->addType('userId', 'string');
 		$this->addType('appId', 'string');
 		$this->addType('identifier', 'string');
+		$this->addType('numberOfImages', 'integer');
 	}
 
 	public function toRow(): array {
@@ -92,6 +92,7 @@ class Task extends Entity {
 			'id' => $task->getId(),
 			'lastUpdated' => time(),
 			'status' => $task->getStatus(),
+			'numberOfImages' => $task->getNumberOfImages(),
 			'input' => $task->getInput(),
 			'userId' => $task->getUserId(),
 			'appId' => $task->getAppId(),
@@ -101,20 +102,9 @@ class Task extends Entity {
 	}
 
 	public function toPublicTask(): OCPTask {
-		$task = new OCPTask($this->getInput(), $this->getAppId(), $this->getuserId(), $this->getIdentifier());
+		$task = new OCPTask($this->getInput(), $this->getAppId(), $this->getNumberOfImages(), $this->getuserId(), $this->getIdentifier());
 		$task->setId($this->getId());
 		$task->setStatus($this->getStatus());
-		$appData = \OC::$server->get(IAppDataFactory::class)->get('core');
-		try {
-			try {
-				$folder = $appData->getFolder('text2image');
-			} catch(NotFoundException) {
-				$folder = $appData->newFolder('text2image');
-			}
-			$task->setOutputImage(new Image(base64_encode($folder->getFile((string)$task->getId())->getContent())));
-		} catch (NotFoundException|NotPermittedException) {
-			// noop
-		}
 		return $task;
 	}
 }
diff --git a/lib/private/TextToImage/Manager.php b/lib/private/TextToImage/Manager.php
index c309b7264e1..a48b202239d 100644
--- a/lib/private/TextToImage/Manager.php
+++ b/lib/private/TextToImage/Manager.php
@@ -139,17 +139,27 @@ class Manager implements IManager {
 					$this->logger->debug('Creating folder in appdata for Text2Image results');
 					$folder = $this->appData->newFolder('text2image');
 				}
-				$this->logger->debug('Creating result file for Text2Image task');
-				$file = $folder->newFile((string) $task->getId());
-				$resource = $file->write();
-				if ($resource === false) {
-					throw new RuntimeException('Text2Image generation using provider ' . $provider->getName() . ' failed: Couldn\'t open file to write.');
+				try {
+					$folder = $folder->getFolder((string) $task->getId());
+				} catch(NotFoundException) {
+					$this->logger->debug('Creating new folder in appdata Text2Image results folder');
+					$folder = $this->appData->newFolder((string) $task->getId());
+				}
+				$this->logger->debug('Creating result files for Text2Image task');
+				$resources = [];
+				for ($i = 0; $i < $task->getNumberOfImages(); $i++) {
+					$resources[] = $folder->newFile((string) $i)->write();
+					if ($resource[count($resources) - 1] === false) {
+						throw new RuntimeException('Text2Image generation using provider ' . $provider->getName() . ' failed: Couldn\'t open file to write.');
+					}
 				}
 				$this->logger->debug('Calling Text2Image provider\'s generate method');
-				$provider->generate($task->getInput(), $resource);
-				if (is_resource($resource)) {
-					// If $resource hasn't been closed yet, we'll do that here
-					fclose($resource);
+				$provider->generate($task->getInput(), $resources);
+				for ($i = 0; $i < $task->getNumberOfImages(); $i++) {
+					if (is_resource($resources[$i])) {
+						// If $resource hasn't been closed yet, we'll do that here
+						fclose($resource[$i]);
+					}
 				}
 				$task->setStatus(Task::STATUS_SUCCESSFUL);
 				$this->logger->debug('Updating Text2Image task in DB');