fi.evolver.ai.spring.provider.replicate.ReplicateRequestGenerator Maven / Gradle / Ivy

Go to download
package fi.evolver.ai.spring.provider.replicate;

import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.ASPECT_RATIO;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.DISABLE_SAFETY_CHECKER;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.GO_FAST;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.GUIDANCE;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.HEIGHT;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.INTERVAL;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.MEGAPIXELS;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.NUM_OUTPUTS;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.OUTPUT_FORMAT;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.OUTPUT_QUALITY;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.PROMPT_STRENGTH;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.PROMPT_UPSAMPLING;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.RAW;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.SAFETY_TOLERANCE;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.SEED;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.STEPS;
import static fi.evolver.ai.spring.provider.replicate.ReplicateRequestParameters.WIDTH;

import java.io.IOException;
import java.io.StringWriter;
import java.io.UncheckedIOException;

import com.fasterxml.jackson.core.JsonGenerator;

import fi.evolver.ai.spring.image.prompt.ImageGenerationPrompt;
import fi.evolver.ai.spring.util.Json;

public class ReplicateRequestGenerator {
	private ReplicateRequestGenerator() {}

	public static String generate(ImageGenerationPrompt prompt, String version) {
		StringWriter writer = new StringWriter();

		try (JsonGenerator json = Json.JSON_FACTORY.createGenerator(writer)) {
			json.useDefaultPrettyPrinter();

			json.writeStartObject();
			json.writeStringField("version", version);
			json.writeObjectFieldStart("input");
			writeModelSpecificParameters(json, prompt);
			json.writeEndObject(); // input
			json.writeEndObject();
		}
		catch (IOException e) {
			throw new UncheckedIOException(e);
		}
		return writer.toString();
	}

	@SuppressWarnings("deprecation")
	private static void writeModelSpecificParameters(JsonGenerator json, ImageGenerationPrompt prompt) throws IOException {
		Json.writeIntegerProperty(json, prompt, SEED);
		json.writeStringField(OUTPUT_FORMAT, prompt.getStringProperty(OUTPUT_FORMAT).orElse("png"));
		Json.writeIntegerProperty(json, prompt, OUTPUT_QUALITY);
		json.writeStringField("prompt", prompt.prompt());
		
		Json.writeIntegerProperty(json, prompt, NUM_OUTPUTS);
		Json.writeDoubleProperty(json, prompt, MEGAPIXELS);
		Json.writeBooleanProperty(json, prompt, GO_FAST);
		Json.writeDoubleProperty(json, prompt, GUIDANCE);
		Json.writeStringProperty(json, prompt, ASPECT_RATIO);
		Json.writeDoubleProperty(json, prompt, PROMPT_STRENGTH);
		Json.writeBooleanProperty(json, prompt, DISABLE_SAFETY_CHECKER);
		Json.writeDoubleProperty(json, prompt, STEPS);
		Json.writeDoubleProperty(json, prompt, INTERVAL);
		Json.writeStringProperty(json, prompt, ASPECT_RATIO);
		Json.writeIntegerProperty(json, prompt, SAFETY_TOLERANCE);
		Json.writeBooleanProperty(json, prompt, PROMPT_UPSAMPLING);
		Json.writeIntegerProperty(json, prompt, WIDTH);
		Json.writeIntegerProperty(json, prompt, HEIGHT);
		Json.writeBooleanProperty(json, prompt, RAW);
	}
}