diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs index 5e43f4ff0d..9c42060b18 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs @@ -54,7 +54,12 @@ pub struct EnqueueImageToGaussianRequest { /// REQUIRED. /// The model to use. pub model: Option, - + + /// OPTIONAL. + /// The provider to use (defaults to Artcraft/Storyteller). + /// Not all (provider, model) combinations are valid. + pub provider: Option, + /// OPTIONAL. /// Text prompt for the gaussian generation. Required. pub prompt: Option, @@ -231,26 +236,32 @@ pub async fn dispatch_request( app_env_configs: &AppEnvConfigs, ) -> Result { - match request.model { - None => { - return Err(GenerateError::BadInput(BadInputReason::NoModelSpecified)); - } - Some(GaussianModel::WorldLabsMarble) => { - return handle_worldlabs_marble( + if request.model.is_none() { + return Err(GenerateError::BadInput(BadInputReason::NoModelSpecified)); + } + + let provider = request.provider.unwrap_or(GenerationProvider::Artcraft); + + match provider { + GenerationProvider::WorldLabs => { + handle_worldlabs_marble( app, app_data_root, app_env_configs, request, worldlabs_creds_manager, - ).await; + ).await } - Some(GaussianModel::Marble0p1Mini) | Some(GaussianModel::Marble0p1Plus) => { - return handle_gaussian_artcraft( + GenerationProvider::Artcraft => { + handle_gaussian_artcraft( request, app, app_env_configs, storyteller_creds_manager, - ).await; + ).await + } + _ => { + Err(GenerateError::NoProviderAvailable) } } }