From 03e30553046eafaef55e0ca6b96745304312f67d Mon Sep 17 00:00:00 2001 From: Aleksandr Kukharenko Date: Thu, 9 Jan 2025 17:28:07 +0200 Subject: [PATCH] handle remote prompt for prodia v2 image --- proxy-router/internal/proxyapi/proxy_sender.go | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/proxy-router/internal/proxyapi/proxy_sender.go b/proxy-router/internal/proxyapi/proxy_sender.go index b6040c7e..876c55bc 100644 --- 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) + } } } }