Skip to content

Commit

Permalink
initial commit
Browse files Browse the repository at this point in the history
  • Loading branch information
@franknoh
franknoh committed Dec 5, 2022
0 parents commit d877104
Show file tree
Hide file tree
Showing 14 changed files with 49,615 additions and 0 deletions.
2 changes: 2 additions & 0 deletions .gitattributes
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
# Auto detect text files and perform LF normalization
* text=auto
29 changes: 29 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
# Compiled class file
*.class

# Log file
*.log

# BlueJ files
*.ctxt

# Mobile Tools for Java (J2ME)
.mtj.tmp/

# Package Files #
*.jar
*.war
*.nar
*.ear
*.zip
*.tar.gz
*.rar

# virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml
hs_err_pid*
replay_pid*

.idea
target
src/main/resources/tmp
src/main/resources/*.pt
21 changes: 21 additions & 0 deletions LICENSE
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
MIT License

Copyright (c) 2022 franknoh

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
12 changes: 12 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# stable-diffusion-java

### how it works
this repo uses [DeepJavaLibrary](https://djl.ai) to run torchscript models on java.
the tokenizer and klms sampler were ported from [stable-diffusion-pytorch](https://github.com/kjsman/stable-diffusion-pytorch) with minimal changes.

### how to use
1. Install python
2. Install torch (see [here](https://pytorch.org/get-started/locally/))
3. Update pom.xml (see [here](https://docs.djl.ai/engines/pytorch/pytorch-engine/index.html))
4. Convert model to torchscript (see [here](https://github.com/franknoh/stable-diffusion-jit))
5. Run the code
42 changes: 42 additions & 0 deletions pom.xml
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>

<groupId>org.example</groupId>
<artifactId>stable-diffusion-processing</artifactId>
<version>1.0-SNAPSHOT</version>

<properties>
<maven.compiler.source>16</maven.compiler.source>
<maven.compiler.target>16</maven.compiler.target>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
</properties>

<dependencies>
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-engine</artifactId>
<version>0.19.0</version>
<scope>runtime</scope>
</dependency>
<dependency>
<groupId>ai.djl</groupId>
<artifactId>api</artifactId>
<version>0.19.0</version>
</dependency>
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-native-cpu</artifactId>
<classifier>win-x86_64</classifier>
<scope>runtime</scope>
<version>1.12.1</version>
</dependency>
<dependency>
<groupId>ai.djl.pytorch</groupId>
<artifactId>pytorch-jni</artifactId>
<version>1.12.1-0.19.0</version>
</dependency>
</dependencies>
</project>
69 changes: 69 additions & 0 deletions src/main/java/org/franknoh/Clip.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,69 @@
package org.franknoh;

import ai.djl.Device;
import ai.djl.MalformedModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;

import java.io.IOException;
import java.nio.file.Paths;
import java.util.List;

public class Clip {
private ZooModel clip;
private Device device;
private Tokenizer tokenizer;
private NDManager manager;
private Predictor<NDList, NDList> clip_predictor;
Clip(Tokenizer tokenizer) {
if(Engine.getInstance().getGpuCount() > 0) {
this.device = Device.gpu();
} else {
this.device = Device.cpu();
}
this.manager = NDManager.newBaseManager(this.device);
Criteria<NDList, NDList> clip_c = Criteria.builder()
.setTypes(NDList.class, NDList.class)
.optModelPath(Paths.get("src/main/resources/clip.pt"))
.optEngine("PyTorch")
.optDevice(this.device)
.build();
try {
this.clip = ModelZoo.loadModel(clip_c);
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new RuntimeException(e);
}
this.tokenizer = tokenizer;
this.clip_predictor = this.clip.newPredictor();
}
public NDArray run(List<Integer> fi_tokens) {
NDList clip_input = new NDList();
int[] tokens_array = new int[fi_tokens.size()];
for (int i = 0; i < fi_tokens.size(); i++) {
tokens_array[i] = fi_tokens.get(i);
}
int[][] tokens_array_2d = new int[][]{tokens_array};
NDArray tokens_ndarray = this.manager.create(tokens_array_2d);
clip_input.add(tokens_ndarray);
NDList clip_output;
try {
clip_output = clip_predictor.predict(clip_input);
} catch (TranslateException e) {
throw new RuntimeException(e);
}
return clip_output.get(0);
}

public NDArray embedText(String fi_text) {
List<Integer> fi_tokens = this.tokenizer.encode(fi_text);
return run(fi_tokens);
}
}
75 changes: 75 additions & 0 deletions src/main/java/org/franknoh/Decoder.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
package org.franknoh;

import ai.djl.Device;
import ai.djl.MalformedModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.ndarray.types.DataType;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;

import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.file.Paths;
import java.util.Random;

public class Decoder {
private ZooModel decoder;
private Device device;
private NDManager manager;
private Predictor<NDList, NDList> diffusion_predictor;
Decoder() {
if(Engine.getInstance().getGpuCount() > 0) {
this.device = Device.gpu();
} else {
this.device = Device.cpu();
}
this.manager = NDManager.newBaseManager(this.device);
Criteria<NDList, NDList> decoder_c = Criteria.builder()
.setTypes(NDList.class, NDList.class)
.optModelPath(Paths.get("src/main/resources/decoder.pt"))
.optEngine("PyTorch")
.optDevice(device)
.build();
try {
this.decoder = decoder_c.loadModel();
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new RuntimeException(e);
}
this.diffusion_predictor = this.decoder.newPredictor();
}
public NDArray run(NDArray latent) {
NDList decoder_input = new NDList(latent);
NDList decoder_output;
try {
decoder_output = diffusion_predictor.predict(decoder_input);
} catch (TranslateException e) {
throw new RuntimeException(e);
}
return decoder_output.get(0);
}
public void saveImage(NDArray t_latent, String name) {
t_latent = t_latent.duplicate();
NDArray t_image = this.run(t_latent).get(0);
t_image = t_image.add(1).mul(127.5f).round().clip(0, 255).toType(DataType.UINT8, false).transpose(1, 2, 0);
OutputStream t_stream = null;
try {
t_stream = new FileOutputStream(Paths.get("src/main/resources/out/"+name+".png").toFile());
} catch (FileNotFoundException e) {
throw new RuntimeException(e);
}
try {
ImageFactory.getInstance().fromNDArray(t_image).save(t_stream, "png");
} catch (IOException e) {
throw new RuntimeException(e);
}
}
}
54 changes: 54 additions & 0 deletions src/main/java/org/franknoh/Diffusion.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package org.franknoh;

import ai.djl.Device;
import ai.djl.MalformedModelException;
import ai.djl.engine.Engine;
import ai.djl.inference.Predictor;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDList;
import ai.djl.ndarray.NDManager;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.TranslateException;

import java.io.IOException;
import java.nio.file.Paths;

public class Diffusion {
private ZooModel diffusion;
private Device device;
private NDManager manager;
private Predictor<NDList, NDList> diffusion_predictor;
Diffusion() {
if(Engine.getInstance().getGpuCount() > 0) {
this.device = Device.gpu();
} else {
this.device = Device.cpu();
}
this.manager = NDManager.newBaseManager(this.device);
Criteria<NDList, NDList> diffusion_c = Criteria.builder()
.setTypes(NDList.class, NDList.class)
.optModelPath(Paths.get("src/main/resources/diffusion.pt"))
.optEngine("PyTorch")
.optDevice(device)
.build();
try {
this.diffusion = diffusion_c.loadModel();
} catch (IOException | ModelNotFoundException | MalformedModelException e) {
throw new RuntimeException(e);
}
this.diffusion_predictor = this.diffusion.newPredictor();
}
public NDArray run(NDArray latent, NDArray context, NDArray time_embedding) {
latent = latent.concat(latent);
NDList diffusion_input = new NDList(latent, context, time_embedding);
NDList diffusion_output;
try {
diffusion_output = diffusion_predictor.predict(diffusion_input);
} catch (TranslateException e) {
throw new RuntimeException(e);
}
return diffusion_output.get(0);
}
}
26 changes: 26 additions & 0 deletions src/main/java/org/franknoh/Generator.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
package org.franknoh;

import ai.djl.Device;
import ai.djl.engine.Engine;
import ai.djl.ndarray.NDArray;
import ai.djl.ndarray.NDManager;

import ai.djl.ndarray.types.DataType;
import ai.djl.ndarray.types.Shape;

public class Generator {
private final NDManager manager;
Generator() {
Device device;
if(Engine.getInstance().getGpuCount() > 0) {
device = Device.gpu();
} else {
device = Device.cpu();
}
this.manager = NDManager.newBaseManager(device);
}

NDArray sample(Shape shape) {
return this.manager.randomNormal(shape, DataType.FLOAT32);
}
}
Loading

0 comments on commit d877104

Please sign in to comment.