Merge pull request #108 from newfla/feat_hidream_O1_image

feat: add HiDream O1 image presets and configuration functions

Author: newfla

cli/src/main.rs

@@ -410,6 +410,8 @@ fn get_preset(args: &Args) -> Preset {
                 .try_into()
                 .unwrap(),
         ),
+        PresetDiscriminants::HiDreamO1ImageDev => Preset::HiDreamO1ImageDev,
+        PresetDiscriminants::HiDreamO1Image => Preset::HiDreamO1Image,
     };
     preset
 }

src/api.rs

@@ -489,8 +489,13 @@ pub struct ModelConfig {
     #[builder(default = "None", private)]
     diffusion_ctx: Option<(*mut sd_ctx_t, sd_ctx_params_t)>,
 
+    /// Hires fix parameters and upscaler model.
     #[builder(default = "Self::hires_init()", setter(custom))]
     hires_params: (Upscaler, HiresParams, Option<CLibPath>),
+
+    /// Extra parameters for sampling, currently used for SDXL sample params, in json string format
+    #[builder(default = "CLibString::default()")]
+    extra_sample_params: CLibString,
 }
 
 impl ModelConfigBuilder {
@@ -780,7 +785,8 @@ impl From<&ModelConfig> for ModelConfigBuilder {
                 value.hires_params.0,
                 value.hires_params.1.clone(),
                 hires_path.as_deref(),
-            );
+            )
+            .extra_sample_params(value.extra_sample_params.clone());
 
         builder.lora_models_internal(value.lora_models.clone());
 
@@ -1248,6 +1254,7 @@ fn gen_img_maybe_progress(
             custom_sigmas: model_config.sigmas.as_mut_ptr(),
             custom_sigmas_count: model_config.sigmas.len() as i32,
             flow_shift: model_config.flow_shift,
+            extra_sample_args: model_config.extra_sample_params.as_ptr(),
         };
         let control_image = sd_image_t {
             width: 0,

src/preset.rs

@@ -9,11 +9,12 @@ use crate::{
         anima, anima2, chroma, chroma_radiance, diff_instruct_star, dream_shaper_xl_2_1_turbo,
         ernie_image, ernie_image_turbo, flux_1_dev, flux_1_mini, flux_1_schnell, flux_2_dev,
         flux_2_klein_4b, flux_2_klein_9b, flux_2_klein_base_4b, flux_2_klein_base_9b,
-        juggernaut_xl_11, nitro_sd_realism, nitro_sd_vibrant, ovis_image, qwen_image, sd_turbo,
-        sdxl_base_1_0, sdxl_turbo_1_0, sdxs512_dream_shaper, segmind_vega, ssd_1b,
-        stable_diffusion_1_4, stable_diffusion_1_5, stable_diffusion_2_1,
-        stable_diffusion_3_5_large, stable_diffusion_3_5_large_turbo, stable_diffusion_3_5_medium,
-        stable_diffusion_3_medium, twinflow_z_image_turbo, z_image_turbo,
+        hi_dream_o1_image, hi_dream_o1_image_dev, juggernaut_xl_11, nitro_sd_realism,
+        nitro_sd_vibrant, ovis_image, qwen_image, sd_turbo, sdxl_base_1_0, sdxl_turbo_1_0,
+        sdxs512_dream_shaper, segmind_vega, ssd_1b, stable_diffusion_1_4, stable_diffusion_1_5,
+        stable_diffusion_2_1, stable_diffusion_3_5_large, stable_diffusion_3_5_large_turbo,
+        stable_diffusion_3_5_medium, stable_diffusion_3_medium, twinflow_z_image_turbo,
+        z_image_turbo,
     },
 };
 
@@ -300,14 +301,18 @@ pub enum Preset {
     Flux2KleinBase9B(Flux2KleinBase9BWeight),
     /// guidance_scale 9. 25 steps. 1024x1024
     SegmindVega,
-    /// cfg__scale 4.0. 30 steps 1024x1024. Vae tiling enabled
+    /// cfg_scale 4.0. 30 steps 1024x1024. Vae tiling enabled
     Anima(AnimaWeight),
-    /// cfg__scale 4.0. 30 steps 1024x1024. Vae tiling enabled
+    /// cfg_scale 4.0. 30 steps 1024x1024. Vae tiling enabled
     Anima2(Anima2Weight),
     /// cfg_scale 5.0. 20 steps 1024x1024. Vae tiling enabled. Flash attention enabled.
     ErnieImage(ErnieImageWeight),
     /// cfg_scale 1.0. 8 steps 1024x1024. Vae tiling enabled. Flash attention enabled.
     ErnieImageTurbo(ErnieImageWeight),
+    /// cfg_scale 1.0. 20 steps 1024x1024.
+    HiDreamO1ImageDev,
+    /// cfg_scale 1.0. 20 steps 1024x1024.
+    HiDreamO1Image,
 }
 
 impl Preset {
@@ -349,6 +354,8 @@ impl Preset {
             Preset::Anima2(sd_type_t) => anima2(sd_type_t),
             Preset::ErnieImage(sd_type_t) => ernie_image(sd_type_t),
             Preset::ErnieImageTurbo(sd_type_t) => ernie_image_turbo(sd_type_t),
+            Preset::HiDreamO1ImageDev => hi_dream_o1_image_dev(),
+            Preset::HiDreamO1Image => hi_dream_o1_image(),
         }
     }
 }

src/preset_builder.rs

@@ -1521,3 +1521,29 @@ fn ernie_image_llm(sd_type: ErnieImageWeight) -> Result<PathBuf, ApiError> {
         ),
     }
 }
+
+pub fn hi_dream_o1_image_dev() -> Result<ConfigsBuilder, ApiError> {
+    let model = download_file_hf_hub(
+        "Comfy-Org/HiDream-O1-Image",
+        "checkpoints/hidream_o1_image_dev_bf16.safetensors",
+    )?;
+    let mut config = ConfigBuilder::default();
+    let mut model_config = ModelConfigBuilder::default();
+    model_config.model(model);
+    config.cfg_scale(1.0).steps(20).height(1024).width(1024);
+
+    Ok((config, model_config))
+}
+
+pub fn hi_dream_o1_image() -> Result<ConfigsBuilder, ApiError> {
+    let model = download_file_hf_hub(
+        "Comfy-Org/HiDream-O1-Image",
+        "checkpoints/hidream_o1_image_bf16.safetensors",
+    )?;
+    let mut config = ConfigBuilder::default();
+    let mut model_config = ModelConfigBuilder::default();
+    model_config.model(model);
+    config.cfg_scale(1.0).steps(20).height(1024).width(1024);
+
+    Ok((config, model_config))
+}