| 1 | import { Body, Controller, Get, Param, Post } from '@nestjs/common' |
| 2 | import { ApiTags } from '@nestjs/swagger' |
| 3 | import { GetToken, Public, TokenInfo } from '@yikart/aitoearn-auth' |
| 4 | import { ApiDoc, ParseObjectIdPipe, UserType } from '@yikart/common' |
| 5 | import { ImageEditDto, ImageGenerationDto } from './image.dto' |
| 6 | import { ImageService } from './image.service' |
| 7 | import { AsyncTaskResponseVo, ImageEditModelParamsVo, ImageGenerationModelParamsVo, ImageResponseVo, TaskStatusResponseVo } from './image.vo' |
| 8 | |
| 9 | @ApiTags('Me/Ai/Image') |
| 10 | @Controller('ai') |
| 11 | export class ImageController { |
| 12 | constructor(private readonly imageService: ImageService) {} |
| 13 | @ApiDoc({ |
| 14 | summary: 'Get Image Generation Model Parameters', |
| 15 | response: [ImageGenerationModelParamsVo], |
| 16 | }) |
| 17 | @Public() |
| 18 | @Get('/models/image/generation') |
| 19 | async getImageGenerationModels(@GetToken() token?: TokenInfo): Promise<ImageGenerationModelParamsVo[]> { |
| 20 | const response = await this.imageService.generationModelConfig({ |
| 21 | userId: token?.id, |
| 22 | userType: UserType.User, |
| 23 | }) |
| 24 | return response.map(item => ImageGenerationModelParamsVo.create(item)) |
| 25 | } |
| 26 | |
| 27 | @ApiDoc({ |
| 28 | summary: 'Get Image Editing Model Parameters', |
| 29 | response: [ImageEditModelParamsVo], |
| 30 | }) |
| 31 | @Public() |
| 32 | @Get('/models/image/edit') |
| 33 | async getImageEditModels(@GetToken() token?: TokenInfo): Promise<ImageEditModelParamsVo[]> { |
| 34 | const response = await this.imageService.editModelConfig({ |
| 35 | userId: token?.id, |
| 36 | userType: UserType.User, |
| 37 | }) |
| 38 | return response.map(item => ImageEditModelParamsVo.create(item)) |
| 39 | } |
| 40 | |
| 41 | @ApiDoc({ |
| 42 | summary: 'Generate AI Image', |
| 43 | body: ImageGenerationDto.schema, |
| 44 | response: ImageResponseVo, |
| 45 | }) |
| 46 | @Post('/image/generate') |
| 47 | async generateImage( |
| 48 | @GetToken() token: TokenInfo, |
| 49 | @Body() body: ImageGenerationDto, |
| 50 | ): Promise<ImageResponseVo> { |
| 51 | const response = await this.imageService.userGeneration({ |
| 52 | userId: token.id, |
| 53 | userType: UserType.User, |
| 54 | ...body, |
| 55 | }) |
| 56 | return ImageResponseVo.create(response) |
| 57 | } |
| 58 | |
| 59 | @ApiDoc({ |
| 60 | summary: 'Edit AI Image', |
| 61 | body: ImageEditDto.schema, |
| 62 | response: ImageResponseVo, |
| 63 | }) |
| 64 | @Post('/image/edit') |
| 65 | async editImage( |
| 66 | @GetToken() token: TokenInfo, |
| 67 | @Body() body: ImageEditDto, |
| 68 | ): Promise<ImageResponseVo> { |
| 69 | const response = await this.imageService.userEdit({ |
| 70 | userId: token.id, |
| 71 | userType: UserType.User, |
| 72 | ...body, |
| 73 | }) |
| 74 | return ImageResponseVo.create(response) |
| 75 | } |
| 76 | |
| 77 | @ApiDoc({ |
| 78 | summary: 'Generate AI Image Asynchronously', |
| 79 | body: ImageGenerationDto.schema, |
| 80 | response: AsyncTaskResponseVo, |
| 81 | }) |
| 82 | @Post('/image/generate/async') |
| 83 | async generateImageAsync( |
| 84 | @GetToken() token: TokenInfo, |
| 85 | @Body() body: ImageGenerationDto, |
| 86 | ) { |
| 87 | const response = await this.imageService.userGenerationAsync({ |
| 88 | userId: token.id, |
| 89 | userType: UserType.User, |
| 90 | ...body, |
| 91 | }) |
| 92 | return AsyncTaskResponseVo.create(response) |
| 93 | } |
| 94 | |
| 95 | @ApiDoc({ |
| 96 | summary: 'Edit AI Image Asynchronously', |
| 97 | body: ImageEditDto.schema, |
| 98 | response: AsyncTaskResponseVo, |
| 99 | }) |
| 100 | @Post('/image/edit/async') |
| 101 | async editImageAsync( |
| 102 | @GetToken() token: TokenInfo, |
| 103 | @Body() body: ImageEditDto, |
| 104 | ) { |
| 105 | const response = await this.imageService.userEditAsync({ |
| 106 | userId: token.id, |
| 107 | userType: UserType.User, |
| 108 | ...body, |
| 109 | }) |
| 110 | return AsyncTaskResponseVo.create(response) |
| 111 | } |
| 112 | |
| 113 | @ApiDoc({ |
| 114 | summary: 'Get Image Task Status', |
| 115 | response: TaskStatusResponseVo, |
| 116 | }) |
| 117 | @Get('/image/task/:logId') |
| 118 | async getImageTaskStatus( |
| 119 | @GetToken() token: TokenInfo, |
| 120 | @Param('logId', ParseObjectIdPipe) logId: string, |
| 121 | ): Promise<TaskStatusResponseVo> { |
| 122 | const response = await this.imageService.getTaskStatus(logId) |
| 123 | return TaskStatusResponseVo.create(response) |
| 124 | } |
| 125 | } |
| 126 |