⚡ Support text-to-image generation

Author: MrGraversen

replicate-tools/default_replicate_models.txt

@@ -1,8 +1,6 @@
 meta/meta-llama-3.1-405b-instruct
 meta/meta-llama-3-70b-instruct
 meta/meta-llama-3-8b-instruct
-thomasmol/whisper-diarization
 black-forest-labs/flux-schnell
-salesforce/blip
-stability-ai/sdxl
-openai/whisper
+black-forest-labs/flux-dev
+black-forest-labs/flux-pro

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

@@ -19,7 +19,7 @@
 
 @Configuration
 @EnableConfigurationProperties(ReplicateProperties.class)
-@Import({Llama3Configuration.class})
+@Import({Llama3Configuration.class, FluxConfiguration.class})
 @ComponentScan({"io.graversen.replicate.service", "io.graversen.replicate.facade"})
 public class ReplicateConfiguration {
     @Bean

spring-boot-starter-replicate/src/main/java/io/graversen/replicate/facade/ReplicateFacade.java

@@ -1,6 +1,7 @@
 package io.graversen.replicate.facade;
 
 import io.graversen.replicate.common.ReplicateModel;
+import io.graversen.replicate.service.CreateImagePrediction;
 import io.graversen.replicate.service.CreateTextPrediction;
 import io.graversen.replicate.service.ReplicateService;
 import lombok.NonNull;
@@ -34,11 +35,30 @@ public CompletableFuture<PredictionResponseAndModel> createPrediction(
                 .whenCompleteAsync(emitPredictionResponseTask, executorService);
     }
 
+    public CompletableFuture<PredictionResponseAndModel> createPrediction(
+            @NonNull ReplicateModel model,
+            @NonNull CreateImagePrediction createPrediction
+    ) {
+        return CompletableFuture
+                .supplyAsync(doCreatePrediction(model, createPrediction), executorService)
+                .thenApplyAsync(checkAndEmitPredictionCreationTask, executorService)
+                .thenApplyAsync(pollPredictionStatusTask, executorService)
+                .whenCompleteAsync(emitPredictionResponseTask, executorService);
+    }
+
     Supplier<Optional<PredictionResponseAndModel>> doCreatePrediction(
             @NonNull ReplicateModel model,
             @NonNull CreateTextPrediction createPrediction
     ) {
         return () -> replicateService.createPrediction(model, createPrediction)
                 .map(predictionResponse -> new PredictionResponseAndModel(predictionResponse, model));
     }
+
+    Supplier<Optional<PredictionResponseAndModel>> doCreatePrediction(
+            @NonNull ReplicateModel model,
+            @NonNull CreateImagePrediction createPrediction
+    ) {
+        return () -> replicateService.createPrediction(model, createPrediction)
+                .map(predictionResponse -> new PredictionResponseAndModel(predictionResponse, model));
+    }
 }

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

@@ -42,10 +42,26 @@ public Optional<PredictionResponse> createPrediction(@NonNull ReplicateModel mod
         return doCreatePrediction(model, mappedPrediction);
     }
 
+    @SneakyThrows
+    public Optional<PredictionResponse> createPrediction(@NonNull ReplicateModel model, @NonNull CreateImagePrediction createPrediction) {
+        final var predictionMapper = predictionMappers.stream()
+                .filter(supportsImagePredictions())
+                .filter(supportsModel(model))
+                .findFirst()
+                .orElseThrow(unsupportedModelError(model));
+
+        final var mappedPrediction = predictionMapper.apply(model, createPrediction);
+        return doCreatePrediction(model, mappedPrediction);
+    }
+
     Predicate<PredictionMapper> supportsTextPredictions() {
         return predictionMapper -> predictionMapper.supportsType(PredictionTypes.TEXT);
     }
 
+    Predicate<PredictionMapper> supportsImagePredictions() {
+        return predictionMapper -> predictionMapper.supportsType(PredictionTypes.IMAGE);
+    }
+
     Predicate<PredictionMapper> supportsModel(@NonNull ReplicateModel model) {
         return predictionMapper -> predictionMapper.supportsModel(model);
     }