✨ Support Flux model family mapping

Author: MrGraversen

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/TextToImageAspectRatio.java

@@ -8,6 +8,18 @@ public class TextToImageAspectRatio {
     @NonNull Integer height;
     @NonNull AspectRatios aspectRatio;
 
+    public static TextToImageAspectRatio square() {
+        return new TextToImageAspectRatio(1080, AspectRatios.RATIO_1_BY_1);
+    }
+
+    public static TextToImageAspectRatio portrait() {
+        return new TextToImageAspectRatio(1350, AspectRatios.RATIO_4_BY_5);
+    }
+
+    public static TextToImageAspectRatio reel() {
+        return new TextToImageAspectRatio(1920, AspectRatios.RATIO_9_BY_16);
+    }
+
     public Integer getWidth() {
         return (int) Math.round(height * ((double) aspectRatio.getWidthRatio() / aspectRatio.getHeightRatio()));
     }

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/common/TextToImagePrompt.java

@@ -0,0 +1,13 @@
+package io.graversen.replicate.common;
+
+import jakarta.annotation.Nullable;
+import lombok.NonNull;
+import lombok.Value;
+
+@Value
+public class TextToImagePrompt {
+    @NonNull String prompt;
+    @Nullable String negativePrompt;
+    @NonNull TextToImageAspectRatio aspectRatio;
+    @Nullable Double promptStrength;
+}

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/configuration/FluxConfiguration.java

@@ -0,0 +1,10 @@
+package io.graversen.replicate.configuration;
+
+import org.springframework.context.annotation.ComponentScan;
+import org.springframework.context.annotation.Configuration;
+
+@Configuration
+@ComponentScan("io.graversen.replicate.flux")
+public class FluxConfiguration {
+
+}

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/flux/FluxModels.java

@@ -0,0 +1,11 @@
+package io.graversen.replicate.flux;
+
+import io.graversen.replicate.common.ReplicateModel;
+import lombok.experimental.UtilityClass;
+
+@UtilityClass
+public class FluxModels {
+    public static final ReplicateModel FLUX_SCHNELL = new ReplicateModel("black-forest-labs", "flux-schnell");
+    public static final ReplicateModel FLUX_DEV = new ReplicateModel("black-forest-labs", "flux-dev");
+    public static final ReplicateModel FLUX_PRO = new ReplicateModel("black-forest-labs", "flux-pro");
+}

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/flux/FluxPredictionMapper.java

@@ -0,0 +1,107 @@
+package io.graversen.replicate.flux;
+
+import io.graversen.replicate.common.BasePredictionMapper;
+import io.graversen.replicate.common.PredictionTypes;
+import io.graversen.replicate.common.ReplicateModel;
+import io.graversen.replicate.models.*;
+import io.graversen.replicate.service.CreateImagePrediction;
+import lombok.NonNull;
+import org.springframework.stereotype.Component;
+
+import java.util.Set;
+import java.util.function.Function;
+
+@Component
+public class FluxPredictionMapper extends BasePredictionMapper<CreateImagePrediction, Object> {
+    @Override
+    protected Set<ReplicateModel> supportedModels() {
+        return Set.of(
+                FluxModels.FLUX_DEV,
+                FluxModels.FLUX_SCHNELL,
+                FluxModels.FLUX_PRO
+        );
+    }
+
+    @Override
+    public Object apply(@NonNull ReplicateModel model, @NonNull CreateImagePrediction createPrediction) {
+        if (model.equals(FluxModels.FLUX_DEV)) {
+            return new BlackForestLabsFluxDevPredictionrequest(
+                    null,
+                    mapFluxDevInput().apply(createPrediction),
+                    null,
+                    null,
+                    null,
+                    null
+            );
+        } else if (model.equals(FluxModels.FLUX_SCHNELL)) {
+            return new BlackForestLabsFluxSchnellPredictionrequest(
+                    null,
+                    mapFluxSchnellInput().apply(createPrediction),
+                    null,
+                    null,
+                    null,
+                    null
+            );
+        } else if (model.equals(FluxModels.FLUX_PRO)) {
+            return new BlackForestLabsFluxProPredictionrequest(
+                    null,
+                    mapFluxProInput().apply(createPrediction),
+                    null,
+                    null,
+                    null,
+                    null
+            );
+        } else {
+            throw new IllegalArgumentException("Unsupported Replicate Model: " + model);
+        }
+    }
+
+    @Override
+    public boolean supportsType(@NonNull PredictionTypes type) {
+        return PredictionTypes.IMAGE.equals(type);
+    }
+
+    Function<CreateImagePrediction, BlackForestLabsFluxDevInput> mapFluxDevInput() {
+        return createImagePrediction -> new BlackForestLabsFluxDevInput(
+                createImagePrediction.getSeed(),
+                null,
+                createImagePrediction.getPrompt().getPrompt(),
+                false,
+                createImagePrediction.getPrompt().getPromptStrength(),
+                null,
+                createImagePrediction.getOutputs(),
+                createImagePrediction.getPrompt().getAspectRatio().getAspectRatioAsString(),
+                "png",
+                100,
+                null,
+                createImagePrediction.getInferenceSteps(),
+                true
+        );
+    }
+
+    Function<CreateImagePrediction, BlackForestLabsFluxSchnellInput> mapFluxSchnellInput() {
+        return createImagePrediction -> new BlackForestLabsFluxSchnellInput(
+                createImagePrediction.getSeed(),
+                createImagePrediction.getPrompt().getPrompt(),
+                false,
+                null,
+                createImagePrediction.getOutputs(),
+                createImagePrediction.getPrompt().getAspectRatio().getAspectRatioAsString(),
+                "png",
+                100,
+                true
+        );
+    }
+
+    Function<CreateImagePrediction, BlackForestLabsFluxProInput> mapFluxProInput() {
+        return createImagePrediction -> new BlackForestLabsFluxProInput(
+                createImagePrediction.getSeed(),
+                createImagePrediction.getInferenceSteps(),
+                createImagePrediction.getPrompt().getPrompt(),
+                createImagePrediction.getPrompt().getPromptStrength(),
+                null,
+                createImagePrediction.getPrompt().getAspectRatio().getAspectRatioAsString(),
+                5
+        );
+    }
+}

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/service/CreateImagePrediction.java

@@ -0,0 +1,17 @@
+package io.graversen.replicate.service;
+
+import io.graversen.replicate.common.TextToImagePrompt;
+import jakarta.annotation.Nullable;
+import lombok.Getter;
+import lombok.NonNull;
+import lombok.RequiredArgsConstructor;
+
+@Getter
+@RequiredArgsConstructor
+public class CreateImagePrediction {
+    private final @NonNull TextToImagePrompt prompt;
+    private final @Nullable Integer outputs;
+    private final @Nullable Integer inferenceSteps;
+    private final @Nullable Integer seed;
+
+}