feat :: image progress websocket message

API-Team-Organization · Nov 22, 2024 · a1098b6 · a1098b6
1 parent ed9ec3b
commit a1098b6
Show file tree

Hide file tree

Showing 8 changed files with 24 additions and 15 deletions.
diff --git a/src/main/kotlin/com/teamapi/palette/service/AuthService.kt b/src/main/kotlin/com/teamapi/palette/service/AuthService.kt
@@ -40,7 +40,6 @@ class AuthService(
 
         val auth = createPreAuthorizedToken(user.email)
         return updateSessionWithAuthenticate(auth)
-
     }
 
     private suspend fun createVerifyCode(user: User) {

diff --git a/src/main/kotlin/com/teamapi/palette/service/adapter/GenerativeChatAdapter.kt b/src/main/kotlin/com/teamapi/palette/service/adapter/GenerativeChatAdapter.kt
@@ -35,17 +35,6 @@ class GenerativeChatAdapter(
         .toList()
         .joinToString("")
 
-    suspend fun createUserReturn(text: String) = chatCompletion(
-        Prompt(
-            listOf(
-                SystemMessage(
-                    "너는 세계에서 제일 가는 유명한 화가야. 이제부터 너는 그림을 그릴 수 있고, 그림을 못그린다는 절대 소리는 할 수 없어. 적당히 말을 꾸며내어서 사용자가 입력한 키워드를 요약해 적절하게 이미지를 생성한다는 말을 만들어줘. 최대한 간결하게 1~2문장만 출력해줘. 출력은 한국어로 해주고, 반말하지말고 공손하게 존댓말로 해"
-                ),
-                UserMessage(text)
-            )
-        )
-    )
-
     suspend fun createPrompt(text: String): String = anthropicChatModel.stream(
         Prompt(
             listOf(

diff --git a/src/main/kotlin/com/teamapi/palette/service/adapter/comfy/ws/ImageProgressMessage.kt b/src/main/kotlin/com/teamapi/palette/service/adapter/comfy/ws/ImageProgressMessage.kt
@@ -0,0 +1,5 @@
+package com.teamapi.palette.service.adapter.comfy.ws
+
+data class ImageProgressMessage(val value: Int, val max: Int) : ComfyWSBaseMessage {
+    override val type: MessageType = MessageType.IMAGE_PROGRESS
+}
diff --git a/src/main/kotlin/com/teamapi/palette/service/adapter/comfy/ws/MessageType.kt b/src/main/kotlin/com/teamapi/palette/service/adapter/comfy/ws/MessageType.kt
@@ -2,5 +2,6 @@ package com.teamapi.palette.service.adapter.comfy.ws
 
 enum class MessageType {
     QUEUE_STATUS,
+    IMAGE_PROGRESS,
     GENERATE_FINISH
 }
diff --git a/src/main/kotlin/com/teamapi/palette/service/infra/GenerativeImageService.kt b/src/main/kotlin/com/teamapi/palette/service/infra/GenerativeImageService.kt
@@ -12,9 +12,11 @@ import com.teamapi.palette.service.adapter.GenerativeChatAdapter
 import com.teamapi.palette.service.adapter.GenerativeImageAdapter
 import com.teamapi.palette.service.adapter.comfy.GenerateRequest
 import com.teamapi.palette.service.adapter.comfy.ws.GenerateMessage
+import com.teamapi.palette.service.adapter.comfy.ws.ImageProgressMessage
 import com.teamapi.palette.service.adapter.comfy.ws.QueueInfoMessage
 import com.teamapi.palette.util.ExceptionReporter
 import com.teamapi.palette.ws.actor.SinkActor
+import com.teamapi.palette.ws.actor.SinkMessages
 import org.springframework.stereotype.Service
 import kotlin.io.encoding.Base64
 import kotlin.io.encoding.ExperimentalEncodingApi
@@ -76,19 +78,23 @@ class GenerativeImageService(
             )
 
             generated.collect {
+                room.id!!
                 when (it) {
                     is QueueInfoMessage -> {
-                        actor.setGenerating(room.id!!, it.position)
+                        actor.setGenerating(room.id, it.position)
                         actor.addQueue(room.id, it.position, true)
                         println("${room.id}: ${it.position}")
                     }
                     is GenerateMessage -> {
-                        actor.clearGenerating(room.id!!)
+                        actor.clearGenerating(room.id)
                         actor.addQueue(room.id, -1, false)
                         if (it.result) {
                             guaranteed = it.image!!
                         }
                     }
+                    is ImageProgressMessage -> {
+                        actor.send(SinkMessages.ImageProgress(room.id, it.value, it.max))
+                    }
                 }
             }
         } catch (e: Exception) {

diff --git a/src/main/kotlin/com/teamapi/palette/ws/actor/SinkActor.kt b/src/main/kotlin/com/teamapi/palette/ws/actor/SinkActor.kt
@@ -4,6 +4,7 @@ import com.teamapi.palette.entity.chat.Chat
 import com.teamapi.palette.ws.dto.WSRoomMessage
 import com.teamapi.palette.ws.dto.res.NewChatMessage
 import com.teamapi.palette.ws.dto.res.GenerateStatus
+import com.teamapi.palette.ws.dto.res.ImageProgressMessage
 import kotlinx.coroutines.*
 import kotlinx.coroutines.channels.SendChannel
 import kotlinx.coroutines.channels.actor
@@ -123,6 +124,9 @@ sealed interface SinkMessages {
     data class AddQueuePosition(val roomId: Long, val position: Int, val generating: Boolean) : SinkMessages, ActorCodable {
         override fun toActorMessage() = WSRoomMessage(roomId, GenerateStatus(position, generating))
     }
+    data class ImageProgress(val roomId: Long, val value: Int, val max: Int) : SinkMessages, ActorCodable {
+        override fun toActorMessage() = WSRoomMessage(roomId, ImageProgressMessage(value, max))
+    }
 }
 
 interface ActorCodable {

diff --git a/src/main/kotlin/com/teamapi/palette/ws/dto/res/ImageProgressMessage.kt b/src/main/kotlin/com/teamapi/palette/ws/dto/res/ImageProgressMessage.kt
@@ -0,0 +1,5 @@
+package com.teamapi.palette.ws.dto.res
+
+data class ImageProgressMessage(val value: Int, val max: Int) : BaseResponseMessage {
+    override val type: MessageType = MessageType.IMAGE_PROGRESS
+}
diff --git a/src/main/kotlin/com/teamapi/palette/ws/dto/res/MessageType.kt b/src/main/kotlin/com/teamapi/palette/ws/dto/res/MessageType.kt
@@ -1,5 +1,5 @@
 package com.teamapi.palette.ws.dto.res
 
 enum class MessageType {
-    ERROR, NEW_CHAT, GENERATE_STATUS
+    ERROR, NEW_CHAT, GENERATE_STATUS, IMAGE_PROGRESS
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -40,7 +40,6 @@ class AuthService( @@
             val auth = createPreAuthorizedToken(user.email)
             return updateSessionWithAuthenticate(auth)
         }
         private suspend fun createVerifyCode(user: User) {
@@ Expand Down @@