From f08cc3fd4a82f0b80f6a03a43138984aec46069a Mon Sep 17 00:00:00 2001 From: saltacc Date: Fri, 2 Feb 2024 21:06:25 -0800 Subject: [PATCH] support docker container & more requirements --- .../comfy_ui/comfy_ui_dependencies.rs | 48 +++++++++++++++++-- .../comfy_ui/comfy_ui_inference_command.rs | 7 +++ .../workflow/comfy_ui/process_job.rs | 30 ++++++++++-- 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs index f2234988f6..e85ca31fb3 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs +++ b/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(), + }, ], } } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs index c878aa5a75..e9e1f0ada5 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs +++ b/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, @@ -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, diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index feef4cc8a1..7b4fad6890 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/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) -> Result { 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) -> Result { 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 = 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 { 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 {} } + // 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...");