Merge pull request #295 from storytold/feature/comfyui

support docker container & more comfyui dependencies

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

@@ -36,12 +36,12 @@ impl RequiredModels {
                     url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_adapter.ckpt".to_string(),
                 },
                 ComfyDependency {
-                    location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/v3_sd15_mm.ckpt"),
-                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_mm.ckpt".to_string(),
+                    location: PathBuf::from("models/embeddings/badhandv4.pt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/badhandv4.pt".to_string(),
                 },
                 ComfyDependency {
-                    location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/v3_sd15_mm.ckpt"),
-                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_mm.ckpt".to_string(),
+                    location: PathBuf::from("models/embeddings/easynegative.safetensors"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/easynegative.safetensors".to_string(),
                 },
                 ComfyDependency {
                     location: PathBuf::from("models/controlnet/v3_sd15_sparsectrl_scribble.ckpt"),
@@ -91,6 +91,46 @@ impl RequiredModels {
                     location: PathBuf::from("models/clip_vision/CLIP-ViT-H-14-laion2B-s32B-b79K.safetensors"),
                     url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/model.safetensors".to_string(),
                 },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/ComfyUI-AnimateDiff-Evolved/models/v3_sd15_mm.ckpt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/v3_sd15_mm.ckpt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/hr16/DWPose-TorchScript-BatchSize5/dw-ll_ucoco_384_bs5.torchscript.pt"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/dw-ll_ucoco_384_bs5.torchscript.pt".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/yzd-v/DWPose/yolox_l.onnx"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/yolox_l.onnx".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/LiheYoung/Depth-Anything/checkpoints/depth_anything_vitl14.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/depth_anything_vitl14.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/lllyasviel/Annotators/sk_model2.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/sk_model2.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/lllyasviel/Annotators/sk_model.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/sk_model.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/lllyasviel/Annotators/netG.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/netG.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/comfyui_controlnet_aux/ckpts/lllyasviel/Annotators/ControlNetHED.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/ControlNetHED.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("custom_nodes/ComfyUI-Frame-Interpolation/ckpts/rife/rife47.pth"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/rife47.pth".to_string(),
+                },
+                ComfyDependency {
+                    location: PathBuf::from("input/example.png"),
+                    url: "https://storage.googleapis.com/vocodes-public/comfyui-deps/example.png".to_string(),
+                },
             ],
         }
     }

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

@@ -38,6 +38,9 @@ pub struct ComfyInferenceCommand {
     /// A single executable script or a much larger bash command.
     executable_or_command: ExecutableOrCommand,
 
+    // Where to mount the filesystem
+    pub(crate) mounts_directory: PathBuf,
+
     /// Config file to use
     config_path: Option<PathBuf>,
 
@@ -105,10 +108,14 @@ impl ComfyInferenceCommand {
                 }
             });
 
+        let mounts_directory = easyenv::get_env_pathbuf_required(
+            "COMFY_MOUNTS_DIRECTORY")?;
+
         Ok(Self {
             comfy_root_code_directory,
             executable_or_command,
             config_path,
+            mounts_directory,
             maybe_virtual_env_activation_command,
             maybe_docker_options,
             maybe_execution_timeout,

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

@@ -82,7 +82,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         .map_err(|e| ProcessSingleJobError::from_io_error(e))?;
 
     // ===================== DOWNLOAD REQUIRED MODELS IF NOT EXIST ===================== //
-    let root_comfy_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
+    let root_comfy_path = model_dependencies.inference_command.mounts_directory.clone();
 
     let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
     let remote_cloud_file_client = match remote_cloud_file_client {
@@ -126,7 +126,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     match all_models {
         Some(models) => {
             for model in &models.dependency_tokens.comfy {
-                let mut dep_path = model_dependencies.inference_command.comfy_root_code_directory.clone();
+                let mut dep_path = model_dependencies.inference_command.mounts_directory.clone();
                 dep_path = dep_path.join(model.location.clone());
                 if !dep_path.exists() {
                     download_file(model.url.clone(), dep_path.clone()).await.map_err(|e| ProcessSingleJobError::Other(e))?;
@@ -140,7 +140,11 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     }
 
     // Download workflow to ComfyRunner
-    let workflow_dir = root_comfy_path.join("../ComfyLauncher");
+    let workflow_dir = root_comfy_path.join("prompt");
+    // make folder if not exist
+    if !workflow_dir.exists() {
+        std::fs::create_dir_all(&workflow_dir).unwrap();
+    }
     let retrieved_workflow_record =  get_weight_by_token(
         job_args.workflow_source,
         false,
@@ -192,6 +196,10 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     match job_args.maybe_sd_model {
         Some(sd_model) => {
             let sd_dir = root_comfy_path.join("models").join("checkpoints");
+            // make if not exist
+            if !sd_dir.exists() {
+                std::fs::create_dir_all(&sd_dir).unwrap();
+            }
             let retrieved_sd_record =  get_weight_by_token(
                 sd_model,
                 false,
@@ -216,7 +224,11 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     let mut maybe_lora_path: Option<PathBuf> = None;
     match job_args.maybe_lora_model {
         Some(lora_model) => {
-            let lora_dir = root_comfy_path.join("models").join("checkpoints");
+            let lora_dir = root_comfy_path.join("models").join("loras");
+            // make if not exist
+            if !lora_dir.exists() {
+                std::fs::create_dir_all(&lora_dir).unwrap();
+            }
             let retrieved_lora_record =  get_weight_by_token(
                 lora_model,
                 false,
@@ -242,6 +254,10 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
     match job_args.maybe_input_file {
         Some(input_file) => {
             let input_dir = root_comfy_path.join("input");
+            // make if not exist
+            if !input_dir.exists() {
+                std::fs::create_dir_all(&input_dir).unwrap();
+            }
             let retrieved_input_record =  get_media_file(
                 input_file,
                 false,
@@ -262,6 +278,12 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result<JobSuccessResu
         None => {}
     }
 
+    // make outputs dir if not exist
+    let output_dir = root_comfy_path.join("output");
+    if !output_dir.exists() {
+        std::fs::create_dir_all(&output_dir).unwrap();
+    }
+
     // ==================== SETUP FOR INFERENCE ==================== //
 
     info!("Ready for ComfyUI inference...");