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

Signed-off-by: Marcel Klehr <mklehr@gmx.net>
pull/40326/head
Marcel Klehr 7 months ago
parent 8968573d9f
commit b7fd5185b6

@ -80,6 +80,7 @@ class TextToImageApiController extends \OCP\AppFramework\OCSController {
* @param string $input Input text
* @param string $appId ID of the app that will execute the task
* @param string $identifier An arbitrary identifier for the task
* @param int $numberOfImages The number of images to generate
*
* @return DataResponse<Http::STATUS_OK, array{task: CoreTextToImageTask}, array{}>|DataResponse<Http::STATUS_PRECONDITION_FAILED, array{message: string}, array{}>
*
@ -89,8 +90,8 @@ class TextToImageApiController extends \OCP\AppFramework\OCSController {
#[PublicPage]
#[UserRateLimit(limit: 20, period: 120)]
#[AnonRateLimit(limit: 5, period: 120)]
public function schedule(string $input, string $appId, string $identifier = ''): DataResponse {
$task = new Task($input, $appId, $this->userId, $identifier);
public function schedule(string $input, string $appId, string $identifier = '', int $numberOfImages = 8): DataResponse {
$task = new Task($input, $appId, $numberOfImages, $this->userId, $identifier);
try {
try {
$this->textToImageManager->runOrScheduleTask($task);
@ -145,6 +146,7 @@ class TextToImageApiController extends \OCP\AppFramework\OCSController {
* This endpoint allows downloading the resulting image of a task
*
* @param int $id The id of the task
* @param int $index The index of the image to retrieve
*
* @return FileDisplayResponse<Http::STATUS_OK, array{'Content-Type': string}>|DataResponse<Http::STATUS_NOT_FOUND|Http::STATUS_INTERNAL_SERVER_ERROR, array{message: string}, array{}>
*
@ -153,15 +155,17 @@ class TextToImageApiController extends \OCP\AppFramework\OCSController {
*/
#[PublicPage]
#[BruteForceProtection(action: 'text2image')]
public function getImage(int $id): DataResponse|FileDisplayResponse {
public function getImage(int $id, int $index): DataResponse|FileDisplayResponse {
try {
$task = $this->textToImageManager->getUserTask($id, $this->userId);
try {
$folder = $this->appData->getFolder('text2image');
} catch(NotFoundException) {
$folder = $this->appData->newFolder('text2image');
$res = new DataResponse(['message' => $this->l->t('Image not found')], Http::STATUS_NOT_FOUND);
$res->throttle(['action' => 'text2image']);
return $res;
}
$file = $folder->getFile((string)$task->getId());
$file = $folder->getFolder((string) $task->getId())->getFile((string) $index);
$info = getimagesizefromstring($file->getContent());
return new FileDisplayResponse($file, Http::STATUS_OK, ['Content-Type' => image_type_to_mime_type($info[2])]);

@ -152,6 +152,7 @@ namespace OCA\Core;
* appId: string,
* input: string,
* identifier: ?string,
* numberOfImages: int
* }
*/
class ResponseDefinitions {

@ -159,7 +159,7 @@ $application->registerRoutes($this, [
['root' => '/text2image', 'name' => 'TextToImageApi#isAvailable', 'url' => '/is_available', 'verb' => 'GET'],
['root' => '/text2image', 'name' => 'TextToImageApi#schedule', 'url' => '/schedule', 'verb' => 'POST'],
['root' => '/text2image', 'name' => 'TextToImageApi#getTask', 'url' => '/task/{id}', 'verb' => 'GET'],
['root' => '/text2image', 'name' => 'TextToImageApi#getImage', 'url' => '/task/{id}/image', 'verb' => 'GET'],
['root' => '/text2image', 'name' => 'TextToImageApi#getImage', 'url' => '/task/{id}/image/{index}', 'verb' => 'GET'],
['root' => '/text2image', 'name' => 'TextToImageApi#deleteTask', 'url' => '/task/{id}', 'verb' => 'DELETE'],
['root' => '/text2image', 'name' => 'TextToImageApi#listTasksByApp', 'url' => '/tasks/app/{appId}', 'verb' => 'GET'],
],

@ -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;
}
}

@ -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');

@ -43,12 +43,12 @@ interface IProvider {
* Processes a text
*
* @param string $prompt The input text
* @param resource $resource The file resource to write the image to
* @param resource[] $resources The file resources to write the images to
* @return void
* @since 28.0.0
* @throws RuntimeException If the text could not be processed
*/
public function generate(string $prompt, $resource): void;
public function generate(string $prompt, array $resources): void;
/**
* The expected runtime for one task with this provider in seconds

@ -25,7 +25,11 @@ declare(strict_types=1);
namespace OCP\TextToImage;
use OCP\Files\AppData\IAppDataFactory;
use OCP\Files\NotFoundException;
use OCP\Files\NotPermittedException;
use OCP\IImage;
use OCP\Image;
/**
* This is a text to image task
@ -35,8 +39,6 @@ use OCP\IImage;
final class Task implements \JsonSerializable {
protected ?int $id = null;
private ?IImage $image = null;
/**
* @since 28.0.0
*/
@ -66,6 +68,7 @@ final class Task implements \JsonSerializable {
/**
* @param string $input
* @param string $appId
* @param int $numberOfImages
* @param string|null $userId
* @param null|string $identifier An arbitrary identifier for this task. max length: 255 chars
* @since 28.0.0
@ -73,25 +76,36 @@ final class Task implements \JsonSerializable {
final public function __construct(
protected string $input,
protected string $appId,
protected int $numberOfImages,
protected ?string $userId,
protected ?string $identifier = '',
) {
}
/**
* @return IImage|null
* @return IImage[]|null
* @since 28.0.0
*/
final public function getOutputImage(): ?IImage {
return $this->image;
final public function getOutputImages(): ?array {
$appData = \OC::$server->get(IAppDataFactory::class)->get('core');
try {
$folder = $appData->getFolder('text2image')->getFolder((string)$this->getId());
$images = [];
for ($i = 0; $i < $this->getNumberOfImages(); $i++) {
$images[] = new Image(base64_encode($folder->getFile((string) $i)->getContent()));
}
return $images;
} catch (NotFoundException|NotPermittedException) {
return null;
}
}
/**
* @param IImage|null $image
* @return int
* @since 28.0.0
*/
final public function setOutputImage(?IImage $image): void {
$this->image = $image;
final public function getNumberOfImages(): int {
return $this->numberOfImages;
}
/**
@ -159,7 +173,7 @@ final class Task implements \JsonSerializable {
}
/**
* @psalm-return array{id: ?int, status: 0|1|2|3|4, userId: ?string, appId: string, input: string, identifier: ?string}
* @psalm-return array{id: ?int, status: 0|1|2|3|4, userId: ?string, appId: string, input: string, identifier: ?string, numberOfImages: int}
* @since 28.0.0
*/
public function jsonSerialize(): array {
@ -168,6 +182,7 @@ final class Task implements \JsonSerializable {
'status' => $this->getStatus(),
'userId' => $this->getUserId(),
'appId' => $this->getAppId(),
'numberOfImages' => $this->getNumberOfImages(),
'input' => $this->getInput(),
'identifier' => $this->getIdentifier(),
];

Loading…
Cancel
Save