Merge pull request #314 from storytold/feature/comfyui-sidecar

Feature/comfyui sidecar

crates/service/job/inference_job/Cargo.toml

@@ -86,7 +86,7 @@ mime = "0.3.17"
 rand = { version = "0.8.5", features = [] }
 jsonpath_lib = "0.3.0"
 walkdir = "2.4.0"
-reqwest = "0.11.22"
+reqwest = { version = "0.11.22", features = ["blocking"] }
 opentelemetry = {version = "0.21.0", features = ["metrics"]}
 opentelemetry-otlp = { version = "0.14.0", features = ["metrics"] }
 opentelemetry_sdk = { version = "0.21.2", features = ["metrics", "rt-tokio"] }

crates/service/job/inference_job/config/inference-job.development.env

@@ -147,11 +147,13 @@ MOCAPNET_INFERENCE_EXECUTABLE_OR_COMMAND="pythonVirtualEnvironment/bin/python -m
 #                                   #
 # ================================= #
 COMFY_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/"
-COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunner.py"
+COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunnerForSidecar.py"
+COMFY_LAUNCH_COMMAND="venv/bin/python /home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/main.py"
 
 # directory to where files should be downloaded to (default to just the comfy directory)
 COMFY_MOUNTS_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/"
 COMFY_VIDEO_PROCESSING_SCRIPT="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/format_video.py"
+COMFY_SETUP_SCRIPT="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/install.py"
 
 # ================================= #
 #                                   #
@@ -186,3 +188,4 @@ HF_DATASETS_CACHE="/tmp/huggingface"
 # nltk also caches data.
 NLTK_DATA="/tmp/nltk_data"
 
+SKIP_TELEMETRY="true"

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs

@@ -1,6 +1,10 @@
 use std::path::PathBuf;
-
+use tokio::process::Command;
+use std::time::{Duration, Instant};
+use actix_web::http::StatusCode;
 use errors::AnyhowResult;
+use reqwest::blocking::Client;
+use tokio::time::sleep;
 
 use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::ComfyInferenceCommand;
 use crate::util::common_commands::ffmpeg_logo_watermark_command::FfmpegLogoWatermarkCommand;
@@ -12,9 +16,56 @@ pub struct ComfyDependencies {
 }
 
 impl ComfyDependencies {
-    pub fn setup() -> AnyhowResult<Self> {
+    pub async fn setup() -> AnyhowResult<Self> {
+        let inference_command = ComfyInferenceCommand::from_env()?;
+        let comfy_setup_script = inference_command.clone().comfy_setup_script;
+        let output = Command::new("python3")
+            .arg(&comfy_setup_script)
+            .output()
+            .await.expect("failed to run comfyui setup script");
+        println!("stdout: {}", String::from_utf8_lossy(&output.stdout));
+
+        // Start the comfy start script as a daemon process
+        let launch_command = format!(
+            "cd {} && {}",
+            inference_command.clone().comfy_root_code_directory.display(),
+            inference_command.clone().comfy_launch_command.display()
+        );
+
+        let _ = Command::new("sh")
+            .arg("-c")
+            .arg(&launch_command)
+            .stdout(std::process::Stdio::inherit())  // Inherit the stdout to stream the output
+            .spawn()
+            .expect("failed to start comfy start script");
+
+        let mut success = false;
+
+        let client = reqwest::Client::new();
+        let start = Instant::now();
+        let timeout = Duration::from_secs(60);
+        while Instant::now().duration_since(start) < timeout {
+            let response = client.get("http://127.0.0.1:8188/prompt").send().await;
+            match response {
+                Ok(r) if r.status() == StatusCode::OK => {
+                    println!("Received 200 response from http://127.0.0.1:8188/prompt");
+                    success = true;
+                    break;
+                },
+                _ => {
+                    println!("Did not receive 200 response, retrying...");
+                    sleep(Duration::from_secs(5)).await;
+                }
+            }
+        }
+
+        if !success {
+            println!("Timeout reached without receiving a 200 response.");
+            panic!("Comfy start failed");
+        }
+
         Ok(Self {
-            inference_command: ComfyInferenceCommand::from_env()?,
+            inference_command,
             dependency_tokens: RequiredModels::init(),
             ffmpeg_watermark_command: FfmpegLogoWatermarkCommand::from_env()?,
         })

crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs

@@ -44,6 +44,10 @@ pub struct ComfyInferenceCommand {
     // Video processing script
     pub(crate) processing_script: PathBuf,
 
+    pub(crate) comfy_setup_script: PathBuf,
+
+    pub(crate) comfy_launch_command: PathBuf,
+
     /// Config file to use
     config_path: Option<PathBuf>,
 
@@ -95,6 +99,12 @@ impl ComfyInferenceCommand {
             }
         };
 
+        let comfy_setup_script = easyenv::get_env_pathbuf_required(
+            "COMFY_SETUP_SCRIPT")?;
+
+        let comfy_launch_command = easyenv::get_env_pathbuf_required(
+            "COMFY_LAUNCH_COMMAND")?;
+
         let maybe_virtual_env_activation_command = easyenv::get_env_string_optional(
             "COMFY_INFERENCE_MAYBE_VENV_COMMAND");
 
@@ -124,6 +134,8 @@ impl ComfyInferenceCommand {
             config_path,
             mounts_directory,
             processing_script,
+            comfy_setup_script,
+            comfy_launch_command,
             maybe_virtual_env_activation_command,
             maybe_docker_options,
             maybe_execution_timeout,

crates/service/job/inference_job/src/job_specific_dependencies.rs

@@ -34,7 +34,7 @@ pub struct JobSpecificDependencies {
 
 impl JobSpecificDependencies {
 
-  pub fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult<Self> {
+  pub async fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult<Self> {
     let mut maybe_rvc_v2_dependencies = None;
     let mut maybe_sad_talker_dependencies = None;
     let mut maybe_svc_dependencies = None;
@@ -101,7 +101,7 @@ impl JobSpecificDependencies {
 
     if scoped_execution.can_run_job(InferenceModelType::ComfyUi) {
       print_with_space("Setting ComfyUI dependencies...");
-      maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup()?);
+      maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup().await?);
     }
 
     if scoped_execution.can_run_job(InferenceModelType::ConvertFbxToGltf) {

crates/service/job/inference_job/src/main.rs

@@ -237,7 +237,7 @@ async fn main() -> AnyhowResult<()> {
 
   let meter = opentelemetry::global::meter("inference-job");
 
-  let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution)?;
+  let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution).await?;
 
   let scoped_tempdir_for_downloads = ScopedTempDirCreator::for_directory(
     easyenv::get_env_pathbuf_or_default(

crates/service/job/inference_job/src/util/instrumentation.rs

@@ -21,6 +21,11 @@ pub fn init_otel_metrics_pipeline(
 ) -> Result<(), opentelemetry::metrics::MetricsError>  {
     info!("Setting up otel metrics pipeline...");
 
+    // check env var to see if we should skip telemetry
+    if std::env::var("SKIP_TELEMETRY").is_ok() {
+        warn!("Skipping telemetry setup due to SKIP_TELEMETRY env var");
+        return Ok(());
+    }
     let provider = opentelemetry_otlp::new_pipeline()
         .metrics(opentelemetry_sdk::runtime::Tokio)
         // TODO: 1. read host from env 2. Single pod of otel-collector is probably not good enough, run daemonset?