fix: edit mode and some controlnet tests

imaginairy/cli/edit.py

@@ -22,7 +22,7 @@ remove_option(edit_options, "allow_compose_phase")
 @click.argument("image_paths", metavar="PATH|URL", required=True, nargs=-1)
 @click.option(
     "--image-strength",
-    default=1,
+    default=0.1,
     show_default=False,
     type=float,
     help="Starting image strength. Between 0 and 1.",
@@ -33,7 +33,7 @@ remove_option(edit_options, "allow_compose_phase")
     "--model",
     help=f"Model to use. Should be one of {', '.join(config.MODEL_SHORT_NAMES)}, or a path to custom weights.",
     show_default=True,
-    default="edit",
+    default="SD-1.5",
 )
 @click.option(
     "--negative-prompt",
@@ -94,7 +94,17 @@ def edit_cmd(
     Same as calling `aimg imagine --model edit --init-image my-dog.jpg --init-image-strength 1` except this command
     can batch edit images.
     """
+    from imaginairy.schema import ControlNetInput
+
     allow_compose_phase = False
+    control_inputs = [
+        ControlNetInput(
+            image=None,
+            image_raw=None,
+            mode="edit",
+        )
+    ]
+
     return _imagine_cmd(
         ctx,
         prompt,
@@ -136,4 +146,5 @@ def edit_cmd(
         arg_schedules,
         make_compilation_animation,
         caption_text,
+        control_inputs=control_inputs,
     )

imaginairy/schema.py

@@ -212,7 +212,7 @@ class ControlNetInput(BaseModel):
     mode: str
     image: Optional[LazyLoadingImage] = None
     image_raw: Optional[LazyLoadingImage] = None
-    strength: int = Field(1, ge=0, le=1000)
+    strength: float = Field(1, ge=0, le=1000)
 
     # @field_validator("image", "image_raw", mode="before")
     # def validate_images(cls, v):

imaginairy/surprise_me.py

@@ -127,7 +127,7 @@ person_prompt_configs = [
 
 
 def surprise_me_prompts(
-    img, person=None, width=None, height=None, steps=30, seed=None, use_controlnet=False
+    img, person=None, width=None, height=None, steps=30, seed=None, use_controlnet=True
 ):
     if isinstance(img, str):
         if img.startswith("http"):
@@ -141,14 +141,14 @@ def surprise_me_prompts(
 
     for prompt_text, strength, kwargs in generic_prompts:
         if use_controlnet:
-            control_input = ControlNetInput(
-                mode="edit",
-            )
+            strength = 5
+            control_input = ControlNetInput(mode="edit", strength=2)
             prompts.append(
                 ImaginePrompt(
                     prompt_text,
+                    negative_prompt="",
                     init_image=img,
-                    init_image_strength=0.05,
+                    init_image_strength=0.3,
                     prompt_strength=strength,
                     control_inputs=[control_input],
                     steps=steps,