Merge remote-tracking branch 'origin/main' into maryhipp/flux-kontext

Author: Mary Hipp

invokeai/app/api/routers/images.py

@@ -99,7 +99,9 @@ async def upload_image(
             raise HTTPException(status_code=400, detail="Invalid resize_to format or size")
 
         try:
-            np_image = pil_to_np(pil_image)
+            # heuristic_resize_fast expects an RGB or RGBA image
+            pil_rgba = pil_image.convert("RGBA")
+            np_image = pil_to_np(pil_rgba)
             np_image = heuristic_resize_fast(np_image, (resize_dims.width, resize_dims.height))
             pil_image = np_to_pil(np_image)
         except Exception:

invokeai/app/api_app.py

@@ -158,7 +158,7 @@ def overridden_redoc() -> HTMLResponse:
 try:
     app.mount("/", NoCacheStaticFiles(directory=Path(web_root_path, "dist"), html=True), name="ui")
 except RuntimeError:
-    logger.warn(f"No UI found at {web_root_path}/dist, skipping UI mount")
+    logger.warning(f"No UI found at {web_root_path}/dist, skipping UI mount")
 app.mount(
     "/static", NoCacheStaticFiles(directory=Path(web_root_path, "static/")), name="static"
 )  # docs favicon is in here

invokeai/app/invocations/baseinvocation.py

@@ -499,7 +499,7 @@ def validate_fields(model_fields: dict[str, FieldInfo], model_type: str) -> None
 
         ui_type = field.json_schema_extra.get("ui_type", None)
         if isinstance(ui_type, str) and ui_type.startswith("DEPRECATED_"):
-            logger.warn(f'"UIType.{ui_type.split("_")[-1]}" is deprecated, ignoring')
+            logger.warning(f'"UIType.{ui_type.split("_")[-1]}" is deprecated, ignoring')
             field.json_schema_extra.pop("ui_type")
     return None
 
@@ -582,14 +582,16 @@ def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
 
         fields: dict[str, tuple[Any, FieldInfo]] = {}
 
+        original_model_fields: dict[str, OriginalModelField] = {}
+
         for field_name, field_info in cls.model_fields.items():
             annotation = field_info.annotation
             assert annotation is not None, f"{field_name} on invocation {invocation_type} has no type annotation."
             assert isinstance(field_info.json_schema_extra, dict), (
                 f"{field_name} on invocation {invocation_type} has a non-dict json_schema_extra, did you forget to use InputField?"
             )
 
-            cls._original_model_fields[field_name] = OriginalModelField(annotation=annotation, field_info=field_info)
+            original_model_fields[field_name] = OriginalModelField(annotation=annotation, field_info=field_info)
 
             validate_field_default(cls.__name__, field_name, invocation_type, annotation, field_info)
 
@@ -613,7 +615,7 @@ def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
                 raise InvalidVersionError(f'Invalid version string for node "{invocation_type}": "{version}"') from e
             uiconfig["version"] = version
         else:
-            logger.warn(f'No version specified for node "{invocation_type}", using "1.0.0"')
+            logger.warning(f'No version specified for node "{invocation_type}", using "1.0.0"')
             uiconfig["version"] = "1.0.0"
 
         cls.UIConfig = UIConfigBase(**uiconfig)
@@ -676,6 +678,7 @@ def wrapper(cls: Type[TBaseInvocation]) -> Type[TBaseInvocation]:
         docstring = cls.__doc__
         new_class = create_model(cls.__qualname__, __base__=cls, __module__=cls.__module__, **fields)  # type: ignore
         new_class.__doc__ = docstring
+        new_class._original_model_fields = original_model_fields
 
         InvocationRegistry.register_invocation(new_class)
 

invokeai/app/invocations/compel.py

@@ -114,6 +114,13 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
 
             c, _options = compel.build_conditioning_tensor_for_conjunction(conjunction)
 
+        del compel
+        del patched_tokenizer
+        del tokenizer
+        del ti_manager
+        del text_encoder
+        del text_encoder_info
+
         c = c.detach().to("cpu")
 
         conditioning_data = ConditioningFieldData(conditionings=[BasicConditioningInfo(embeds=c)])
@@ -222,7 +229,10 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
             else:
                 c_pooled = None
 
+        del compel
+        del patched_tokenizer
         del tokenizer
+        del ti_manager
         del text_encoder
         del text_encoder_info
 

invokeai/app/invocations/fields.py

@@ -438,7 +438,7 @@ class WithWorkflow:
     workflow = None
 
     def __init_subclass__(cls) -> None:
-        logger.warn(
+        logger.warning(
             f"{cls.__module__.split('.')[0]}.{cls.__name__}: WithWorkflow is deprecated. Use `context.workflow` to access the workflow."
         )
         super().__init_subclass__()
@@ -579,7 +579,7 @@ def InputField(
 
     if default_factory is not _Unset and default_factory is not None:
         default = default_factory()
-        logger.warn('"default_factory" is not supported, calling it now to set "default"')
+        logger.warning('"default_factory" is not supported, calling it now to set "default"')
 
     # These are the args we may wish pass to the pydantic `Field()` function
     field_args = {

invokeai/app/services/config/config_default.py

@@ -24,7 +24,6 @@
 INIT_FILE = Path("invokeai.yaml")
 DB_FILE = Path("invokeai.db")
 LEGACY_INIT_FILE = Path("invokeai.init")
-DEVICE = Literal["auto", "cpu", "cuda", "cuda:1", "mps"]
 PRECISION = Literal["auto", "float16", "bfloat16", "float32"]
 ATTENTION_TYPE = Literal["auto", "normal", "xformers", "sliced", "torch-sdp"]
 ATTENTION_SLICE_SIZE = Literal["auto", "balanced", "max", 1, 2, 3, 4, 5, 6, 7, 8]
@@ -93,7 +92,7 @@ class InvokeAIAppConfig(BaseSettings):
         vram: DEPRECATED: This setting is no longer used. It has been replaced by `max_cache_vram_gb`, but most users will not need to use this config since automatic cache size limits should work well in most cases. This config setting will be removed once the new model cache behavior is stable.
         lazy_offload: DEPRECATED: This setting is no longer used. Lazy-offloading is enabled by default. This config setting will be removed once the new model cache behavior is stable.
         pytorch_cuda_alloc_conf: Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to "backend:cudaMallocAsync" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.
-        device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `cuda:1`, `mps`
+        device: Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)
         precision: Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.<br>Valid values: `auto`, `float16`, `bfloat16`, `float32`
         sequential_guidance: Whether to calculate guidance in serial instead of in parallel, lowering memory requirements.
         attention_type: Attention type.<br>Valid values: `auto`, `normal`, `xformers`, `sliced`, `torch-sdp`
@@ -176,7 +175,7 @@ class InvokeAIAppConfig(BaseSettings):
     pytorch_cuda_alloc_conf: Optional[str] = Field(default=None,            description="Configure the Torch CUDA memory allocator. This will impact peak reserved VRAM usage and performance. Setting to \"backend:cudaMallocAsync\" works well on many systems. The optimal configuration is highly dependent on the system configuration (device type, VRAM, CUDA driver version, etc.), so must be tuned experimentally.")
 
     # DEVICE
-    device:                      DEVICE = Field(default="auto",             description="Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.")
+    device:                      str = Field(default="auto",                description="Preferred execution device. `auto` will choose the device depending on the hardware platform and the installed torch capabilities.<br>Valid values: `auto`, `cpu`, `cuda`, `mps`, `cuda:N` (where N is a device number)", pattern=r"^(auto|cpu|mps|cuda(:\d+)?)$")
     precision:                PRECISION = Field(default="auto",             description="Floating point precision. `float16` will consume half the memory of `float32` but produce slightly lower-quality images. The `auto` setting will guess the proper precision based on your video card and operating system.")
 
     # GENERATION

invokeai/app/services/images/images_default.py

@@ -78,7 +78,7 @@ def create(
                         board_id=board_id, image_name=image_name
                     )
                 except Exception as e:
-                    self.__invoker.services.logger.warn(f"Failed to add image to board {board_id}: {str(e)}")
+                    self.__invoker.services.logger.warning(f"Failed to add image to board {board_id}: {str(e)}")
             self.__invoker.services.image_files.save(
                 image_name=image_name, image=image, metadata=metadata, workflow=workflow, graph=graph
             )

invokeai/app/services/model_install/model_install_default.py

@@ -148,7 +148,7 @@ def stop(self, invoker: Optional[Invoker] = None) -> None:
     def _clear_pending_jobs(self) -> None:
         for job in self.list_jobs():
             if not job.in_terminal_state:
-                self._logger.warning("Cancelling job {job.id}")
+                self._logger.warning(f"Cancelling job {job.id}")
                 self.cancel_job(job)
         while True:
             try:

invokeai/app/services/session_processor/session_processor_default.py

@@ -1,3 +1,4 @@
+import gc
 import traceback
 from contextlib import suppress
 from threading import BoundedSemaphore, Thread
@@ -439,6 +440,12 @@ def _process(
                         poll_now_event.wait(self._polling_interval)
                         continue
 
+                    # GC-ing here can reduce peak memory usage of the invoke process by freeing allocated memory blocks.
+                    # Most queue items take seconds to execute, so the relative cost of a GC is very small.
+                    # Python will never cede allocated memory back to the OS, so anything we can do to reduce the peak
+                    # allocation is well worth it.
+                    gc.collect()
+
                     self._invoker.services.logger.info(
                         f"Executing queue item {self._queue_item.item_id}, session {self._queue_item.session_id}"
                     )

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -104,11 +104,7 @@ def _get_highest_priority(self, queue_id: str) -> int:
         return cast(Union[int, None], cursor.fetchone()[0]) or 0
 
     async def enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool) -> EnqueueBatchResult:
-        return await asyncio.to_thread(self._enqueue_batch, queue_id, batch, prepend)
-
-    def _enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool) -> EnqueueBatchResult:
         try:
-            cursor = self._conn.cursor()
             # TODO: how does this work in a multi-user scenario?
             current_queue_size = self._get_current_queue_size(queue_id)
             max_queue_size = self.__invoker.services.configuration.max_queue_size
@@ -118,28 +114,29 @@ def _enqueue_batch(self, queue_id: str, batch: Batch, prepend: bool) -> EnqueueB
             if prepend:
                 priority = self._get_highest_priority(queue_id) + 1
 
-            requested_count = calc_session_count(batch)
-            values_to_insert = prepare_values_to_insert(
+            requested_count = await asyncio.to_thread(
+                calc_session_count,
+                batch=batch,
+            )
+            values_to_insert = await asyncio.to_thread(
+                prepare_values_to_insert,
                 queue_id=queue_id,
                 batch=batch,
                 priority=priority,
                 max_new_queue_items=max_new_queue_items,
             )
             enqueued_count = len(values_to_insert)
 
-            if requested_count > enqueued_count:
-                values_to_insert = values_to_insert[:max_new_queue_items]
-
-            cursor.executemany(
-                """--sql
-                INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                """,
-                values_to_insert,
-            )
-            self._conn.commit()
+            with self._conn:
+                cursor = self._conn.cursor()
+                cursor.executemany(
+                    """--sql
+                    INSERT INTO session_queue (queue_id, session, session_id, batch_id, field_values, priority, workflow, origin, destination, retried_from_item_id)
+                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    """,
+                    values_to_insert,
+                )
         except Exception:
-            self._conn.rollback()
             raise
         enqueue_result = EnqueueBatchResult(
             queue_id=queue_id,