kiroshi 0.1.0

An opinionated image diffusion server for Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use crate::Error;
use crate::protocol;

use tokio::fs;
use tokio::io;
use tokio::process;
use tokio::time;

use std::path::Path;
use std::sync::Arc;

#[derive(Debug, Clone)]
pub struct Server {
    _container: Arc<Container>,
}

#[derive(Debug)]
struct Container(String);

impl Server {
    pub async fn run(
        image_models_dir: impl AsRef<Path>,
        video_models_dir: impl AsRef<Path>,
    ) -> Result<Server, Error> {
        let image_models = {
            let image_models_dir = image_models_dir.as_ref();
            fs::create_dir_all(&image_models_dir).await?;

            format!(
                "{host}:/models/image",
                host = image_models_dir.to_string_lossy()
            )
        };

        let video_models = {
            let video_models_dir = video_models_dir.as_ref();
            fs::create_dir_all(&video_models_dir).await?;

            format!(
                "{host}:/models/video",
                host = video_models_dir.to_string_lossy()
            )
        };

        let mut process = process::Command::new("docker")
            .arg("create")
            .args(["-t", "--rm"])
            .args(["--gpus", "all"])
            .args(["-p", "9149:9149"])
            .args(["-v", &image_models])
            .args(["-v", &video_models])
            .arg("ghcr.io/hecrj/kiroshi/server:latest")
            .stdout(std::process::Stdio::piped())
            .stdin(std::process::Stdio::null())
            .stderr(std::process::Stdio::null())
            .spawn()?;

        let container = {
            use io::AsyncBufReadExt;

            let output = io::BufReader::new(process.stdout.take().expect("piped stdout"));

            let mut lines = output.lines();

            lines.next_line().await?.ok_or(Error::DockerFailed)?
        };

        let _start = process::Command::new("docker")
            .args(["start", &container])
            .output()
            .await?;

        let _logs = process::Command::new("docker")
            .args(["logs", "-f", &container])
            .spawn()?;

        // Wait until server is accepting connections
        while ping().await.is_err() {
            time::sleep(time::Duration::from_millis(500)).await;
        }

        Ok(Server {
            _container: Arc::new(Container(container)),
        })
    }
}

impl Drop for Container {
    fn drop(&mut self) {
        use std::process;

        let _ = process::Command::new("docker")
            .args(["stop", &self.0])
            .stdin(process::Stdio::null())
            .stdout(process::Stdio::null())
            .stderr(process::Stdio::null())
            .spawn();
    }
}

pub const PING: protocol::Plug<protocol::Never, bool> = protocol::Plug::new("ping");

async fn ping() -> Result<bool, Error> {
    let mut stream = protocol::connect(PING).await?;

    Ok(stream.read().await?)
}