java项目怎么集成stable diffusion图文生成算法

bash复制代码
pip install transformers torch

from transformers import StableDiffusionPipeline  
  
def generate_image(prompt):  
    pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")  
    image = pipeline(prompt, num_inference_steps=50, guidance_scale=7.5)[0]['sample']  
    # 这里为了简化，我们假设只是打印出图像数据（实际中应该保存或发送图像）  
    print(f"Generated image data for prompt: {prompt}")  
    # 在实际应用中，我们可能需要将图像保存到文件或使用其他方式返回  
  
if __name__ == "__main__":  
    import sys  
    if len(sys.argv) > 1:  
        prompt = ' '.join(sys.argv[1:])  
        generate_image(prompt)  
    else:  
        print("Usage: python stable_diffusion.py <prompt>")

import java.io.BufferedReader;  
import java.io.IOException;  
import java.io.InputStreamReader;  
  
public class StableDiffusionJava {  
    public static void main(String[] args) {  
        if (args.length < 1) {  
            System.out.println("Usage: java StableDiffusionJava <prompt>");  
            return;  
        }  
  
        String prompt = String.join(" ", args);  
        String pythonScriptPath = "python stable_diffusion.py";  
        try {  
            ProcessBuilder pb = new ProcessBuilder(pythonScriptPath, prompt);  
            Process p = pb.start();  
  
            BufferedReader reader = new BufferedReader(new InputStreamReader(p.getInputStream()));  
            String line;  
            while ((line = reader.readLine()) != null) {  
                System.out.println(line);  
            }  
  
            int exitCode = p.waitFor();  
            System.out.println("Exited with error code : " + exitCode);  
  
        } catch (IOException | InterruptedException e) {  
            e.printStackTrace();  
        }  
    }  
}

bash复制代码
pip install transformers flask

# stable_diffusion_server.py  
from flask import Flask, request, jsonify  
from transformers import StableDiffusionPipeline  
from PIL import Image  
import io  
import base64  
app = Flask(__name__)  
pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")  
@app.route('/generate', methods=['POST'])  
def generate_image():  
    data = request.json  
    prompt = data.get('prompt', 'A beautiful landscape')  
    num_inference_steps = data.get('num_inference_steps', 50)  
    guidance_scale = data.get('guidance_scale', 7.5)  
    try:  
        images = pipeline(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale)  
        # 假设我们只发送第一张生成的图像  
        image = images[0]['sample']  
        # 将PIL图像转换为Base64字符串  
        buffered = io.BytesIO()  
        image.save(buffered, format="PNG")  
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")  
        return jsonify({'image_base64': img_str})  
    except Exception as e:  
        return jsonify({'error': str(e)}), 500  
if __name__ == '__main__':  
    app.run(host='0.0.0.0', port=5000)

<dependency>  
    <groupId>org.json</groupId>  
    <artifactId>json</artifactId>  
    <version>20210307</version>  
</dependency>

// StableDiffusionClient.java  
import java.io.BufferedReader;  
import java.io.InputStreamReader;  
import java.net.HttpURLConnection;  
import java.net.URL;  
import java.nio.charset.StandardCharsets;  
import java.util.HashMap;  
import java.util.Map;  
import org.json.JSONObject;  
  
public class StableDiffusionClient {  
  
    public static void main(String[] args) {  
        String urlString = "http://localhost:5000/generate";  
        Map<String, Object> data = new HashMap<>();  
        data.put("prompt", "A colorful sunset over the ocean");  
        data.put("num_inference_steps", 50);  
        data.put("guidance_scale", 7.5);  
  
        try {  
            URL url = new URL(urlString);  
            HttpURLConnection con = (HttpURLConnection) url.openConnection();  
  
            con.setRequestMethod("POST");  
            con.setRequestProperty("Content-Type", "application/json; utf-8");  
            con.setRequestProperty("Accept", "application/json");  
            con.setDoOutput(true);  
  
            String jsonInputString = new JSONObject(data).toString();  
            byte[] postData = jsonInputString.getBytes(StandardCharsets.UTF_8);  
  
            try (java.io.OutputStream os = con.getOutputStream()) {  
                os.write(postData);  
            }  
  
            int responseCode = con.getResponseCode();  
            System.out.println("POST Response Code : " + responseCode);  
  
            BufferedReader in = new BufferedReader(  
                    new InputStreamReader(con.getInputStream()));  
            String inputLine;  
            StringBuffer response = new StringBuffer();  
  
            while ((inputLine = in.readLine()) != null) {  
                response.append(inputLine);  
            }  
            in.close();  
  
            // 打印接收到的JSON响应  
            System.out.println(response.toString());  
  
            // 解析JSON并获取图像Base64字符串（如果需要）  
            JSONObject jsonObj = new JSONObject(response.toString());  
            String imageBase64 = jsonObj.getString("image_base64");  
            System.out.println("Image Base64: " + imageBase64);  
  
        } catch (Exception e) {  
            e.printStackTrace();  
        }  
    }  
}