Merge branch 'main' into fix_rocm_documentation

Author: ArKam

invokeai/app/api/routers/session_queue.py

@@ -7,7 +7,6 @@
 from invokeai.app.api.dependencies import ApiDependencies
 from invokeai.app.services.session_processor.session_processor_common import SessionProcessorStatus
 from invokeai.app.services.session_queue.session_queue_common import (
-    QUEUE_ORDER_BY,
     Batch,
     BatchStatus,
     CancelAllExceptCurrentResult,
@@ -92,21 +91,18 @@ async def list_all_queue_items(
 
 @session_queue_router.get(
     "/{queue_id}/item_ids",
-    operation_id="get_queue_itemIds",
+    operation_id="get_queue_item_ids",
     responses={
         200: {"model": ItemIdsResult},
     },
 )
 async def get_queue_item_ids(
     queue_id: str = Path(description="The queue id to perform this operation on"),
-    order_by: QUEUE_ORDER_BY = Query(default="created_at", description="The sort field"),
     order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"),
 ) -> ItemIdsResult:
     """Gets all queue item ids that match the given parameters"""
     try:
-        return ApiDependencies.invoker.services.session_queue.get_queue_item_ids(
-            queue_id=queue_id, order_by=order_by, order_dir=order_dir
-        )
+        return ApiDependencies.invoker.services.session_queue.get_queue_item_ids(queue_id=queue_id, order_dir=order_dir)
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Unexpected error while listing all queue item ids: {e}")
 
@@ -130,7 +126,9 @@ async def get_queue_items_by_item_ids(
         queue_items: list[SessionQueueItem] = []
         for item_id in item_ids:
             try:
-                queue_item = session_queue_service.get_queue_item(item_id)
+                queue_item = session_queue_service.get_queue_item(item_id=item_id)
+                if queue_item.queue_id != queue_id:  # Auth protection for items from other queues
+                    continue
                 queue_items.append(queue_item)
             except Exception:
                 # Skip missing queue items - they may have been deleted between item id fetch and queue item fetch
@@ -376,7 +374,10 @@ async def get_queue_item(
 ) -> SessionQueueItem:
     """Gets a queue item"""
     try:
-        return ApiDependencies.invoker.services.session_queue.get_queue_item(item_id)
+        queue_item = ApiDependencies.invoker.services.session_queue.get_queue_item(item_id=item_id)
+        if queue_item.queue_id != queue_id:
+            raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
+        return queue_item
     except SessionQueueItemNotFoundError:
         raise HTTPException(status_code=404, detail=f"Queue item with id {item_id} not found in queue {queue_id}")
     except Exception as e:

invokeai/app/invocations/fields.py

@@ -1,11 +1,12 @@
 from enum import Enum
 from typing import Any, Callable, Optional, Tuple
 
-from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter, model_validator
+from pydantic import BaseModel, ConfigDict, Field, RootModel, TypeAdapter
 from pydantic.fields import _Unset
 from pydantic_core import PydanticUndefined
 
 from invokeai.app.util.metaenum import MetaEnum
+from invokeai.backend.image_util.segment_anything.shared import BoundingBox
 from invokeai.backend.util.logging import InvokeAILogger
 
 logger = InvokeAILogger.get_logger()
@@ -331,14 +332,9 @@ class ConditioningField(BaseModel):
     )
 
 
-class BoundingBoxField(BaseModel):
+class BoundingBoxField(BoundingBox):
     """A bounding box primitive value."""
 
-    x_min: int = Field(ge=0, description="The minimum x-coordinate of the bounding box (inclusive).")
-    x_max: int = Field(ge=0, description="The maximum x-coordinate of the bounding box (exclusive).")
-    y_min: int = Field(ge=0, description="The minimum y-coordinate of the bounding box (inclusive).")
-    y_max: int = Field(ge=0, description="The maximum y-coordinate of the bounding box (exclusive).")
-
     score: Optional[float] = Field(
         default=None,
         ge=0.0,
@@ -347,21 +343,6 @@ class BoundingBoxField(BaseModel):
         "when the bounding box was produced by a detector and has an associated confidence score.",
     )
 
-    @model_validator(mode="after")
-    def check_coords(self):
-        if self.x_min > self.x_max:
-            raise ValueError(f"x_min ({self.x_min}) is greater than x_max ({self.x_max}).")
-        if self.y_min > self.y_max:
-            raise ValueError(f"y_min ({self.y_min}) is greater than y_max ({self.y_max}).")
-        return self
-
-    def tuple(self) -> Tuple[int, int, int, int]:
-        """
-        Returns the bounding box as a tuple suitable for use with PIL's `Image.crop()` method.
-        This method returns a tuple of the form (left, upper, right, lower) == (x_min, y_min, x_max, y_max).
-        """
-        return (self.x_min, self.y_min, self.x_max, self.y_max)
-
 
 class MetadataField(RootModel[dict[str, Any]]):
     """

invokeai/app/invocations/segment_anything.py

@@ -1,72 +1,75 @@
-from enum import Enum
+from itertools import zip_longest
 from pathlib import Path
 from typing import Literal
 
 import numpy as np
 import torch
 from PIL import Image
-from pydantic import BaseModel, Field
-from transformers import AutoProcessor
+from pydantic import BaseModel, Field, model_validator
 from transformers.models.sam import SamModel
 from transformers.models.sam.processing_sam import SamProcessor
+from transformers.models.sam2 import Sam2Model
+from transformers.models.sam2.processing_sam2 import Sam2Processor
 
 from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation
 from invokeai.app.invocations.fields import BoundingBoxField, ImageField, InputField, TensorField
 from invokeai.app.invocations.primitives import MaskOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.image_util.segment_anything.mask_refinement import mask_to_polygon, polygon_to_mask
+from invokeai.backend.image_util.segment_anything.segment_anything_2_pipeline import SegmentAnything2Pipeline
 from invokeai.backend.image_util.segment_anything.segment_anything_pipeline import SegmentAnythingPipeline
-
-SegmentAnythingModelKey = Literal["segment-anything-base", "segment-anything-large", "segment-anything-huge"]
+from invokeai.backend.image_util.segment_anything.shared import SAMInput, SAMPoint
+
+SegmentAnythingModelKey = Literal[
+    "segment-anything-base",
+    "segment-anything-large",
+    "segment-anything-huge",
+    "segment-anything-2-tiny",
+    "segment-anything-2-small",
+    "segment-anything-2-base",
+    "segment-anything-2-large",
+]
 SEGMENT_ANYTHING_MODEL_IDS: dict[SegmentAnythingModelKey, str] = {
     "segment-anything-base": "facebook/sam-vit-base",
     "segment-anything-large": "facebook/sam-vit-large",
     "segment-anything-huge": "facebook/sam-vit-huge",
+    "segment-anything-2-tiny": "facebook/sam2.1-hiera-tiny",
+    "segment-anything-2-small": "facebook/sam2.1-hiera-small",
+    "segment-anything-2-base": "facebook/sam2.1-hiera-base-plus",
+    "segment-anything-2-large": "facebook/sam2.1-hiera-large",
 }
 
 
-class SAMPointLabel(Enum):
-    negative = -1
-    neutral = 0
-    positive = 1
-
-
-class SAMPoint(BaseModel):
-    x: int = Field(..., description="The x-coordinate of the point")
-    y: int = Field(..., description="The y-coordinate of the point")
-    label: SAMPointLabel = Field(..., description="The label of the point")
-
-
 class SAMPointsField(BaseModel):
-    points: list[SAMPoint] = Field(..., description="The points of the object")
+    points: list[SAMPoint] = Field(..., description="The points of the object", min_length=1)
 
-    def to_list(self) -> list[list[int]]:
+    def to_list(self) -> list[list[float]]:
         return [[point.x, point.y, point.label.value] for point in self.points]
 
 
 @invocation(
     "segment_anything",
     title="Segment Anything",
-    tags=["prompt", "segmentation"],
+    tags=["prompt", "segmentation", "sam", "sam2"],
     category="segmentation",
-    version="1.2.0",
+    version="1.3.0",
 )
 class SegmentAnythingInvocation(BaseInvocation):
-    """Runs a Segment Anything Model."""
+    """Runs a Segment Anything Model (SAM or SAM2)."""
 
     # Reference:
     # - https://arxiv.org/pdf/2304.02643
     # - https://huggingface.co/docs/transformers/v4.43.3/en/model_doc/grounding-dino#grounded-sam
     # - https://github.com/NielsRogge/Transformers-Tutorials/blob/a39f33ac1557b02ebfb191ea7753e332b5ca933f/Grounding%20DINO/GroundingDINO_with_Segment_Anything.ipynb
 
-    model: SegmentAnythingModelKey = InputField(description="The Segment Anything model to use.")
+    model: SegmentAnythingModelKey = InputField(description="The Segment Anything model to use (SAM or SAM2).")
     image: ImageField = InputField(description="The image to segment.")
     bounding_boxes: list[BoundingBoxField] | None = InputField(
-        default=None, description="The bounding boxes to prompt the SAM model with."
+        default=None, description="The bounding boxes to prompt the model with."
     )
     point_lists: list[SAMPointsField] | None = InputField(
         default=None,
-        description="The list of point lists to prompt the SAM model with. Each list of points represents a single object.",
+        description="The list of point lists to prompt the model with. Each list of points represents a single object.",
     )
     apply_polygon_refinement: bool = InputField(
         description="Whether to apply polygon refinement to the masks. This will smooth the edges of the masks slightly and ensure that each mask consists of a single closed polygon (before merging).",
@@ -77,14 +80,18 @@ class SegmentAnythingInvocation(BaseInvocation):
         default="all",
     )
 
+    @model_validator(mode="after")
+    def validate_points_and_boxes_len(self):
+        if self.point_lists is not None and self.bounding_boxes is not None:
+            if len(self.point_lists) != len(self.bounding_boxes):
+                raise ValueError("If both point_lists and bounding_boxes are provided, they must have the same length.")
+        return self
+
     @torch.no_grad()
     def invoke(self, context: InvocationContext) -> MaskOutput:
         # The models expect a 3-channel RGB image.
         image_pil = context.images.get_pil(self.image.image_name, mode="RGB")
 
-        if self.point_lists is not None and self.bounding_boxes is not None:
-            raise ValueError("Only one of point_lists or bounding_box can be provided.")
-
         if (not self.bounding_boxes or len(self.bounding_boxes) == 0) and (
             not self.point_lists or len(self.point_lists) == 0
         ):
@@ -111,26 +118,38 @@ def _load_sam_model(model_path: Path):
             # model, and figure out how to make it work in the pipeline.
             # torch_dtype=TorchDevice.choose_torch_dtype(),
         )
-
-        sam_processor = AutoProcessor.from_pretrained(model_path, local_files_only=True)
-        assert isinstance(sam_processor, SamProcessor)
+        sam_processor = SamProcessor.from_pretrained(model_path, local_files_only=True)
         return SegmentAnythingPipeline(sam_model=sam_model, sam_processor=sam_processor)
 
+    @staticmethod
+    def _load_sam_2_model(model_path: Path):
+        sam2_model = Sam2Model.from_pretrained(model_path, local_files_only=True)
+        sam2_processor = Sam2Processor.from_pretrained(model_path, local_files_only=True)
+        return SegmentAnything2Pipeline(sam2_model=sam2_model, sam2_processor=sam2_processor)
+
     def _segment(self, context: InvocationContext, image: Image.Image) -> list[torch.Tensor]:
-        """Use Segment Anything (SAM) to generate masks given an image + a set of bounding boxes."""
-        # Convert the bounding boxes to the SAM input format.
-        sam_bounding_boxes = (
-            [[bb.x_min, bb.y_min, bb.x_max, bb.y_max] for bb in self.bounding_boxes] if self.bounding_boxes else None
-        )
-        sam_points = [p.to_list() for p in self.point_lists] if self.point_lists else None
+        """Use Segment Anything (SAM or SAM2) to generate masks given an image + a set of bounding boxes."""
+
+        source = SEGMENT_ANYTHING_MODEL_IDS[self.model]
+        inputs: list[SAMInput] = []
+        for bbox_field, point_field in zip_longest(self.bounding_boxes or [], self.point_lists or [], fillvalue=None):
+            inputs.append(
+                SAMInput(
+                    bounding_box=bbox_field,
+                    points=point_field.points if point_field else None,
+                )
+            )
 
-        with (
-            context.models.load_remote_model(
-                source=SEGMENT_ANYTHING_MODEL_IDS[self.model], loader=SegmentAnythingInvocation._load_sam_model
-            ) as sam_pipeline,
-        ):
-            assert isinstance(sam_pipeline, SegmentAnythingPipeline)
-            masks = sam_pipeline.segment(image=image, bounding_boxes=sam_bounding_boxes, point_lists=sam_points)
+        if "sam2" in source:
+            loader = SegmentAnythingInvocation._load_sam_2_model
+            with context.models.load_remote_model(source=source, loader=loader) as pipeline:
+                assert isinstance(pipeline, SegmentAnything2Pipeline)
+                masks = pipeline.segment(image=image, inputs=inputs)
+        else:
+            loader = SegmentAnythingInvocation._load_sam_model
+            with context.models.load_remote_model(source=source, loader=loader) as pipeline:
+                assert isinstance(pipeline, SegmentAnythingPipeline)
+                masks = pipeline.segment(image=image, inputs=inputs)
 
         masks = self._process_masks(masks)
         if self.apply_polygon_refinement:

invokeai/app/services/bulk_download/bulk_download_default.py

@@ -150,4 +150,15 @@ def get_path(self, bulk_download_item_name: str) -> str:
     def _is_valid_path(self, path: Union[str, Path]) -> bool:
         """Validates the path given for a bulk download."""
         path = path if isinstance(path, Path) else Path(path)
-        return path.exists()
+
+        # Resolve the path to handle any path traversal attempts (e.g., ../)
+        resolved_path = path.resolve()
+
+        # The path may not traverse out of the bulk downloads folder or its subfolders
+        does_not_traverse = resolved_path.parent == self._bulk_downloads_folder.resolve()
+
+        # The path must exist and be a .zip file
+        does_exist = resolved_path.exists()
+        is_zip_file = resolved_path.suffix == ".zip"
+
+        return does_exist and is_zip_file and does_not_traverse

invokeai/app/services/session_queue/session_queue_base.py

@@ -2,7 +2,6 @@
 from typing import Any, Coroutine, Optional
 
 from invokeai.app.services.session_queue.session_queue_common import (
-    QUEUE_ORDER_BY,
     Batch,
     BatchStatus,
     CancelAllExceptCurrentResult,
@@ -149,15 +148,14 @@ def list_all_queue_items(
     def get_queue_item_ids(
         self,
         queue_id: str,
-        order_by: QUEUE_ORDER_BY = "created_at",
         order_dir: SQLiteDirection = SQLiteDirection.Descending,
     ) -> ItemIdsResult:
         """Gets all queue item ids that match the given parameters"""
         pass
 
     @abstractmethod
     def get_queue_item(self, item_id: int) -> SessionQueueItem:
-        """Gets a session queue item by ID"""
+        """Gets a session queue item by ID for a given queue"""
         pass
 
     @abstractmethod

invokeai/app/services/session_queue/session_queue_common.py

@@ -174,7 +174,6 @@ def validate_graph(cls, v: Graph):
 
 DEFAULT_QUEUE_ID = "default"
 
-QUEUE_ORDER_BY = Literal["created_at", "completed_at"]
 QUEUE_ITEM_STATUS = Literal["pending", "in_progress", "completed", "failed", "canceled"]
 
 

invokeai/app/services/session_queue/session_queue_sqlite.py

@@ -10,7 +10,6 @@
 from invokeai.app.services.session_queue.session_queue_common import (
     DEFAULT_QUEUE_ID,
     QUEUE_ITEM_STATUS,
-    QUEUE_ORDER_BY,
     Batch,
     BatchStatus,
     CancelAllExceptCurrentResult,
@@ -623,15 +622,14 @@ def list_all_queue_items(
     def get_queue_item_ids(
         self,
         queue_id: str,
-        order_by: QUEUE_ORDER_BY = "created_at",
         order_dir: SQLiteDirection = SQLiteDirection.Descending,
     ) -> ItemIdsResult:
         with self._db.transaction() as cursor_:
             query = f"""--sql
                 SELECT item_id
                 FROM session_queue
                 WHERE queue_id = ?
-                ORDER BY {order_by} {order_dir.value}
+                ORDER BY created_at {order_dir.value}
                 """
             query_params = [queue_id]