enable cog-video for staff only

crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/workflow_payload.rs

@@ -211,6 +211,10 @@ pub struct WorkflowArgs {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub use_cinematic: Option<bool>,
 
+    #[serde(rename = "uv")]
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub use_cogvideo: Option<bool>,
+
     #[serde(rename = "s")]
     #[serde(skip_serializing_if = "Option::is_none")]
     pub strength: Option<f32>,

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

@@ -86,6 +86,8 @@ pub struct InferenceArgs<'s> {
 
     pub use_cinematic: bool,
 
+    pub use_cogvideo: bool,
+
     pub maybe_strength: Option<f32>,
 
     pub frame_skip: Option<u8>,
@@ -308,6 +310,10 @@ impl ComfyInferenceCommand {
             command.push_str(" --enable-cinematic ");
         }
 
+        if args.use_cogvideo {
+            command.push_str(" --enable-cogvideo ");
+        }
+
         if let Some(strength) = args.maybe_strength {
             command.push_str(" --strength ");
             command.push_str(&strength.to_string());

crates/service/job/inference_job/src/job/job_types/workflow/video_style_transfer/process_video_style_transfer_job.rs

@@ -402,7 +402,6 @@ pub async fn process_video_style_transfer_job(deps: &JobDependencies, job: &Avai
 
     info!("Running ComfyUI inference...");
 
-    // tokio::time::sleep(Duration::from_secs(30)).await;
     let command_exit_status = comfy_deps
         .inference_command
         .execute_inference(
@@ -418,6 +417,7 @@ pub async fn process_video_style_transfer_job(deps: &JobDependencies, job: &Avai
                 lipsync_enabled,
                 disable_lcm: comfy_args.disable_lcm.unwrap_or(false),
                 use_cinematic: comfy_args.use_cinematic.unwrap_or(false),
+                use_cogvideo: comfy_args.use_cogvideo.unwrap_or(false),
                 maybe_strength: comfy_args.strength,
                 frame_skip: comfy_args.frame_skip,
                 global_ipa_image_filename: global_ipa_image
@@ -435,7 +435,6 @@ pub async fn process_video_style_transfer_job(deps: &JobDependencies, job: &Avai
                     .flatten(),
         }).await;
 
-    // tokio::time::sleep(Duration::from_secs(60)).await;
     let inference_duration = Instant::now().duration_since(inference_start_time);
 
     info!("Inference command exited with status: {:?}", command_exit_status);
@@ -457,8 +456,6 @@ pub async fn process_video_style_transfer_job(deps: &JobDependencies, job: &Avai
         info!("Comfy output video dimensions: {}x{}", dimensions.width, dimensions.height);
     }
 
-    // tokio::time::sleep(Duration::from_secs(10)).await;
-    // preview_frames_uploader.await.map_err(|e| ProcessSingleJobError::Other(anyhow!("bucket uploader error: {:?}", e)))?;
     let _ = preview_cancellation_tx.send(());
 
     // ==================== CHECK OUTPUT FILE ======================== //

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/coordinate_workflow_args.rs

@@ -28,6 +28,10 @@ pub struct CoordinatedWorkflowArgs {
   /// Don't let ordinary users do this.
   pub use_cinematic: Option<bool>,
 
+  /// Use cogvideo
+  /// Only for staff
+  pub use_cogvideo: Option<bool>,
+
   /// Remove watermark from the output
   /// Only for premium accounts
   pub remove_watermark: Option<bool>,
@@ -74,6 +78,7 @@ fn handle_flags(args: &mut CoordinatedWorkflowArgs, is_staff: bool) {
     // Non-staff cannot use these workflows
     args.disable_lcm = None;
     args.remove_watermark = None;
+    args.use_cogvideo = None;
   }
 
   if args.use_upscaler == Some(true) {
@@ -109,7 +114,7 @@ fn handle_flags(args: &mut CoordinatedWorkflowArgs, is_staff: bool) {
 
 #[cfg(test)]
 mod tests {
-  use super::{CoordinatedWorkflowArgs, coordinate_workflow_args};
+  use super::{coordinate_workflow_args, CoordinatedWorkflowArgs};
 
   mod prompts {
     use super::*;
@@ -183,6 +188,7 @@ mod tests {
         use_upscaler: None,
         disable_lcm: None,
         use_cinematic: None,
+        use_cogvideo: None,
         remove_watermark: None,
       };
 
@@ -196,6 +202,7 @@ mod tests {
       assert_eq!(coordinated_args.use_face_detailer, None);
       assert_eq!(coordinated_args.use_upscaler, None);
       assert_eq!(coordinated_args.disable_lcm, None);
+      assert_eq!(coordinated_args.use_cogvideo, None);
     }
 
     #[test]
@@ -208,6 +215,7 @@ mod tests {
         use_upscaler: Some(true),
         disable_lcm: None,
         use_cinematic: None,
+        use_cogvideo: None,
         remove_watermark: None,
       };
 
@@ -218,6 +226,7 @@ mod tests {
       assert_eq!(coordinated_args.use_upscaler, Some(true));
       assert_eq!(coordinated_args.disable_lcm, None);
       assert_eq!(coordinated_args.use_cinematic, None);
+      assert_eq!(coordinated_args.use_cogvideo, None);
     }
 
     #[test]
@@ -230,6 +239,7 @@ mod tests {
         use_upscaler: Some(true),
         disable_lcm: None,
         use_cinematic: Some(true),
+        use_cogvideo: None,
         remove_watermark: None,
       };
 
@@ -240,6 +250,7 @@ mod tests {
       assert_eq!(coordinated_args.use_upscaler, Some(true));
       assert_eq!(coordinated_args.disable_lcm, None);
       assert_eq!(coordinated_args.use_cinematic, None);
+      assert_eq!(coordinated_args.use_cogvideo, None);
     }
 
     #[test]
@@ -252,6 +263,7 @@ mod tests {
         use_upscaler: None,
         disable_lcm: Some(true),
         use_cinematic: Some(true),
+        use_cogvideo: None,
         remove_watermark: None,
       };
 
@@ -262,6 +274,7 @@ mod tests {
       assert_eq!(coordinated_args.use_upscaler, None);
       assert_eq!(coordinated_args.disable_lcm, None);
       assert_eq!(coordinated_args.use_cinematic, Some(true));
+      assert_eq!(coordinated_args.use_cogvideo, None);
     }
 
     #[test]
@@ -274,6 +287,7 @@ mod tests {
         use_upscaler: None,
         disable_lcm: None,
         use_cinematic: None,
+        use_cogvideo: None,
         remove_watermark: None,
       };
 
@@ -284,6 +298,7 @@ mod tests {
       assert_eq!(coordinated_args.use_upscaler, None);
       assert_eq!(coordinated_args.disable_lcm, None);
       assert_eq!(coordinated_args.use_cinematic, Some(true));
+      assert_eq!(coordinated_args.use_cogvideo, None);
     }
   }
 }

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_studio_workflow_handler.rs

@@ -253,6 +253,11 @@ pub async fn enqueue_studio_workflow_handler(
     ]),
     use_face_detailer: request.use_face_detailer,
     use_upscaler: request.use_upscaler,
+    use_cogvideo: lazy_any_option_true(&[
+      Box::new(|| request.use_cogvideo),
+      Box::new(|| get_request_header_optional(&http_request, "USE-COGVIDEO")
+        .map(|value| str_to_bool(&value)))
+    ]),
     remove_watermark,
   };
 
@@ -314,6 +319,7 @@ pub async fn enqueue_studio_workflow_handler(
     trim_end_seconds: None,
     target_fps: None,
     generate_fast_previews: Some(false),
+    use_cogvideo: coordinated_args.use_cogvideo,
   };
 
   info!("Creating ComfyUI job record...");

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/enqueue_video_style_transfer_workflow_handler.rs

@@ -240,6 +240,11 @@ pub async fn enqueue_video_style_transfer_workflow_handler(
     ]),
     use_face_detailer: request.use_face_detailer,
     use_upscaler: request.use_upscaler,
+    use_cogvideo: lazy_any_option_true(&[
+      Box::new(|| request.use_cogvideo),
+      Box::new(|| get_request_header_optional(&http_request, "USE-COGVIDEO")
+        .map(|value| str_to_bool(&value)))
+    ]),
     remove_watermark: request.remove_watermark,
   };
 
@@ -314,6 +319,7 @@ pub async fn enqueue_video_style_transfer_workflow_handler(
     trim_end_seconds: None,
     target_fps: None,
     generate_fast_previews: Some(false),
+    use_cogvideo: coordinated_args.use_cogvideo,
   };
 
   info!("Creating ComfyUI job record...");

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue/vst_common/vst_request.rs

@@ -76,6 +76,9 @@ pub struct VstRequest {
   /// Only for premium accounts
   pub use_upscaler: Option<bool>,
 
+  /// Use cogvideo
+  pub use_cogvideo: Option<bool>,
+
   /// Optional visibility setting override.
   pub creator_set_visibility: Option<Visibility>,
 

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui_handler.rs

@@ -289,6 +289,7 @@ pub async fn enqueue_comfy_ui_handler(
         frame_skip: None,
         watermark_type: None,
         generate_fast_previews: Some(false),
+        use_cogvideo: None,
     };
 
     info!("Creating ComfyUI job record...");

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_video_style_transfer_handler.rs

@@ -381,6 +381,7 @@ pub async fn enqueue_video_style_transfer_handler(
         ]),
         use_face_detailer: request.use_face_detailer,
         use_upscaler: request.use_upscaler,
+        use_cogvideo: None,
         remove_watermark: request.remove_watermark,
     };
 
@@ -442,6 +443,7 @@ pub async fn enqueue_video_style_transfer_handler(
         target_fps: None,
         watermark_type: None,
         generate_fast_previews: Some(false),
+        use_cogvideo: None,
     };
 
     info!("Creating ComfyUI job record...");

crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_workflow_upload_request.rs

@@ -259,6 +259,7 @@ pub async fn enqueue_workflow_upload_request(
         frame_skip: None,
         watermark_type: None,
         generate_fast_previews: Some(false),
+        use_cogvideo: None,
     };
    
     // create the inference args here