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) + } } } }