1
0

Manager.php 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341
  1. <?php
  2. declare(strict_types=1);
  3. /**
  4. * @copyright Copyright (c) 2023 Marcel Klehr <mklehr@gmx.net>
  5. *
  6. * @author Marcel Klehr <mklehr@gmx.net>
  7. *
  8. * @license GNU AGPL version 3 or any later version
  9. *
  10. * This program is free software: you can redistribute it and/or modify
  11. * it under the terms of the GNU Affero General Public License as
  12. * published by the Free Software Foundation, either version 3 of the
  13. * License, or (at your option) any later version.
  14. *
  15. * This program is distributed in the hope that it will be useful,
  16. * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17. * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  18. * GNU Affero General Public License for more details.
  19. *
  20. * You should have received a copy of the GNU Affero General Public License
  21. * along with this program. If not, see <http://www.gnu.org/licenses/>.
  22. */
  23. namespace OC\TextToImage;
  24. use OC\AppFramework\Bootstrap\Coordinator;
  25. use OC\TextToImage\Db\Task as DbTask;
  26. use OC\TextToImage\Db\TaskMapper;
  27. use OCP\AppFramework\Db\DoesNotExistException;
  28. use OCP\AppFramework\Db\MultipleObjectsReturnedException;
  29. use OCP\BackgroundJob\IJobList;
  30. use OCP\DB\Exception;
  31. use OCP\Files\AppData\IAppDataFactory;
  32. use OCP\Files\IAppData;
  33. use OCP\Files\NotFoundException;
  34. use OCP\Files\NotPermittedException;
  35. use OCP\IConfig;
  36. use OCP\IServerContainer;
  37. use OCP\PreConditionNotMetException;
  38. use OCP\TextToImage\Exception\TaskFailureException;
  39. use OCP\TextToImage\Exception\TaskNotFoundException;
  40. use OCP\TextToImage\IManager;
  41. use OCP\TextToImage\IProvider;
  42. use OCP\TextToImage\IProviderWithUserId;
  43. use OCP\TextToImage\Task;
  44. use Psr\Log\LoggerInterface;
  45. use RuntimeException;
  46. use Throwable;
  47. class Manager implements IManager {
  48. /** @var ?list<IProvider> */
  49. private ?array $providers = null;
  50. private IAppData $appData;
  51. public function __construct(
  52. private IServerContainer $serverContainer,
  53. private Coordinator $coordinator,
  54. private LoggerInterface $logger,
  55. private IJobList $jobList,
  56. private TaskMapper $taskMapper,
  57. private IConfig $config,
  58. IAppDataFactory $appDataFactory,
  59. ) {
  60. $this->appData = $appDataFactory->get('core');
  61. }
  62. /**
  63. * @inheritDoc
  64. */
  65. public function getProviders(): array {
  66. $context = $this->coordinator->getRegistrationContext();
  67. if ($context === null) {
  68. return [];
  69. }
  70. if ($this->providers !== null) {
  71. return $this->providers;
  72. }
  73. $this->providers = [];
  74. foreach ($context->getTextToImageProviders() as $providerServiceRegistration) {
  75. $class = $providerServiceRegistration->getService();
  76. try {
  77. /** @var IProvider $provider */
  78. $provider = $this->serverContainer->get($class);
  79. $this->providers[] = $provider;
  80. } catch (Throwable $e) {
  81. $this->logger->error('Failed to load Text to image provider ' . $class, [
  82. 'exception' => $e,
  83. ]);
  84. }
  85. }
  86. return $this->providers;
  87. }
  88. /**
  89. * @inheritDoc
  90. */
  91. public function hasProviders(): bool {
  92. $context = $this->coordinator->getRegistrationContext();
  93. if ($context === null) {
  94. return false;
  95. }
  96. return count($context->getTextToImageProviders()) > 0;
  97. }
  98. /**
  99. * @inheritDoc
  100. */
  101. public function runTask(Task $task): void {
  102. $this->logger->debug('Running TextToImage Task');
  103. if (!$this->hasProviders()) {
  104. throw new PreConditionNotMetException('No text to image provider is installed that can handle this task');
  105. }
  106. $providers = $this->getPreferredProviders();
  107. foreach ($providers as $provider) {
  108. $this->logger->debug('Trying to run Text2Image provider '.$provider::class);
  109. try {
  110. $task->setStatus(Task::STATUS_RUNNING);
  111. $completionExpectedAt = new \DateTime('now');
  112. $completionExpectedAt->add(new \DateInterval('PT'.$provider->getExpectedRuntime().'S'));
  113. $task->setCompletionExpectedAt($completionExpectedAt);
  114. if ($task->getId() === null) {
  115. $this->logger->debug('Inserting Text2Image task into DB');
  116. $taskEntity = $this->taskMapper->insert(DbTask::fromPublicTask($task));
  117. $task->setId($taskEntity->getId());
  118. } else {
  119. $this->logger->debug('Updating Text2Image task in DB');
  120. $this->taskMapper->update(DbTask::fromPublicTask($task));
  121. }
  122. try {
  123. $folder = $this->appData->getFolder('text2image');
  124. } catch(NotFoundException) {
  125. $this->logger->debug('Creating folder in appdata for Text2Image results');
  126. $folder = $this->appData->newFolder('text2image');
  127. }
  128. try {
  129. $folder = $folder->getFolder((string) $task->getId());
  130. } catch(NotFoundException) {
  131. $this->logger->debug('Creating new folder in appdata Text2Image results folder');
  132. $folder = $folder->newFolder((string) $task->getId());
  133. }
  134. $this->logger->debug('Creating result files for Text2Image task');
  135. $resources = [];
  136. $files = [];
  137. for ($i = 0; $i < $task->getNumberOfImages(); $i++) {
  138. $file = $folder->newFile((string) $i);
  139. $files[] = $file;
  140. $resource = $file->write();
  141. if ($resource !== false && $resource !== true && is_resource($resource)) {
  142. $resources[] = $resource;
  143. } else {
  144. throw new RuntimeException('Text2Image generation using provider "' . $provider->getName() . '" failed: Couldn\'t open file to write.');
  145. }
  146. }
  147. $this->logger->debug('Calling Text2Image provider\'s generate method');
  148. if ($provider instanceof IProviderWithUserId) {
  149. $provider->setUserId($task->getUserId());
  150. }
  151. $provider->generate($task->getInput(), $resources);
  152. for ($i = 0; $i < $task->getNumberOfImages(); $i++) {
  153. if (is_resource($resources[$i])) {
  154. // If $resource hasn't been closed yet, we'll do that here
  155. fclose($resources[$i]);
  156. }
  157. }
  158. $task->setStatus(Task::STATUS_SUCCESSFUL);
  159. $this->logger->debug('Updating Text2Image task in DB');
  160. $this->taskMapper->update(DbTask::fromPublicTask($task));
  161. return;
  162. } catch (\RuntimeException|\Throwable $e) {
  163. for ($i = 0; $i < $task->getNumberOfImages(); $i++) {
  164. if (isset($files, $files[$i])) {
  165. try {
  166. $files[$i]->delete();
  167. } catch(NotPermittedException $e) {
  168. $this->logger->warning('Failed to clean up Text2Image result file after error', ['exception' => $e]);
  169. }
  170. }
  171. }
  172. $this->logger->info('Text2Image generation using provider "' . $provider->getName() . '" failed', ['exception' => $e]);
  173. $task->setStatus(Task::STATUS_FAILED);
  174. try {
  175. $this->taskMapper->update(DbTask::fromPublicTask($task));
  176. } catch (Exception $e) {
  177. $this->logger->warning('Failed to update database after Text2Image error', ['exception' => $e]);
  178. }
  179. throw new TaskFailureException('Text2Image generation using provider "' . $provider->getName() . '" failed: ' . $e->getMessage(), 0, $e);
  180. }
  181. }
  182. $task->setStatus(Task::STATUS_FAILED);
  183. try {
  184. $this->taskMapper->update(DbTask::fromPublicTask($task));
  185. } catch (Exception $e) {
  186. $this->logger->warning('Failed to update database after Text2Image error', ['exception' => $e]);
  187. }
  188. throw new TaskFailureException('Could not run task');
  189. }
  190. /**
  191. * @inheritDoc
  192. */
  193. public function scheduleTask(Task $task): void {
  194. if (!$this->hasProviders()) {
  195. throw new PreConditionNotMetException('No text to image provider is installed that can handle this task');
  196. }
  197. $this->logger->debug('Scheduling Text2Image Task');
  198. $task->setStatus(Task::STATUS_SCHEDULED);
  199. $completionExpectedAt = new \DateTime('now');
  200. $completionExpectedAt->add(new \DateInterval('PT'.$this->getPreferredProviders()[0]->getExpectedRuntime().'S'));
  201. $task->setCompletionExpectedAt($completionExpectedAt);
  202. $taskEntity = DbTask::fromPublicTask($task);
  203. $this->taskMapper->insert($taskEntity);
  204. $task->setId($taskEntity->getId());
  205. $this->jobList->add(TaskBackgroundJob::class, [
  206. 'taskId' => $task->getId()
  207. ]);
  208. }
  209. /**
  210. * @inheritDoc
  211. */
  212. public function runOrScheduleTask(Task $task) : void {
  213. if (!$this->hasProviders()) {
  214. throw new PreConditionNotMetException('No text to image provider is installed that can handle this task');
  215. }
  216. $providers = $this->getPreferredProviders();
  217. $maxExecutionTime = (int) ini_get('max_execution_time');
  218. // Offload the task to a background job if the expected runtime of the likely provider is longer than 80% of our max execution time
  219. if ($providers[0]->getExpectedRuntime() > $maxExecutionTime * 0.8) {
  220. $this->scheduleTask($task);
  221. return;
  222. }
  223. $this->runTask($task);
  224. }
  225. /**
  226. * @inheritDoc
  227. */
  228. public function deleteTask(Task $task): void {
  229. $taskEntity = DbTask::fromPublicTask($task);
  230. $this->taskMapper->delete($taskEntity);
  231. $this->jobList->remove(TaskBackgroundJob::class, [
  232. 'taskId' => $task->getId()
  233. ]);
  234. }
  235. /**
  236. * Get a task from its id
  237. *
  238. * @param int $id The id of the task
  239. * @return Task
  240. * @throws RuntimeException If the query failed
  241. * @throws TaskNotFoundException If the task could not be found
  242. */
  243. public function getTask(int $id): Task {
  244. try {
  245. $taskEntity = $this->taskMapper->find($id);
  246. return $taskEntity->toPublicTask();
  247. } catch (DoesNotExistException $e) {
  248. throw new TaskNotFoundException('Could not find task with the provided id');
  249. } catch (MultipleObjectsReturnedException $e) {
  250. throw new RuntimeException('Could not uniquely identify task with given id', 0, $e);
  251. } catch (Exception $e) {
  252. throw new RuntimeException('Failure while trying to find task by id: ' . $e->getMessage(), 0, $e);
  253. }
  254. }
  255. /**
  256. * Get a task from its user id and task id
  257. * If userId is null, this can only get a task that was scheduled anonymously
  258. *
  259. * @param int $id The id of the task
  260. * @param string|null $userId The user id that scheduled the task
  261. * @return Task
  262. * @throws RuntimeException If the query failed
  263. * @throws TaskNotFoundException If the task could not be found
  264. */
  265. public function getUserTask(int $id, ?string $userId): Task {
  266. try {
  267. $taskEntity = $this->taskMapper->findByIdAndUser($id, $userId);
  268. return $taskEntity->toPublicTask();
  269. } catch (DoesNotExistException $e) {
  270. throw new TaskNotFoundException('Could not find task with the provided id and user id');
  271. } catch (MultipleObjectsReturnedException $e) {
  272. throw new RuntimeException('Could not uniquely identify task with given id and user id', 0, $e);
  273. } catch (Exception $e) {
  274. throw new RuntimeException('Failure while trying to find task by id and user id: ' . $e->getMessage(), 0, $e);
  275. }
  276. }
  277. /**
  278. * Get a list of tasks scheduled by a specific user for a specific app
  279. * and optionally with a specific identifier.
  280. * This cannot be used to get anonymously scheduled tasks
  281. *
  282. * @param string $userId
  283. * @param string $appId
  284. * @param string|null $identifier
  285. * @return Task[]
  286. * @throws RuntimeException
  287. */
  288. public function getUserTasksByApp(?string $userId, string $appId, ?string $identifier = null): array {
  289. try {
  290. $taskEntities = $this->taskMapper->findUserTasksByApp($userId, $appId, $identifier);
  291. return array_map(static function (DbTask $taskEntity) {
  292. return $taskEntity->toPublicTask();
  293. }, $taskEntities);
  294. } catch (Exception $e) {
  295. throw new RuntimeException('Failure while trying to find tasks by appId and identifier: ' . $e->getMessage(), 0, $e);
  296. }
  297. }
  298. /**
  299. * @return list<IProvider>
  300. */
  301. private function getPreferredProviders() {
  302. $providers = $this->getProviders();
  303. $json = $this->config->getAppValue('core', 'ai.text2image_provider', '');
  304. if ($json !== '') {
  305. try {
  306. $id = json_decode($json, true, 512, JSON_THROW_ON_ERROR);
  307. $provider = current(array_filter($providers, fn ($provider) => $provider->getId() === $id));
  308. if ($provider !== false && $provider !== null) {
  309. $providers = [$provider];
  310. }
  311. } catch (\JsonException $e) {
  312. $this->logger->warning('Failed to decode Text2Image setting `ai.text2image_provider`', ['exception' => $e]);
  313. }
  314. }
  315. return $providers;
  316. }
  317. }