Java集成stable diffusion 的方法 - IT评测·应用市场-qidao123.com

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

复制代码