CAMEL-24540: camel-huggingface - honour model revision for sentence-embeddings/text-to-image and set the text-to-image OUTPUT header

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

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);
     }
 }

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:

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
             )

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));
+    }
+}