diff --git a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java index d6e9649f401e2..40fb875d91cae 100644 --- a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java +++ b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/SentenceEmbeddingsPredictor.java @@ -103,7 +103,7 @@ protected String getRequirements() { @Override protected String getPythonScript() { - return loadPythonScript("sentence_embeddings.py", config.getDevice(), config.getModelId()); + return loadPythonScript("sentence_embeddings.py", config.getDevice(), config.getModelId(), config.getRevision()); } @Override diff --git a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java index a928854acd1ac..561e7cd0d381b 100644 --- a/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java +++ b/components/camel-ai/camel-huggingface/src/main/java/org/apache/camel/component/huggingface/tasks/TextToImagePredictor.java @@ -22,6 +22,7 @@ import ai.djl.modality.Output; import org.apache.camel.Exchange; import org.apache.camel.RuntimeCamelException; +import org.apache.camel.component.huggingface.HuggingFaceConstants; import org.apache.camel.component.huggingface.HuggingFaceEndpoint; /** @@ -92,7 +93,7 @@ protected String getRequirements() { @Override protected String getPythonScript() { - return loadPythonScript("text_to_image.py", config.getModelId(), config.getDevice()); + return loadPythonScript("text_to_image.py", config.getModelId(), config.getRevision(), config.getDevice()); } @Override @@ -117,5 +118,6 @@ protected void processOutput(Exchange exchange, Output output) throws Exception } exchange.getMessage().setBody(imageBytes); exchange.getMessage().setHeader("Content-Type", "image/png"); + exchange.getMessage().setHeader(HuggingFaceConstants.OUTPUT, imageBytes); } } diff --git a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py index 6fe310f35f245..9ee585c2dde64 100644 --- a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py +++ b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/sentence_embeddings.py @@ -31,7 +31,7 @@ def handle(inputs: Input): device = '%s' if device == 'auto': device = 'cuda' if torch.cuda.is_available() else 'cpu' - model = SentenceTransformer('%s', device=device) + model = SentenceTransformer('%s', device=device, revision='%s') logging.debug("Model initialized") if inputs.content.size() == 0: diff --git a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py index e8b41788c2b52..aa7eae92638cc 100644 --- a/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py +++ b/components/camel-ai/camel-huggingface/src/main/resources/org/apache/camel/component/huggingface/tasks/text_to_image.py @@ -31,6 +31,7 @@ def handle(inputs: Input): logging.debug("Initializing pipeline") pipe = StableDiffusionPipeline.from_pretrained( '%s', + revision='%s', torch_dtype=torch.float32, # CPU-safe safety_checker=None ) diff --git a/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/RevisionAndOutputHeaderTest.java b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/RevisionAndOutputHeaderTest.java new file mode 100644 index 0000000000000..535a7fa69737d --- /dev/null +++ b/components/camel-ai/camel-huggingface/src/test/java/org/apache/camel/component/huggingface/tasks/RevisionAndOutputHeaderTest.java @@ -0,0 +1,90 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.camel.component.huggingface.tasks; + +import ai.djl.modality.Output; +import org.apache.camel.Exchange; +import org.apache.camel.component.huggingface.HuggingFaceConfiguration; +import org.apache.camel.component.huggingface.HuggingFaceConstants; +import org.apache.camel.component.huggingface.HuggingFaceEndpoint; +import org.apache.camel.impl.DefaultCamelContext; +import org.apache.camel.support.DefaultExchange; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The sentence-embeddings and text-to-image tasks must honour the configured model revision (they were the only two + * that dropped it), and the text-to-image task must publish its result on the OUTPUT header as its Javadoc promises. + */ +class RevisionAndOutputHeaderTest { + + private DefaultCamelContext context; + + @BeforeEach + void setUp() { + context = new DefaultCamelContext(); + } + + @AfterEach + void tearDown() { + context.stop(); + } + + private HuggingFaceEndpoint endpoint(HuggingFaceConfiguration config) { + HuggingFaceEndpoint endpoint = new HuggingFaceEndpoint(null, null, config); + endpoint.setCamelContext(context); + return endpoint; + } + + @Test + void sentenceEmbeddingsScriptPinsRevision() { + HuggingFaceConfiguration config = new HuggingFaceConfiguration(); + config.setModelId("sentence-transformers/all-MiniLM-L6-v2"); + config.setRevision("v1.5"); + SentenceEmbeddingsPredictor predictor = new SentenceEmbeddingsPredictor(endpoint(config)); + assertTrue(predictor.getPythonScript().contains("revision='v1.5'"), + "the generated script must pin the configured revision"); + } + + @Test + void textToImageScriptPinsRevision() { + HuggingFaceConfiguration config = new HuggingFaceConfiguration(); + config.setModelId("stabilityai/stable-diffusion"); + config.setRevision("fp16"); + TextToImagePredictor predictor = new TextToImagePredictor(endpoint(config)); + assertTrue(predictor.getPythonScript().contains("revision='fp16'"), + "the generated script must pin the configured revision"); + } + + @Test + void textToImagePublishesTheImageOnTheOutputHeader() throws Exception { + HuggingFaceConfiguration config = new HuggingFaceConfiguration(); + TextToImagePredictor predictor = new TextToImagePredictor(endpoint(config)); + Exchange exchange = new DefaultExchange(context); + Output output = new Output(); + byte[] image = { 1, 2, 3, 4 }; + output.add("data", image); + + predictor.processOutput(exchange, output); + + assertArrayEquals(image, exchange.getMessage().getHeader(HuggingFaceConstants.OUTPUT, byte[].class)); + } +}