feature: adds --composition-strength parameter to cli (#416)

imaginairy/api/generate_refiners.py

@@ -284,9 +284,9 @@ def _generate_single_image(
             if comp_image is not None:
                 result_images["composition"] = comp_img_orig
                 result_images["composition-upscaled"] = comp_image
-                comp_cutoff = 0.50
-                first_step = int((prompt.steps) * comp_cutoff)
-                noise_step = int((prompt.steps - 1) * comp_cutoff)
+                composition_strength = prompt.composition_strength
+                first_step = int((prompt.steps) * composition_strength)
+                noise_step = int((prompt.steps - 1) * composition_strength)
                 log_img(comp_img_orig, "comp_image")
                 log_img(comp_image, "comp_image_upscaled")
                 comp_image_t = pillow_img_to_torch_image(comp_image)

imaginairy/cli/edit.py

@@ -84,6 +84,7 @@ def edit_cmd(
     arg_schedules,
     make_compilation_animation,
     caption_text,
+    composition_strength,
 ):
     """
     Edit an image via AI.
@@ -140,5 +141,6 @@ def edit_cmd(
         arg_schedules,
         make_compilation_animation,
         caption_text,
+        composition_strength,
         control_inputs=control_inputs,
     )

imaginairy/cli/imagine.py

@@ -116,6 +116,7 @@ def imagine_cmd(
     arg_schedules,
     make_compilation_animation,
     caption_text,
+    composition_strength,
     control_image,
     control_image_raw,
     control_strength,
@@ -209,6 +210,7 @@ def imagine_cmd(
         arg_schedules,
         make_compilation_animation,
         caption_text,
+        composition_strength,
         control_inputs=control_inputs,
         videogen=videogen,
     )

imaginairy/cli/shared.py

@@ -59,6 +59,7 @@ def _imagine_cmd(
     model_weights_path,
     model_architecture,
     prompt_library_path,
+    composition_strength,
     version=False,
     make_gif=False,
     make_compare_gif=False,
@@ -186,6 +187,7 @@ def _imagine_cmd(
                     allow_compose_phase=allow_compose_phase,
                     model_weights=model_weights_path,
                     caption_text=caption_text,
+                    composition_strength=composition_strength,
                 )
                 from imaginairy.utils.prompt_schedules import (
                     parse_schedule_strs,
@@ -501,7 +503,14 @@ common_options = [
         "--caption-text",
         "caption_text",
         default=None,
-        help="Specify the text to write onto the image",
+        help="Specify the text to write onto the image.",
         type=str,
     ),
+    click.option(
+        "--composition-strength",
+        default=None,
+        show_default=False,
+        type=float,
+        help=("Strength of the composition phase."),
+    ),
 ]

imaginairy/schema.py

@@ -299,6 +299,7 @@ class ImaginePrompt(BaseModel, protected_namespaces=()):
     caption_text: str = Field(
         "", description="text to be overlaid on the image", validate_default=True
     )
+    composition_strength: float = Field(default=0.5, ge=0, le=1, validate_default=True)
     inpaint_method: InpaintMethod = "finetune"
 
     def __init__(
@@ -329,6 +330,7 @@ class ImaginePrompt(BaseModel, protected_namespaces=()):
         is_intermediate: bool = False,
         collect_progress_latents: bool = False,
         caption_text: str = "",
+        composition_strength: float | None = 0.5,
         inpaint_method: InpaintMethod = "finetune",
     ):
         super().__init__(
@@ -357,6 +359,7 @@ class ImaginePrompt(BaseModel, protected_namespaces=()):
             is_intermediate=is_intermediate,
             collect_progress_latents=collect_progress_latents,
             caption_text=caption_text,
+            composition_strength=composition_strength,
             inpaint_method=inpaint_method,
         )