handle remote prompt for prodia v2 image

MorpheusAIs · Jan 9, 2025 · 03e3055 · 03e3055
1 parent d50a0a4
commit 03e3055
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/proxy-router/internal/proxyapi/proxy_sender.go b/proxy-router/internal/proxyapi/proxy_sender.go
@@ -603,7 +603,15 @@ func (p *ProxyServiceSender) rpcRequestStreamV2(
 					responses = append(responses, videoGenerationResult)
 					chunk = gcs.NewChunkVideo(&videoGenerationResult)
 				} else {
-					return nil, ttftMs, totalTokens, lib.WrapError(ErrInvalidResponse, err)
+					var imageGenerationResult gcs.ImageRawContentResult
+					err = json.Unmarshal(aiResponse, &imageGenerationResult)
+					if err == nil && imageGenerationResult.ImageRawContent != "" {
+						totalTokens += 1
+						responses = append(responses, imageGenerationResult)
+						chunk = gcs.NewChunkImageRawContent(&imageGenerationResult)
+					} else {
+						return nil, ttftMs, totalTokens, lib.WrapError(ErrInvalidResponse, err)
+					}
 				}
 			}
 		}