first commit

2026-04-26 21:52:23 +03:00
commit 880f412e2c
2662 changed files with 866266 additions and 0 deletions
--- a/core/providers/huggingface/images.go
+++ b/core/providers/huggingface/images.go
@@ -0,0 +1,603 @@
+package huggingface
+
+import (
+	"encoding/base64"
+	"fmt"
+	"net/http"
+	"strconv"
+	"strings"
+
+	"github.com/bytedance/sonic"
+	nebiusProvider "github.com/maximhq/bifrost/core/providers/nebius"
+	providerUtils "github.com/maximhq/bifrost/core/providers/utils"
+	schemas "github.com/maximhq/bifrost/core/schemas"
+)
+
+// Models that support multiple images (image_urls)
+var falAIMultiImageEditModels = map[string]bool{
+	"fal-ai/flux-2/edit":     true,
+	"fal-ai/flux-2-pro/edit": true,
+}
+
+// Models that only support single image (image_url)
+var falAISingleImageEditModels = map[string]bool{
+	"fal-ai/flux-pro/kontext":        true,
+	"fal-ai/flux/dev/image-to-image": true,
+}
+
+// ToHuggingFaceImageGenerationRequest converts a Bifrost image generation request to provider-specific format
+func ToHuggingFaceImageGenerationRequest(bifrostReq *schemas.BifrostImageGenerationRequest) (providerUtils.RequestBodyWithExtraParams, error) {
+	if bifrostReq == nil || bifrostReq.Input == nil {
+		return nil, fmt.Errorf("bifrost request is nil or input is nil")
+	}
+
+	inferenceProvider, model, nameErr := splitIntoModelProvider(bifrostReq.Model)
+	if nameErr != nil {
+		return nil, nameErr
+	}
+
+	switch inferenceProvider {
+	case nebius:
+		req := &nebiusProvider.NebiusImageGenerationRequest{
+			Model:  &model,
+			Prompt: &bifrostReq.Input.Prompt,
+		}
+
+		if bifrostReq.Params != nil {
+			if bifrostReq.Params.ResponseFormat != nil {
+				req.ResponseFormat = bifrostReq.Params.ResponseFormat
+			}
+
+			if bifrostReq.Params.Size != nil && strings.ToLower(*bifrostReq.Params.Size) != "auto" {
+				size := strings.Split(strings.ToLower(*bifrostReq.Params.Size), "x")
+				if len(size) != 2 {
+					return nil, fmt.Errorf("invalid size format: expected 'WIDTHxHEIGHT', got %q", *bifrostReq.Params.Size)
+				}
+
+				width, err := strconv.Atoi(size[0])
+				if err != nil {
+					return nil, fmt.Errorf("invalid width in size %q: %w", *bifrostReq.Params.Size, err)
+				}
+
+				height, err := strconv.Atoi(size[1])
+				if err != nil {
+					return nil, fmt.Errorf("invalid height in size %q: %w", *bifrostReq.Params.Size, err)
+				}
+
+				req.Width = &width
+				req.Height = &height
+			}
+			if bifrostReq.Params.OutputFormat != nil {
+				req.ResponseExtension = bifrostReq.Params.OutputFormat
+			}
+
+			// Handle nebius inconsistency - normalize ResponseExtension case-insensitively
+			if req.ResponseExtension != nil && strings.ToLower(*req.ResponseExtension) == "jpeg" {
+				req.ResponseExtension = schemas.Ptr("jpg")
+			}
+
+			// Map seed from direct field
+			if bifrostReq.Params.Seed != nil {
+				req.Seed = bifrostReq.Params.Seed
+			}
+
+			// Map negative_prompt from direct field
+			if bifrostReq.Params.NegativePrompt != nil {
+				req.NegativePrompt = bifrostReq.Params.NegativePrompt
+			}
+
+			// Handle extra params for nebius
+			if bifrostReq.Params.ExtraParams != nil {
+				req.ExtraParams = bifrostReq.Params.ExtraParams
+				// Map num_inference_steps
+				if v, ok := schemas.SafeExtractIntPointer(bifrostReq.Params.ExtraParams["num_inference_steps"]); ok {
+					delete(req.ExtraParams, "num_inference_steps")
+					req.NumInferenceSteps = v
+				}
+
+				// Map guidance_scale
+				if v, ok := schemas.SafeExtractIntPointer(bifrostReq.Params.ExtraParams["guidance_scale"]); ok {
+					delete(req.ExtraParams, "guidance_scale")
+					req.GuidanceScale = v
+				}
+
+				// Map loras
+				if lorasValue, exists := bifrostReq.Params.ExtraParams["loras"]; exists && lorasValue != nil {
+					delete(req.ExtraParams, "loras")
+					if lorasArray, ok := lorasValue.([]interface{}); ok {
+						for _, item := range lorasArray {
+							if loraMap, ok := item.(map[string]interface{}); ok {
+								if url, ok := schemas.SafeExtractString(loraMap["url"]); ok {
+									if scale, ok := schemas.SafeExtractInt(loraMap["scale"]); ok {
+										req.Loras = append(req.Loras, nebiusProvider.NebiusLora{URL: url, Scale: scale})
+									}
+								}
+							}
+						}
+					}
+				}
+			}
+		}
+		return req, nil
+
+	case hfInference:
+		req := &HuggingFaceHFInferenceImageGenerationRequest{
+			Inputs: bifrostReq.Input.Prompt,
+		}
+		if bifrostReq.Params != nil {
+			req.ExtraParams = bifrostReq.Params.ExtraParams
+		}
+		return req, nil
+
+	case falAI:
+		req := &HuggingFaceFalAIImageGenerationRequest{
+			Prompt: bifrostReq.Input.Prompt,
+		}
+
+		if bifrostReq.Params != nil {
+			// Map n to num_images for fal-ai
+			if bifrostReq.Params.N != nil {
+				req.NumImages = bifrostReq.Params.N
+			}
+
+			// Pass through response_format
+			if bifrostReq.Params.ResponseFormat != nil {
+				req.ResponseFormat = bifrostReq.Params.ResponseFormat
+			}
+
+			// Pass through output_format
+			if bifrostReq.Params.OutputFormat != nil {
+				if strings.ToLower(*bifrostReq.Params.OutputFormat) == "jpg" {
+					req.OutputFormat = schemas.Ptr("jpeg")
+				} else {
+					req.OutputFormat = bifrostReq.Params.OutputFormat
+				}
+			}
+
+			// Convert size from "WxH" format to fal-ai's image_size object
+			if bifrostReq.Params.Size != nil && strings.ToLower(*bifrostReq.Params.Size) != "auto" {
+				size := strings.Split(*bifrostReq.Params.Size, "x")
+				if len(size) == 2 {
+					width, err := strconv.Atoi(size[0])
+					if err == nil {
+						height, err := strconv.Atoi(size[1])
+						if err == nil {
+							req.ImageSize = &HuggingFaceFalAISize{
+								Width:  width,
+								Height: height,
+							}
+						}
+					}
+				}
+			}
+
+			if bifrostReq.Params.ResponseFormat != nil && *bifrostReq.Params.ResponseFormat == "b64_json" {
+				req.SyncMode = schemas.Ptr(true)
+			}
+
+			if bifrostReq.Params.Moderation != nil && *bifrostReq.Params.Moderation == "low" {
+				req.EnableSafetyChecker = schemas.Ptr(false)
+			}
+
+			// Map seed from direct field
+			if bifrostReq.Params.Seed != nil {
+				req.Seed = bifrostReq.Params.Seed
+			}
+
+			// Map negative_prompt from direct field
+			if bifrostReq.Params.NegativePrompt != nil {
+				req.NegativePrompt = bifrostReq.Params.NegativePrompt
+			}
+
+			// Map num_inference_steps from direct field
+			if bifrostReq.Params.NumInferenceSteps != nil {
+				req.NumInferenceSteps = bifrostReq.Params.NumInferenceSteps
+			}
+
+			// Parse fal-ai specific params from ExtraParams
+			if bifrostReq.Params.ExtraParams != nil {
+				req.ExtraParams = bifrostReq.Params.ExtraParams
+				// Map guidance_scale
+				if v, ok := schemas.SafeExtractFloat64Pointer(bifrostReq.Params.ExtraParams["guidance_scale"]); ok {
+					delete(req.ExtraParams, "guidance_scale")
+					req.GuidanceScale = v
+				}
+
+				// Map acceleration
+				if v, ok := schemas.SafeExtractStringPointer(bifrostReq.Params.ExtraParams["acceleration"]); ok {
+					delete(req.ExtraParams, "acceleration")
+					req.Acceleration = v
+				}
+
+				// Map enable_prompt_expansion
+				if v, ok := schemas.SafeExtractBoolPointer(bifrostReq.Params.ExtraParams["enable_prompt_expansion"]); ok {
+					delete(req.ExtraParams, "enable_prompt_expansion")
+					req.EnablePromptExpansion = v
+				}
+
+				// Map enable_safety_checker
+				if v, ok := schemas.SafeExtractBoolPointer(bifrostReq.Params.ExtraParams["enable_safety_checker"]); ok {
+					delete(req.ExtraParams, "enable_safety_checker")
+					req.EnableSafetyChecker = v
+				}
+			}
+		}
+		return req, nil
+
+	case together:
+		req := &HuggingFaceTogetherImageGenerationRequest{
+			Prompt: bifrostReq.Input.Prompt,
+			Model:  model,
+		}
+
+		if bifrostReq.Params != nil {
+			req.ExtraParams = bifrostReq.Params.ExtraParams
+			if bifrostReq.Params.ResponseFormat != nil {
+				req.ResponseFormat = bifrostReq.Params.ResponseFormat
+			}
+
+			if bifrostReq.Params.Size != nil {
+				req.Size = bifrostReq.Params.Size
+			}
+
+			if bifrostReq.Params.N != nil {
+				req.N = bifrostReq.Params.N
+			}
+			if bifrostReq.Params.ResponseFormat != nil && *bifrostReq.Params.ResponseFormat == "b64_json" {
+				req.ResponseFormat = schemas.Ptr("base64")
+			}
+			if bifrostReq.Params.NumInferenceSteps != nil {
+				req.Steps = bifrostReq.Params.NumInferenceSteps
+			}
+		}
+		return req, nil
+
+	default:
+		return nil, fmt.Errorf("unsupported inference provider for image generation: %s", inferenceProvider)
+	}
+}
+
+// ToHuggingFaceImageStreamRequest converts a Bifrost image generation request to fal-ai streaming format
+func ToHuggingFaceImageStreamRequest(bifrostReq *schemas.BifrostImageGenerationRequest) (*HuggingFaceFalAIImageStreamRequest, error) {
+	if bifrostReq == nil || bifrostReq.Input == nil {
+		return nil, fmt.Errorf("bifrost request is nil or input is nil")
+	}
+
+	req := &HuggingFaceFalAIImageStreamRequest{
+		Prompt: bifrostReq.Input.Prompt,
+	}
+
+	if bifrostReq.Params != nil {
+		req.ExtraParams = bifrostReq.Params.ExtraParams
+		// Map n to num_images for fal-ai
+		if bifrostReq.Params.N != nil {
+			req.NumImages = bifrostReq.Params.N
+		}
+
+		// Pass through response_format
+		if bifrostReq.Params.ResponseFormat != nil {
+			req.ResponseFormat = bifrostReq.Params.ResponseFormat
+		}
+
+		// Pass through output_format
+		// Convert "jpg" to "jpeg" for fal-ai (fal-ai only accepts "jpeg", "png", "webp")
+		if bifrostReq.Params.OutputFormat != nil {
+			if strings.ToLower(*bifrostReq.Params.OutputFormat) == "jpg" {
+				req.OutputFormat = schemas.Ptr("jpeg")
+			} else {
+				req.OutputFormat = bifrostReq.Params.OutputFormat
+			}
+		}
+
+		// Convert size from "WxH" format to fal-ai's image_size object
+		if bifrostReq.Params.Size != nil && strings.ToLower(*bifrostReq.Params.Size) != "auto" {
+			size := strings.Split(*bifrostReq.Params.Size, "x")
+			if len(size) == 2 {
+				width, err := strconv.Atoi(size[0])
+				if err == nil {
+					height, err := strconv.Atoi(size[1])
+					if err == nil {
+						req.ImageSize = &HuggingFaceFalAISize{
+							Width:  width,
+							Height: height,
+						}
+					}
+				}
+			}
+		}
+		if bifrostReq.Params.Seed != nil {
+			req.Seed = bifrostReq.Params.Seed
+		}
+		if bifrostReq.Params.NumInferenceSteps != nil {
+			req.NumInferenceSteps = bifrostReq.Params.NumInferenceSteps
+		}
+		if bifrostReq.Params.ResponseFormat != nil && *bifrostReq.Params.ResponseFormat == "b64_json" {
+			req.SyncMode = schemas.Ptr(true)
+		}
+		if bifrostReq.Params.Moderation != nil && *bifrostReq.Params.Moderation == "low" {
+			req.EnableSafetyChecker = schemas.Ptr(false)
+		}
+
+		// Parse fal-ai specific params from ExtraParams
+		if bifrostReq.Params.ExtraParams != nil {
+			if v, ok := schemas.SafeExtractFloat64Pointer(bifrostReq.Params.ExtraParams["guidance_scale"]); ok {
+				delete(req.ExtraParams, "guidance_scale")
+				req.GuidanceScale = v
+			}
+			if v, ok := schemas.SafeExtractStringPointer(bifrostReq.Params.ExtraParams["acceleration"]); ok {
+				delete(req.ExtraParams, "acceleration")
+				req.Acceleration = v
+			}
+			if v, ok := schemas.SafeExtractBoolPointer(bifrostReq.Params.ExtraParams["enable_prompt_expansion"]); ok {
+				delete(req.ExtraParams, "enable_prompt_expansion")
+				req.EnablePromptExpansion = v
+			}
+			if v, ok := schemas.SafeExtractBoolPointer(bifrostReq.Params.ExtraParams["enable_safety_checker"]); ok {
+				delete(req.ExtraParams, "enable_safety_checker")
+				req.EnableSafetyChecker = v
+			}
+		}
+	}
+
+	return req, nil
+}
+
+// UnmarshalHuggingFaceImageGenerationResponse unmarshals HuggingFace image generation response to Bifrost format
+func UnmarshalHuggingFaceImageGenerationResponse(data []byte, model string) (*schemas.BifrostImageGenerationResponse, error) {
+	if data == nil {
+		return nil, fmt.Errorf("response data is nil")
+	}
+	inferenceProvider, _, err := splitIntoModelProvider(model)
+	if err != nil {
+		return nil, err
+	}
+
+	switch inferenceProvider {
+	case nebius:
+		// Unmarshal into Nebius response format
+		var nebiusResponse nebiusProvider.NebiusImageGenerationResponse
+		if err := sonic.Unmarshal(data, &nebiusResponse); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal Nebius response: %w", err)
+		}
+
+		// Convert to Bifrost format using Nebius converter
+		bifrostResponse := nebiusProvider.ToBifrostImageResponse(&nebiusResponse)
+		if bifrostResponse == nil {
+			return nil, fmt.Errorf("failed to convert Nebius response to Bifrost format")
+		}
+
+		// Set model field (Nebius converter doesn't set it, similar to embeddings pattern)
+		if bifrostResponse.Model == "" {
+			bifrostResponse.Model = model
+		}
+
+		return bifrostResponse, nil
+
+	case hfInference:
+		// Handle raw byte data - encode to base64
+		b64Data := base64.StdEncoding.EncodeToString(data)
+		return &schemas.BifrostImageGenerationResponse{
+			Model: model,
+			Data: []schemas.ImageData{
+				{
+					B64JSON: b64Data,
+					Index:   0,
+				},
+			},
+		}, nil
+
+	case falAI:
+		// Handle fal-ai JSON response
+		var falResponse HuggingFaceFalAIImageGenerationResponse
+		if err := sonic.Unmarshal(data, &falResponse); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal fal-ai response: %w", err)
+		}
+
+		imageData := make([]schemas.ImageData, len(falResponse.Images))
+		for i, img := range falResponse.Images {
+			// Handle both URL and base64 responses
+			imageData[i] = schemas.ImageData{
+				URL:     img.URL,
+				B64JSON: img.B64JSON,
+				Index:   i,
+			}
+		}
+
+		return &schemas.BifrostImageGenerationResponse{
+			Model: model,
+			Data:  imageData,
+		}, nil
+
+	case together:
+		// Handle together JSON response
+		var togetherResponse HuggingFaceTogetherImageGenerationResponse
+		if err := sonic.Unmarshal(data, &togetherResponse); err != nil {
+			return nil, fmt.Errorf("failed to unmarshal together response: %w", err)
+		}
+
+		imageData := make([]schemas.ImageData, len(togetherResponse.Data))
+		for i, img := range togetherResponse.Data {
+			imageData[i] = schemas.ImageData{
+				B64JSON: img.B64JSON,
+				URL:     img.URL,
+				Index:   i,
+			}
+		}
+
+		return &schemas.BifrostImageGenerationResponse{
+			Model: model,
+			Data:  imageData,
+		}, nil
+
+	default:
+		return nil, fmt.Errorf("unsupported inference provider: %s", inferenceProvider)
+	}
+}
+
+// imageBytesToBase64DataURL converts raw image bytes to base64 data URL format
+func imageBytesToBase64DataURL(imageBytes []byte) string {
+	mimeType := http.DetectContentType(imageBytes)
+	b64Data := base64.StdEncoding.EncodeToString(imageBytes)
+	return fmt.Sprintf("data:%s;base64,%s", mimeType, b64Data)
+}
+
+// mapFalAIImageEditParams maps common parameters from Bifrost request to fal-ai request
+func mapFalAIImageEditParams(bifrostReq *schemas.BifrostImageEditRequest, req *HuggingFaceFalAIImageEditRequest) {
+	if bifrostReq.Params == nil {
+		return
+	}
+
+	// Map n to num_images for fal-ai
+	if bifrostReq.Params.N != nil {
+		req.NumImages = bifrostReq.Params.N
+	}
+
+	// Pass through output_format
+	if bifrostReq.Params.OutputFormat != nil {
+		if strings.ToLower(*bifrostReq.Params.OutputFormat) == "jpg" {
+			req.OutputFormat = schemas.Ptr("jpeg")
+		} else {
+			req.OutputFormat = bifrostReq.Params.OutputFormat
+		}
+	}
+
+	if bifrostReq.Params.ResponseFormat != nil && *bifrostReq.Params.ResponseFormat == "b64_json" {
+		req.SyncMode = schemas.Ptr(true)
+	}
+
+	// Convert size from "WxH" format to fal-ai's image_size object
+	if bifrostReq.Params.Size != nil && strings.ToLower(*bifrostReq.Params.Size) != "auto" {
+		size := strings.Split(*bifrostReq.Params.Size, "x")
+		if len(size) == 2 {
+			width, err := strconv.Atoi(size[0])
+			if err == nil {
+				height, err := strconv.Atoi(size[1])
+				if err == nil {
+					req.ImageSize = &HuggingFaceFalAISize{
+						Width:  width,
+						Height: height,
+					}
+				}
+			}
+		}
+	}
+
+	// Pass-through num_inference_steps
+	if bifrostReq.Params.NumInferenceSteps != nil {
+		req.NumInferenceSteps = bifrostReq.Params.NumInferenceSteps
+	}
+
+	// Pass-through seed
+	if bifrostReq.Params.Seed != nil {
+		req.Seed = bifrostReq.Params.Seed
+	}
+
+	// Parse fal-ai specific params from ExtraParams
+	if bifrostReq.Params.ExtraParams != nil {
+		// Map guidance_scale
+		if v, ok := schemas.SafeExtractFloat64Pointer(bifrostReq.Params.ExtraParams["guidance_scale"]); ok {
+			delete(req.ExtraParams, "guidance_scale")
+			req.GuidanceScale = v
+		}
+
+		// Map acceleration
+		if v, ok := schemas.SafeExtractStringPointer(bifrostReq.Params.ExtraParams["acceleration"]); ok {
+			delete(req.ExtraParams, "acceleration")
+			req.Acceleration = v
+		}
+
+		// Map enable_safety_checker
+		if v, ok := schemas.SafeExtractBoolPointer(bifrostReq.Params.ExtraParams["enable_safety_checker"]); ok {
+			delete(req.ExtraParams, "enable_safety_checker")
+			req.EnableSafetyChecker = v
+		}
+	}
+}
+
+// ToHuggingFaceImageEditRequest converts a Bifrost image edit request to fal-ai format
+func ToHuggingFaceImageEditRequest(bifrostReq *schemas.BifrostImageEditRequest) (*HuggingFaceFalAIImageEditRequest, error) {
+	if bifrostReq == nil || bifrostReq.Input == nil {
+		return nil, fmt.Errorf("bifrost request is nil or input is nil")
+	}
+
+	if len(bifrostReq.Input.Images) == 0 {
+		return nil, fmt.Errorf("at least one image is required")
+	}
+
+	// Convert images to base64 data URLs
+	imageURLs := make([]string, 0, len(bifrostReq.Input.Images))
+	for _, img := range bifrostReq.Input.Images {
+		if len(img.Image) == 0 {
+			continue
+		}
+		imageURLs = append(imageURLs, imageBytesToBase64DataURL(img.Image))
+	}
+
+	if len(imageURLs) == 0 {
+		return nil, fmt.Errorf("no valid images found")
+	}
+
+	// Extract model name to determine image field strategy
+	_, modelName, err := splitIntoModelProvider(bifrostReq.Model)
+	if err != nil {
+		return nil, fmt.Errorf("failed to split model name: %w", err)
+	}
+
+	req := &HuggingFaceFalAIImageEditRequest{
+		Prompt: bifrostReq.Input.Prompt,
+	}
+
+	// Check for explicit override in ExtraParams
+	var useMultiImage *bool
+	if bifrostReq.Params != nil && bifrostReq.Params.ExtraParams != nil {
+		req.ExtraParams = bifrostReq.Params.ExtraParams
+		if v, ok := schemas.SafeExtractBoolPointer(bifrostReq.Params.ExtraParams["use_image_urls"]); ok {
+			delete(req.ExtraParams, "use_image_urls")
+			useMultiImage = v
+		}
+	}
+
+	// Determine which image field to use based on model capabilities
+	if useMultiImage != nil {
+		// Explicit override from user
+		if *useMultiImage {
+			req.ImageURLs = imageURLs
+		} else if len(imageURLs) == 1 {
+			req.ImageURL = &imageURLs[0]
+		} else {
+			return nil, fmt.Errorf("use_image_urls is false but multiple images provided (%d images)", len(imageURLs))
+		}
+	} else if falAIMultiImageEditModels[modelName] {
+		// Model supports multiple images - always use image_urls
+		req.ImageURLs = imageURLs
+	} else if falAISingleImageEditModels[modelName] {
+		// Model only supports single image - validate and use image_url
+		if len(imageURLs) == 1 {
+			req.ImageURL = &imageURLs[0]
+		} else {
+			return nil, fmt.Errorf("model %s only supports single image, got %d images", modelName, len(imageURLs))
+		}
+	} else {
+		// Unknown model - fallback to count-based logic
+		if len(imageURLs) == 1 {
+			req.ImageURL = &imageURLs[0]
+		} else {
+			req.ImageURLs = imageURLs
+		}
+	}
+
+	// Map common parameters
+	mapFalAIImageEditParams(bifrostReq, req)
+	return req, nil
+}
+
+// extractImagesFromStreamResponse extracts images from a fal-ai streaming response.
+// Handles both API envelope structure (Data.Images) and legacy flattened format (top-level Images).
+func extractImagesFromStreamResponse(response *HuggingFaceFalAIImageStreamResponse) []FalAIImage {
+	// Prefer Data.Images if available (API envelope structure)
+	if response.Data != nil && len(response.Data.Images) > 0 {
+		return response.Data.Images
+	}
+	// Fall back to top-level Images (legacy format)
+	return response.Images
+}