diff --git a/Makefile b/Makefile index f1e81429e73..ecf101f1d55 100644 --- a/Makefile +++ b/Makefile @@ -12,12 +12,13 @@ help: @echo "mypy-all Run mypy ignoring the config in pyproject.tom but still ignoring missing imports" @echo "test Run the unit tests." @echo "update-config-docstring Update the app's config docstring so mkdocs can autogenerate it correctly." - @echo "frontend-install Install the pnpm modules needed for the front end" - @echo "frontend-build Build the frontend in order to run on localhost:9090" + @echo "frontend-install Install the pnpm modules needed for the frontend" + @echo "frontend-build Build the frontend for localhost:9090" + @echo "frontend-test Run the frontend test suite once" @echo "frontend-dev Run the frontend in developer mode on localhost:5173" @echo "frontend-typegen Generate types for the frontend from the OpenAPI schema" - @echo "frontend-prettier Format the frontend using lint:prettier" - @echo "wheel Build the wheel for the current version" + @echo "frontend-lint Run frontend checks and fixable lint/format steps" + @echo "wheel Build the wheel for the current version" @echo "tag-release Tag the GitHub repository with the current version (use at release time only!)" @echo "openapi Generate the OpenAPI schema for the app, outputting to stdout" @echo "docs Serve the mkdocs site with live reload" @@ -57,6 +58,10 @@ frontend-install: frontend-build: cd invokeai/frontend/web && pnpm build +# Run the frontend test suite once +frontend-test: + cd invokeai/frontend/web && pnpm run test:run + # Run the frontend in dev mode frontend-dev: cd invokeai/frontend/web && pnpm dev diff --git a/invokeai/app/api/dependencies.py b/invokeai/app/api/dependencies.py index 339a0ceadb4..532f7d75905 100644 --- a/invokeai/app/api/dependencies.py +++ b/invokeai/app/api/dependencies.py @@ -46,10 +46,12 @@ from invokeai.app.services.workflow_records.workflow_records_sqlite import SqliteWorkflowRecordsStorage from invokeai.app.services.workflow_thumbnails.workflow_thumbnails_disk import WorkflowThumbnailFileStorageDisk from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( + AnimaConditioningInfo, BasicConditioningInfo, CogView4ConditioningInfo, ConditioningFieldData, FLUXConditioningInfo, + QwenImageConditioningInfo, SD3ConditioningInfo, SDXLConditioningInfo, ZImageConditioningInfo, @@ -140,6 +142,8 @@ def initialize( SD3ConditioningInfo, CogView4ConditioningInfo, ZImageConditioningInfo, + QwenImageConditioningInfo, + AnimaConditioningInfo, ], ephemeral=True, ), diff --git a/invokeai/app/api/routers/auth.py b/invokeai/app/api/routers/auth.py index 8008da2bf53..e0b0c885cd2 100644 --- a/invokeai/app/api/routers/auth.py +++ b/invokeai/app/api/routers/auth.py @@ -104,7 +104,10 @@ async def get_setup_status() -> SetupStatusResponse: # In multiuser mode, check if an admin exists user_service = ApiDependencies.invoker.services.users setup_required = not user_service.has_admin() - admin_email = user_service.get_admin_email() + + # Only expose admin_email during initial setup to avoid leaking + # administrator identity on public deployments. + admin_email = user_service.get_admin_email() if setup_required else None return SetupStatusResponse( setup_required=setup_required, @@ -158,6 +161,7 @@ async def login( user_id=user.user_id, email=user.email, is_admin=user.is_admin, + remember_me=request.remember_me, ) token = create_access_token(token_data, expires_delta) diff --git a/invokeai/app/api/routers/board_images.py b/invokeai/app/api/routers/board_images.py index cb5e0ab51ab..f94e4f2437c 100644 --- a/invokeai/app/api/routers/board_images.py +++ b/invokeai/app/api/routers/board_images.py @@ -1,12 +1,53 @@ from fastapi import Body, HTTPException from fastapi.routing import APIRouter +from invokeai.app.api.auth_dependencies import CurrentUserOrDefault from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.services.images.images_common import AddImagesToBoardResult, RemoveImagesFromBoardResult board_images_router = APIRouter(prefix="/v1/board_images", tags=["boards"]) +def _assert_board_write_access(board_id: str, current_user: CurrentUserOrDefault) -> None: + """Raise 403 if the current user may not mutate the given board. + + Write access is granted when ANY of these hold: + - The user is an admin. + - The user owns the board. + - The board visibility is Public (public boards accept contributions from any user). + """ + from invokeai.app.services.board_records.board_records_common import BoardVisibility + + try: + board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id) + except Exception: + raise HTTPException(status_code=404, detail="Board not found") + if current_user.is_admin: + return + if board.user_id == current_user.user_id: + return + if board.board_visibility == BoardVisibility.Public: + return + raise HTTPException(status_code=403, detail="Not authorized to modify this board") + + +def _assert_image_direct_owner(image_name: str, current_user: CurrentUserOrDefault) -> None: + """Raise 403 if the current user is not the direct owner of the image. + + This is intentionally stricter than _assert_image_owner in images.py: + board ownership is NOT sufficient here. Allowing a user to add someone + else's image to their own board would grant them mutation rights via the + board-ownership fallback in _assert_image_owner, escalating read access + into write access. + """ + if current_user.is_admin: + return + owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name) + if owner is not None and owner == current_user.user_id: + return + raise HTTPException(status_code=403, detail="Not authorized to move this image") + + @board_images_router.post( "/", operation_id="add_image_to_board", @@ -17,14 +58,17 @@ response_model=AddImagesToBoardResult, ) async def add_image_to_board( + current_user: CurrentUserOrDefault, board_id: str = Body(description="The id of the board to add to"), image_name: str = Body(description="The name of the image to add"), ) -> AddImagesToBoardResult: """Creates a board_image""" + _assert_board_write_access(board_id, current_user) + _assert_image_direct_owner(image_name, current_user) try: added_images: set[str] = set() affected_boards: set[str] = set() - old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" + old_board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) or "none" ApiDependencies.invoker.services.board_images.add_image_to_board(board_id=board_id, image_name=image_name) added_images.add(image_name) affected_boards.add(board_id) @@ -48,13 +92,16 @@ async def add_image_to_board( response_model=RemoveImagesFromBoardResult, ) async def remove_image_from_board( + current_user: CurrentUserOrDefault, image_name: str = Body(description="The name of the image to remove", embed=True), ) -> RemoveImagesFromBoardResult: """Removes an image from its board, if it had one""" try: + old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" + if old_board_id != "none": + _assert_board_write_access(old_board_id, current_user) removed_images: set[str] = set() affected_boards: set[str] = set() - old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name) removed_images.add(image_name) affected_boards.add("none") @@ -64,6 +111,8 @@ async def remove_image_from_board( affected_boards=list(affected_boards), ) + except HTTPException: + raise except Exception: raise HTTPException(status_code=500, detail="Failed to remove image from board") @@ -78,16 +127,21 @@ async def remove_image_from_board( response_model=AddImagesToBoardResult, ) async def add_images_to_board( + current_user: CurrentUserOrDefault, board_id: str = Body(description="The id of the board to add to"), image_names: list[str] = Body(description="The names of the images to add", embed=True), ) -> AddImagesToBoardResult: """Adds a list of images to a board""" + _assert_board_write_access(board_id, current_user) try: added_images: set[str] = set() affected_boards: set[str] = set() for image_name in image_names: try: - old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" + _assert_image_direct_owner(image_name, current_user) + old_board_id = ( + ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) or "none" + ) ApiDependencies.invoker.services.board_images.add_image_to_board( board_id=board_id, image_name=image_name, @@ -96,12 +150,16 @@ async def add_images_to_board( affected_boards.add(board_id) affected_boards.add(old_board_id) + except HTTPException: + raise except Exception: pass return AddImagesToBoardResult( added_images=list(added_images), affected_boards=list(affected_boards), ) + except HTTPException: + raise except Exception: raise HTTPException(status_code=500, detail="Failed to add images to board") @@ -116,6 +174,7 @@ async def add_images_to_board( response_model=RemoveImagesFromBoardResult, ) async def remove_images_from_board( + current_user: CurrentUserOrDefault, image_names: list[str] = Body(description="The names of the images to remove", embed=True), ) -> RemoveImagesFromBoardResult: """Removes a list of images from their board, if they had one""" @@ -125,15 +184,21 @@ async def remove_images_from_board( for image_name in image_names: try: old_board_id = ApiDependencies.invoker.services.images.get_dto(image_name).board_id or "none" + if old_board_id != "none": + _assert_board_write_access(old_board_id, current_user) ApiDependencies.invoker.services.board_images.remove_image_from_board(image_name=image_name) removed_images.add(image_name) affected_boards.add("none") affected_boards.add(old_board_id) + except HTTPException: + raise except Exception: pass return RemoveImagesFromBoardResult( removed_images=list(removed_images), affected_boards=list(affected_boards), ) + except HTTPException: + raise except Exception: raise HTTPException(status_code=500, detail="Failed to remove images from board") diff --git a/invokeai/app/api/routers/boards.py b/invokeai/app/api/routers/boards.py index e93bb8b2a9b..6897e90aff4 100644 --- a/invokeai/app/api/routers/boards.py +++ b/invokeai/app/api/routers/boards.py @@ -6,7 +6,7 @@ from invokeai.app.api.auth_dependencies import CurrentUserOrDefault from invokeai.app.api.dependencies import ApiDependencies -from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecordOrderBy +from invokeai.app.services.board_records.board_records_common import BoardChanges, BoardRecordOrderBy, BoardVisibility from invokeai.app.services.boards.boards_common import BoardDTO from invokeai.app.services.image_records.image_records_common import ImageCategory from invokeai.app.services.shared.pagination import OffsetPaginatedResults @@ -56,7 +56,14 @@ async def get_board( except Exception: raise HTTPException(status_code=404, detail="Board not found") - if not current_user.is_admin and result.user_id != current_user.user_id: + # Admins can access any board. + # Owners can access their own boards. + # Shared and public boards are visible to all authenticated users. + if ( + not current_user.is_admin + and result.user_id != current_user.user_id + and result.board_visibility == BoardVisibility.Private + ): raise HTTPException(status_code=403, detail="Not authorized to access this board") return result @@ -188,7 +195,11 @@ async def list_all_board_image_names( except Exception: raise HTTPException(status_code=404, detail="Board not found") - if not current_user.is_admin and board.user_id != current_user.user_id: + if ( + not current_user.is_admin + and board.user_id != current_user.user_id + and board.board_visibility == BoardVisibility.Private + ): raise HTTPException(status_code=403, detail="Not authorized to access this board") image_names = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board( @@ -196,4 +207,15 @@ async def list_all_board_image_names( categories, is_intermediate, ) + + # For uncategorized images (board_id="none"), filter to only the caller's + # images so that one user cannot enumerate another's uncategorized images. + # Admin users can see all uncategorized images. + if board_id == "none" and not current_user.is_admin: + image_names = [ + name + for name in image_names + if ApiDependencies.invoker.services.image_records.get_user_id(name) == current_user.user_id + ] + return image_names diff --git a/invokeai/app/api/routers/images.py b/invokeai/app/api/routers/images.py index 6b11762c9ec..a3ae6fce82b 100644 --- a/invokeai/app/api/routers/images.py +++ b/invokeai/app/api/routers/images.py @@ -38,6 +38,96 @@ IMAGE_MAX_AGE = 31536000 +def _assert_image_owner(image_name: str, current_user: CurrentUserOrDefault) -> None: + """Raise 403 if the current user does not own the image and is not an admin. + + Ownership is satisfied when ANY of these hold: + - The user is an admin. + - The user is the image's direct owner (image_records.user_id). + - The user owns the board the image sits on. + - The image sits on a Public board (public boards grant mutation rights). + """ + from invokeai.app.services.board_records.board_records_common import BoardVisibility + + if current_user.is_admin: + return + owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name) + if owner is not None and owner == current_user.user_id: + return + + # Check whether the user owns the board the image belongs to, + # or the board is Public (public boards grant mutation rights). + board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) + if board_id is not None: + try: + board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id) + if board.user_id == current_user.user_id: + return + if board.board_visibility == BoardVisibility.Public: + return + except Exception: + pass + + raise HTTPException(status_code=403, detail="Not authorized to modify this image") + + +def _assert_image_read_access(image_name: str, current_user: CurrentUserOrDefault) -> None: + """Raise 403 if the current user may not view the image. + + Access is granted when ANY of these hold: + - The user is an admin. + - The user owns the image. + - The image sits on a shared or public board. + """ + from invokeai.app.services.board_records.board_records_common import BoardVisibility + + if current_user.is_admin: + return + + owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name) + if owner is not None and owner == current_user.user_id: + return + + # Check whether the image's board makes it visible to other users. + board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) + if board_id is not None: + try: + board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id) + if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public): + return + except Exception: + pass + + raise HTTPException(status_code=403, detail="Not authorized to access this image") + + +def _assert_board_read_access(board_id: str, current_user: CurrentUserOrDefault) -> None: + """Raise 403 if the current user may not read images from this board. + + Access is granted when ANY of these hold: + - The user is an admin. + - The user owns the board. + - The board visibility is Shared or Public. + """ + from invokeai.app.services.board_records.board_records_common import BoardVisibility + + if current_user.is_admin: + return + + try: + board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id) + except Exception: + raise HTTPException(status_code=404, detail="Board not found") + + if board.user_id == current_user.user_id: + return + + if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public): + return + + raise HTTPException(status_code=403, detail="Not authorized to access this board") + + class ResizeToDimensions(BaseModel): width: int = Field(..., gt=0) height: int = Field(..., gt=0) @@ -83,6 +173,22 @@ async def upload_image( ), ) -> ImageDTO: """Uploads an image for the current user""" + # If uploading into a board, verify the user has write access. + # Public boards allow uploads from any authenticated user. + if board_id is not None: + from invokeai.app.services.board_records.board_records_common import BoardVisibility + + try: + board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id) + except Exception: + raise HTTPException(status_code=404, detail="Board not found") + if ( + not current_user.is_admin + and board.user_id != current_user.user_id + and board.board_visibility != BoardVisibility.Public + ): + raise HTTPException(status_code=403, detail="Not authorized to upload to this board") + if not file.content_type or not file.content_type.startswith("image"): raise HTTPException(status_code=415, detail="Not an image") @@ -165,9 +271,11 @@ async def create_image_upload_entry( @images_router.delete("/i/{image_name}", operation_id="delete_image", response_model=DeleteImagesResult) async def delete_image( + current_user: CurrentUserOrDefault, image_name: str = Path(description="The name of the image to delete"), ) -> DeleteImagesResult: """Deletes an image""" + _assert_image_owner(image_name, current_user) deleted_images: set[str] = set() affected_boards: set[str] = set() @@ -189,26 +297,31 @@ async def delete_image( @images_router.delete("/intermediates", operation_id="clear_intermediates") -async def clear_intermediates() -> int: - """Clears all intermediates""" +async def clear_intermediates( + current_user: CurrentUserOrDefault, +) -> int: + """Clears all intermediates. Requires admin.""" + if not current_user.is_admin: + raise HTTPException(status_code=403, detail="Only admins can clear all intermediates") try: count_deleted = ApiDependencies.invoker.services.images.delete_intermediates() return count_deleted except Exception: raise HTTPException(status_code=500, detail="Failed to clear intermediates") - pass @images_router.get("/intermediates", operation_id="get_intermediates_count") -async def get_intermediates_count() -> int: - """Gets the count of intermediate images""" +async def get_intermediates_count( + current_user: CurrentUserOrDefault, +) -> int: + """Gets the count of intermediate images. Non-admin users only see their own intermediates.""" try: - return ApiDependencies.invoker.services.images.get_intermediates_count() + user_id = None if current_user.is_admin else current_user.user_id + return ApiDependencies.invoker.services.images.get_intermediates_count(user_id=user_id) except Exception: raise HTTPException(status_code=500, detail="Failed to get intermediates") - pass @images_router.patch( @@ -217,10 +330,12 @@ async def get_intermediates_count() -> int: response_model=ImageDTO, ) async def update_image( + current_user: CurrentUserOrDefault, image_name: str = Path(description="The name of the image to update"), image_changes: ImageRecordChanges = Body(description="The changes to apply to the image"), ) -> ImageDTO: """Updates an image""" + _assert_image_owner(image_name, current_user) try: return ApiDependencies.invoker.services.images.update(image_name, image_changes) @@ -234,9 +349,11 @@ async def update_image( response_model=ImageDTO, ) async def get_image_dto( + current_user: CurrentUserOrDefault, image_name: str = Path(description="The name of image to get"), ) -> ImageDTO: """Gets an image's DTO""" + _assert_image_read_access(image_name, current_user) try: return ApiDependencies.invoker.services.images.get_dto(image_name) @@ -250,9 +367,11 @@ async def get_image_dto( response_model=Optional[MetadataField], ) async def get_image_metadata( + current_user: CurrentUserOrDefault, image_name: str = Path(description="The name of image to get"), ) -> Optional[MetadataField]: """Gets an image's metadata""" + _assert_image_read_access(image_name, current_user) try: return ApiDependencies.invoker.services.images.get_metadata(image_name) @@ -269,8 +388,11 @@ class WorkflowAndGraphResponse(BaseModel): "/i/{image_name}/workflow", operation_id="get_image_workflow", response_model=WorkflowAndGraphResponse ) async def get_image_workflow( + current_user: CurrentUserOrDefault, image_name: str = Path(description="The name of image whose workflow to get"), ) -> WorkflowAndGraphResponse: + _assert_image_read_access(image_name, current_user) + try: workflow = ApiDependencies.invoker.services.images.get_workflow(image_name) graph = ApiDependencies.invoker.services.images.get_graph(image_name) @@ -306,8 +428,12 @@ async def get_image_workflow( async def get_image_full( image_name: str = Path(description="The name of full-resolution image file to get"), ) -> Response: - """Gets a full-resolution image file""" + """Gets a full-resolution image file. + This endpoint is intentionally unauthenticated because browsers load images + via tags which cannot send Bearer tokens. Image names are UUIDs, + providing security through unguessability. + """ try: path = ApiDependencies.invoker.services.images.get_path(image_name) with open(path, "rb") as f: @@ -335,8 +461,12 @@ async def get_image_full( async def get_image_thumbnail( image_name: str = Path(description="The name of thumbnail image file to get"), ) -> Response: - """Gets a thumbnail image file""" + """Gets a thumbnail image file. + This endpoint is intentionally unauthenticated because browsers load images + via tags which cannot send Bearer tokens. Image names are UUIDs, + providing security through unguessability. + """ try: path = ApiDependencies.invoker.services.images.get_path(image_name, thumbnail=True) with open(path, "rb") as f: @@ -354,9 +484,11 @@ async def get_image_thumbnail( response_model=ImageUrlsDTO, ) async def get_image_urls( + current_user: CurrentUserOrDefault, image_name: str = Path(description="The name of the image whose URL to get"), ) -> ImageUrlsDTO: """Gets an image and thumbnail URL""" + _assert_image_read_access(image_name, current_user) try: image_url = ApiDependencies.invoker.services.images.get_url(image_name) @@ -392,6 +524,11 @@ async def list_image_dtos( ) -> OffsetPaginatedResults[ImageDTO]: """Gets a list of image DTOs for the current user""" + # Validate that the caller can read from this board before listing its images. + # "none" is a sentinel for uncategorized images and is handled by the SQL layer. + if board_id is not None and board_id != "none": + _assert_board_read_access(board_id, current_user) + image_dtos = ApiDependencies.invoker.services.images.get_many( offset, limit, @@ -410,6 +547,7 @@ async def list_image_dtos( @images_router.post("/delete", operation_id="delete_images_from_list", response_model=DeleteImagesResult) async def delete_images_from_list( + current_user: CurrentUserOrDefault, image_names: list[str] = Body(description="The list of names of images to delete", embed=True), ) -> DeleteImagesResult: try: @@ -417,24 +555,31 @@ async def delete_images_from_list( affected_boards: set[str] = set() for image_name in image_names: try: + _assert_image_owner(image_name, current_user) image_dto = ApiDependencies.invoker.services.images.get_dto(image_name) board_id = image_dto.board_id or "none" ApiDependencies.invoker.services.images.delete(image_name) deleted_images.add(image_name) affected_boards.add(board_id) + except HTTPException: + raise except Exception: pass return DeleteImagesResult( deleted_images=list(deleted_images), affected_boards=list(affected_boards), ) + except HTTPException: + raise except Exception: raise HTTPException(status_code=500, detail="Failed to delete images") @images_router.delete("/uncategorized", operation_id="delete_uncategorized_images", response_model=DeleteImagesResult) -async def delete_uncategorized_images() -> DeleteImagesResult: - """Deletes all images that are uncategorized""" +async def delete_uncategorized_images( + current_user: CurrentUserOrDefault, +) -> DeleteImagesResult: + """Deletes all uncategorized images owned by the current user (or all if admin)""" image_names = ApiDependencies.invoker.services.board_images.get_all_board_image_names_for_board( board_id="none", categories=None, is_intermediate=None @@ -445,9 +590,13 @@ async def delete_uncategorized_images() -> DeleteImagesResult: affected_boards: set[str] = set() for image_name in image_names: try: + _assert_image_owner(image_name, current_user) ApiDependencies.invoker.services.images.delete(image_name) deleted_images.add(image_name) affected_boards.add("none") + except HTTPException: + # Skip images not owned by the current user + pass except Exception: pass return DeleteImagesResult( @@ -464,6 +613,7 @@ class ImagesUpdatedFromListResult(BaseModel): @images_router.post("/star", operation_id="star_images_in_list", response_model=StarredImagesResult) async def star_images_in_list( + current_user: CurrentUserOrDefault, image_names: list[str] = Body(description="The list of names of images to star", embed=True), ) -> StarredImagesResult: try: @@ -471,23 +621,29 @@ async def star_images_in_list( affected_boards: set[str] = set() for image_name in image_names: try: + _assert_image_owner(image_name, current_user) updated_image_dto = ApiDependencies.invoker.services.images.update( image_name, changes=ImageRecordChanges(starred=True) ) starred_images.add(image_name) affected_boards.add(updated_image_dto.board_id or "none") + except HTTPException: + raise except Exception: pass return StarredImagesResult( starred_images=list(starred_images), affected_boards=list(affected_boards), ) + except HTTPException: + raise except Exception: raise HTTPException(status_code=500, detail="Failed to star images") @images_router.post("/unstar", operation_id="unstar_images_in_list", response_model=UnstarredImagesResult) async def unstar_images_in_list( + current_user: CurrentUserOrDefault, image_names: list[str] = Body(description="The list of names of images to unstar", embed=True), ) -> UnstarredImagesResult: try: @@ -495,17 +651,22 @@ async def unstar_images_in_list( affected_boards: set[str] = set() for image_name in image_names: try: + _assert_image_owner(image_name, current_user) updated_image_dto = ApiDependencies.invoker.services.images.update( image_name, changes=ImageRecordChanges(starred=False) ) unstarred_images.add(image_name) affected_boards.add(updated_image_dto.board_id or "none") + except HTTPException: + raise except Exception: pass return UnstarredImagesResult( unstarred_images=list(unstarred_images), affected_boards=list(affected_boards), ) + except HTTPException: + raise except Exception: raise HTTPException(status_code=500, detail="Failed to unstar images") @@ -523,6 +684,7 @@ class ImagesDownloaded(BaseModel): "/download", operation_id="download_images_from_list", response_model=ImagesDownloaded, status_code=202 ) async def download_images_from_list( + current_user: CurrentUserOrDefault, background_tasks: BackgroundTasks, image_names: Optional[list[str]] = Body( default=None, description="The list of names of images to download", embed=True @@ -533,6 +695,16 @@ async def download_images_from_list( ) -> ImagesDownloaded: if (image_names is None or len(image_names) == 0) and board_id is None: raise HTTPException(status_code=400, detail="No images or board id specified.") + + # Validate that the caller can read every image they are requesting. + # For a board_id request, check board visibility; for explicit image names, + # check each image individually. + if board_id: + _assert_board_read_access(board_id, current_user) + if image_names: + for name in image_names: + _assert_image_read_access(name, current_user) + bulk_download_item_id: str = ApiDependencies.invoker.services.bulk_download.generate_item_id(board_id) background_tasks.add_task( @@ -540,6 +712,7 @@ async def download_images_from_list( image_names, board_id, bulk_download_item_id, + current_user.user_id, ) return ImagesDownloaded(bulk_download_item_name=bulk_download_item_id + ".zip") @@ -558,11 +731,21 @@ async def download_images_from_list( }, ) async def get_bulk_download_item( + current_user: CurrentUserOrDefault, background_tasks: BackgroundTasks, bulk_download_item_name: str = Path(description="The bulk_download_item_name of the bulk download item to get"), ) -> FileResponse: - """Gets a bulk download zip file""" + """Gets a bulk download zip file. + + Requires authentication. The caller must be the user who initiated the + download (tracked by the bulk download service) or an admin. + """ try: + # Verify the caller owns this download (or is an admin) + owner = ApiDependencies.invoker.services.bulk_download.get_owner(bulk_download_item_name) + if owner is not None and owner != current_user.user_id and not current_user.is_admin: + raise HTTPException(status_code=403, detail="Not authorized to access this download") + path = ApiDependencies.invoker.services.bulk_download.get_path(bulk_download_item_name) response = FileResponse( @@ -574,6 +757,8 @@ async def get_bulk_download_item( response.headers["Cache-Control"] = f"max-age={IMAGE_MAX_AGE}" background_tasks.add_task(ApiDependencies.invoker.services.bulk_download.delete, bulk_download_item_name) return response + except HTTPException: + raise except Exception: raise HTTPException(status_code=404) @@ -594,6 +779,10 @@ async def get_image_names( ) -> ImageNamesResult: """Gets ordered list of image names with metadata for optimistic updates""" + # Validate that the caller can read from this board before listing its images. + if board_id is not None and board_id != "none": + _assert_board_read_access(board_id, current_user) + try: result = ApiDependencies.invoker.services.images.get_image_names( starred_first=starred_first, @@ -617,6 +806,7 @@ async def get_image_names( responses={200: {"model": list[ImageDTO]}}, ) async def get_images_by_names( + current_user: CurrentUserOrDefault, image_names: list[str] = Body(embed=True, description="Object containing list of image names to fetch DTOs for"), ) -> list[ImageDTO]: """Gets image DTOs for the specified image names. Maintains order of input names.""" @@ -628,8 +818,12 @@ async def get_images_by_names( image_dtos: list[ImageDTO] = [] for name in image_names: try: + _assert_image_read_access(name, current_user) dto = image_service.get_dto(name) image_dtos.append(dto) + except HTTPException: + # Skip images the user is not authorized to view + continue except Exception: # Skip missing images - they may have been deleted between name fetch and DTO fetch continue diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 9536b7fc8bd..822d9655fe7 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -858,7 +858,7 @@ def generate_html(title: str, heading: str, repo_id: str, is_error: bool, messag "/install", operation_id="list_model_installs", ) -async def list_model_installs() -> List[ModelInstallJob]: +async def list_model_installs(current_admin: AdminUserOrDefault) -> List[ModelInstallJob]: """Return the list of model install jobs. Install jobs have a numeric `id`, a `status`, and other fields that provide information on @@ -890,7 +890,9 @@ async def list_model_installs() -> List[ModelInstallJob]: 404: {"description": "No such job"}, }, ) -async def get_model_install_job(id: int = Path(description="Model install id")) -> ModelInstallJob: +async def get_model_install_job( + current_admin: AdminUserOrDefault, id: int = Path(description="Model install id") +) -> ModelInstallJob: """ Return model install job corresponding to the given source. See the documentation for 'List Model Install Jobs' for information on the format of the return value. @@ -933,7 +935,9 @@ async def cancel_model_install_job( }, status_code=201, ) -async def pause_model_install_job(id: int = Path(description="Model install job ID")) -> ModelInstallJob: +async def pause_model_install_job( + current_admin: AdminUserOrDefault, id: int = Path(description="Model install job ID") +) -> ModelInstallJob: """Pause the model install job corresponding to the given job ID.""" installer = ApiDependencies.invoker.services.model_manager.install try: @@ -953,7 +957,9 @@ async def pause_model_install_job(id: int = Path(description="Model install job }, status_code=201, ) -async def resume_model_install_job(id: int = Path(description="Model install job ID")) -> ModelInstallJob: +async def resume_model_install_job( + current_admin: AdminUserOrDefault, id: int = Path(description="Model install job ID") +) -> ModelInstallJob: """Resume a paused model install job corresponding to the given job ID.""" installer = ApiDependencies.invoker.services.model_manager.install try: @@ -973,7 +979,9 @@ async def resume_model_install_job(id: int = Path(description="Model install job }, status_code=201, ) -async def restart_failed_model_install_job(id: int = Path(description="Model install job ID")) -> ModelInstallJob: +async def restart_failed_model_install_job( + current_admin: AdminUserOrDefault, id: int = Path(description="Model install job ID") +) -> ModelInstallJob: """Restart failed or non-resumable file downloads for the given job.""" installer = ApiDependencies.invoker.services.model_manager.install try: @@ -994,6 +1002,7 @@ async def restart_failed_model_install_job(id: int = Path(description="Model ins status_code=201, ) async def restart_model_install_file( + current_admin: AdminUserOrDefault, id: int = Path(description="Model install job ID"), file_source: AnyHttpUrl = Body(description="File download URL to restart"), ) -> ModelInstallJob: diff --git a/invokeai/app/api/routers/recall_parameters.py b/invokeai/app/api/routers/recall_parameters.py index 0af3fd29b0c..ec08adba2e8 100644 --- a/invokeai/app/api/routers/recall_parameters.py +++ b/invokeai/app/api/routers/recall_parameters.py @@ -7,6 +7,7 @@ from fastapi.routing import APIRouter from pydantic import BaseModel, ConfigDict, Field +from invokeai.app.api.auth_dependencies import CurrentUserOrDefault from invokeai.app.api.dependencies import ApiDependencies from invokeai.backend.image_util.controlnet_processor import process_controlnet_image from invokeai.backend.model_manager.taxonomy import ModelType @@ -291,12 +292,58 @@ def resolve_ip_adapter_models(ip_adapters: list[IPAdapterRecallParameter]) -> li return resolved_adapters +def _assert_recall_image_access(parameters: "RecallParameter", current_user: CurrentUserOrDefault) -> None: + """Validate that the caller can read every image referenced in the recall parameters. + + Control layers and IP adapters may reference image_name fields. Without this + check an attacker who knows another user's image UUID could use the recall + endpoint to extract image dimensions and — for ControlNet preprocessors — mint + a derived processed image they can then fetch. + """ + from invokeai.app.services.board_records.board_records_common import BoardVisibility + + image_names: list[str] = [] + if parameters.control_layers: + for layer in parameters.control_layers: + if layer.image_name is not None: + image_names.append(layer.image_name) + if parameters.ip_adapters: + for adapter in parameters.ip_adapters: + if adapter.image_name is not None: + image_names.append(adapter.image_name) + + if not image_names: + return + + # Admin can access all images + if current_user.is_admin: + return + + for image_name in image_names: + owner = ApiDependencies.invoker.services.image_records.get_user_id(image_name) + if owner is not None and owner == current_user.user_id: + continue + + # Check board visibility + board_id = ApiDependencies.invoker.services.board_image_records.get_board_for_image(image_name) + if board_id is not None: + try: + board = ApiDependencies.invoker.services.boards.get_dto(board_id=board_id) + if board.board_visibility in (BoardVisibility.Shared, BoardVisibility.Public): + continue + except Exception: + pass + + raise HTTPException(status_code=403, detail=f"Not authorized to access image {image_name}") + + @recall_parameters_router.post( "/{queue_id}", operation_id="update_recall_parameters", response_model=dict[str, Any], ) async def update_recall_parameters( + current_user: CurrentUserOrDefault, queue_id: str = Path(..., description="The queue id to perform this operation on"), parameters: RecallParameter = Body(..., description="Recall parameters to update"), ) -> dict[str, Any]: @@ -328,6 +375,10 @@ async def update_recall_parameters( """ logger = ApiDependencies.invoker.services.logger + # Validate image access before processing — prevents information leakage + # (dimensions) and derived-image minting via ControlNet preprocessors. + _assert_recall_image_access(parameters, current_user) + try: # Get only the parameters that were actually provided (non-None values) provided_params = {k: v for k, v in parameters.model_dump().items() if v is not None} @@ -335,14 +386,14 @@ async def update_recall_parameters( if not provided_params: return {"status": "no_parameters_provided", "updated_count": 0} - # Store each parameter in client state using a consistent key format + # Store each parameter in client state scoped to the current user updated_count = 0 for param_key, param_value in provided_params.items(): # Convert parameter values to JSON strings for storage value_str = json.dumps(param_value) try: ApiDependencies.invoker.services.client_state_persistence.set_by_key( - queue_id, f"recall_{param_key}", value_str + current_user.user_id, f"recall_{param_key}", value_str ) updated_count += 1 except Exception as e: @@ -396,7 +447,9 @@ async def update_recall_parameters( logger.info( f"Emitting recall_parameters_updated event for queue {queue_id} with {len(provided_params)} parameters" ) - ApiDependencies.invoker.services.events.emit_recall_parameters_updated(queue_id, provided_params) + ApiDependencies.invoker.services.events.emit_recall_parameters_updated( + queue_id, current_user.user_id, provided_params + ) logger.info("Successfully emitted recall_parameters_updated event") except Exception as e: logger.error(f"Error emitting recall parameters event: {e}", exc_info=True) @@ -425,6 +478,7 @@ async def update_recall_parameters( response_model=dict[str, Any], ) async def get_recall_parameters( + current_user: CurrentUserOrDefault, queue_id: str = Path(..., description="The queue id to retrieve parameters for"), ) -> dict[str, Any]: """ diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index 403e7727cb4..41a5a411c7a 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -44,7 +44,8 @@ def sanitize_queue_item_for_user( """Sanitize queue item for non-admin users viewing other users' items. For non-admin users viewing queue items belonging to other users, - the field_values, session graph, and workflow should be hidden/cleared to protect privacy. + only timestamps, status, and error information are exposed. All other + fields (user identity, generation parameters, graphs, workflows) are stripped. Args: queue_item: The queue item to sanitize @@ -58,15 +59,25 @@ def sanitize_queue_item_for_user( if is_admin or queue_item.user_id == current_user_id: return queue_item - # For non-admins viewing other users' items, clear sensitive fields - # Create a shallow copy to avoid mutating the original + # For non-admins viewing other users' items, strip everything except + # item_id, queue_id, status, and timestamps sanitized_item = queue_item.model_copy(deep=False) + sanitized_item.user_id = "redacted" + sanitized_item.user_display_name = None + sanitized_item.user_email = None + sanitized_item.batch_id = "redacted" + sanitized_item.session_id = "redacted" + sanitized_item.origin = None + sanitized_item.destination = None + sanitized_item.priority = 0 sanitized_item.field_values = None + sanitized_item.retried_from_item_id = None sanitized_item.workflow = None - # Clear the session graph by replacing it with an empty graph execution state - # This prevents information leakage through the generation graph + sanitized_item.error_type = None + sanitized_item.error_message = None + sanitized_item.error_traceback = None sanitized_item.session = GraphExecutionState( - id=queue_item.session.id, + id="redacted", graph=Graph(), ) return sanitized_item @@ -126,12 +137,16 @@ async def list_all_queue_items( }, ) async def get_queue_item_ids( + current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to perform this operation on"), order_dir: SQLiteDirection = Query(default=SQLiteDirection.Descending, description="The order of sort"), ) -> ItemIdsResult: - """Gets all queue item ids that match the given parameters""" + """Gets all queue item ids that match the given parameters. Non-admin users only see their own items.""" try: - return ApiDependencies.invoker.services.session_queue.get_queue_item_ids(queue_id=queue_id, order_dir=order_dir) + user_id = None if current_user.is_admin else current_user.user_id + return ApiDependencies.invoker.services.session_queue.get_queue_item_ids( + queue_id=queue_id, order_dir=order_dir, user_id=user_id + ) except Exception as e: raise HTTPException(status_code=500, detail=f"Unexpected error while listing all queue item ids: {e}") @@ -376,11 +391,15 @@ async def prune( }, ) async def get_current_queue_item( + current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to perform this operation on"), ) -> Optional[SessionQueueItem]: """Gets the currently execution queue item""" try: - return ApiDependencies.invoker.services.session_queue.get_current(queue_id) + item = ApiDependencies.invoker.services.session_queue.get_current(queue_id) + if item is not None: + item = sanitize_queue_item_for_user(item, current_user.user_id, current_user.is_admin) + return item except Exception as e: raise HTTPException(status_code=500, detail=f"Unexpected error while getting current queue item: {e}") @@ -393,11 +412,15 @@ async def get_current_queue_item( }, ) async def get_next_queue_item( + current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to perform this operation on"), ) -> Optional[SessionQueueItem]: """Gets the next queue item, without executing it""" try: - return ApiDependencies.invoker.services.session_queue.get_next(queue_id) + item = ApiDependencies.invoker.services.session_queue.get_next(queue_id) + if item is not None: + item = sanitize_queue_item_for_user(item, current_user.user_id, current_user.is_admin) + return item except Exception as e: raise HTTPException(status_code=500, detail=f"Unexpected error while getting next queue item: {e}") @@ -413,9 +436,10 @@ async def get_queue_status( current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to perform this operation on"), ) -> SessionQueueAndProcessorStatus: - """Gets the status of the session queue""" + """Gets the status of the session queue. Non-admin users see only their own counts and cannot see current item details unless they own it.""" try: - queue = ApiDependencies.invoker.services.session_queue.get_queue_status(queue_id, user_id=current_user.user_id) + user_id = None if current_user.is_admin else current_user.user_id + queue = ApiDependencies.invoker.services.session_queue.get_queue_status(queue_id, user_id=user_id) processor = ApiDependencies.invoker.services.session_processor.get_status() return SessionQueueAndProcessorStatus(queue=queue, processor=processor) except Exception as e: @@ -430,12 +454,16 @@ async def get_queue_status( }, ) async def get_batch_status( + current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to perform this operation on"), batch_id: str = Path(description="The batch to get the status of"), ) -> BatchStatus: - """Gets the status of the session queue""" + """Gets the status of a batch. Non-admin users only see their own batches.""" try: - return ApiDependencies.invoker.services.session_queue.get_batch_status(queue_id=queue_id, batch_id=batch_id) + user_id = None if current_user.is_admin else current_user.user_id + return ApiDependencies.invoker.services.session_queue.get_batch_status( + queue_id=queue_id, batch_id=batch_id, user_id=user_id + ) except Exception as e: raise HTTPException(status_code=500, detail=f"Unexpected error while getting batch status: {e}") @@ -529,13 +557,15 @@ async def cancel_queue_item( responses={200: {"model": SessionQueueCountsByDestination}}, ) async def counts_by_destination( + current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to query"), destination: str = Query(description="The destination to query"), ) -> SessionQueueCountsByDestination: - """Gets the counts of queue items by destination""" + """Gets the counts of queue items by destination. Non-admin users only see their own items.""" try: + user_id = None if current_user.is_admin else current_user.user_id return ApiDependencies.invoker.services.session_queue.get_counts_by_destination( - queue_id=queue_id, destination=destination + queue_id=queue_id, destination=destination, user_id=user_id ) except Exception as e: raise HTTPException(status_code=500, detail=f"Unexpected error while fetching counts by destination: {e}") diff --git a/invokeai/app/api/routers/workflows.py b/invokeai/app/api/routers/workflows.py index 7e34660a1df..1c88a77a3f6 100644 --- a/invokeai/app/api/routers/workflows.py +++ b/invokeai/app/api/routers/workflows.py @@ -74,7 +74,9 @@ async def update_workflow( raise HTTPException(status_code=404, detail="Workflow not found") if not current_user.is_admin and existing.user_id != current_user.user_id: raise HTTPException(status_code=403, detail="Not authorized to update this workflow") - return ApiDependencies.invoker.services.workflow_records.update(workflow=workflow) + # Pass user_id for defense-in-depth SQL scoping; admins pass None to allow any. + user_id = None if current_user.is_admin else current_user.user_id + return ApiDependencies.invoker.services.workflow_records.update(workflow=workflow, user_id=user_id) @workflows_router.delete( @@ -99,7 +101,8 @@ async def delete_workflow( except WorkflowThumbnailFileNotFoundException: # It's OK if the workflow has no thumbnail file. We can still delete the workflow. pass - ApiDependencies.invoker.services.workflow_records.delete(workflow_id) + user_id = None if current_user.is_admin else current_user.user_id + ApiDependencies.invoker.services.workflow_records.delete(workflow_id, user_id=user_id) @workflows_router.post( @@ -259,8 +262,12 @@ async def delete_workflow_thumbnail( async def get_workflow_thumbnail( workflow_id: str = Path(description="The id of the workflow thumbnail to get"), ) -> FileResponse: - """Gets a workflow's thumbnail image""" + """Gets a workflow's thumbnail image. + This endpoint is intentionally unauthenticated because browsers load images + via tags which cannot send Bearer tokens. Workflow IDs are UUIDs, + providing security through unguessability. + """ try: path = ApiDependencies.invoker.services.workflow_thumbnails.get_path(workflow_id) @@ -298,8 +305,9 @@ async def update_workflow_is_public( if config.multiuser and not current_user.is_admin and existing.user_id != current_user.user_id: raise HTTPException(status_code=403, detail="Not authorized to update this workflow") + user_id = None if current_user.is_admin else current_user.user_id return ApiDependencies.invoker.services.workflow_records.update_is_public( - workflow_id=workflow_id, is_public=is_public + workflow_id=workflow_id, is_public=is_public, user_id=user_id ) @@ -368,7 +376,18 @@ async def counts_by_category( operation_id="update_opened_at", ) async def update_opened_at( + current_user: CurrentUserOrDefault, workflow_id: str = Path(description="The workflow to update"), ) -> None: """Updates the opened_at field of a workflow""" - ApiDependencies.invoker.services.workflow_records.update_opened_at(workflow_id) + try: + existing = ApiDependencies.invoker.services.workflow_records.get(workflow_id) + except WorkflowNotFoundError: + raise HTTPException(status_code=404, detail="Workflow not found") + + config = ApiDependencies.invoker.services.configuration + if config.multiuser and not current_user.is_admin and existing.user_id != current_user.user_id: + raise HTTPException(status_code=403, detail="Not authorized to update this workflow") + + user_id = None if current_user.is_admin else current_user.user_id + ApiDependencies.invoker.services.workflow_records.update_opened_at(workflow_id, user_id=user_id) diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index fcead54eb1e..5783b804c0b 100644 --- a/invokeai/app/api/sockets.py +++ b/invokeai/app/api/sockets.py @@ -121,6 +121,11 @@ async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> b Returns True to accept the connection, False to reject it. Stores user_id in the internal socket users dict for later use. + + In multiuser mode, connections without a valid token are rejected outright + so that anonymous clients cannot subscribe to queue rooms and observe + queue activity belonging to other users. In single-user mode, unauthenticated + connections are accepted as the system admin user. """ # Extract token from auth data or headers token = None @@ -137,6 +142,23 @@ async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> b if token: token_data = verify_token(token) if token_data: + # In multiuser mode, also verify the backing user record still + # exists and is active — mirrors the REST auth check in + # auth_dependencies.py. A deleted or deactivated user whose + # JWT has not yet expired must not be allowed to open a socket. + if self._is_multiuser_enabled(): + try: + from invokeai.app.api.dependencies import ApiDependencies + + user = ApiDependencies.invoker.services.users.get(token_data.user_id) + if user is None or not user.is_active: + logger.warning(f"Rejecting socket {sid}: user {token_data.user_id} not found or inactive") + return False + except Exception: + # If user service is unavailable, fail closed + logger.warning(f"Rejecting socket {sid}: unable to verify user record") + return False + # Store user_id and is_admin in socket users dict self._socket_users[sid] = { "user_id": token_data.user_id, @@ -147,14 +169,37 @@ async def _handle_connect(self, sid: str, environ: dict, auth: dict | None) -> b ) return True - # If no valid token, store system user for backward compatibility + # No valid token provided. In multiuser mode this is not allowed — reject + # the connection so anonymous clients cannot subscribe to queue rooms. + # In single-user mode, fall through and accept the socket as system admin. + if self._is_multiuser_enabled(): + logger.warning( + f"Rejecting socket {sid} connection: multiuser mode is enabled and no valid auth token was provided" + ) + return False + self._socket_users[sid] = { "user_id": "system", - "is_admin": False, + "is_admin": True, } - logger.debug(f"Socket {sid} connected as system user (no valid token)") + logger.debug(f"Socket {sid} connected as system admin (single-user mode)") return True + @staticmethod + def _is_multiuser_enabled() -> bool: + """Check whether multiuser mode is enabled. Fails closed if configuration + is not yet initialized, which should not happen in practice but prevents + accidentally opening the socket during startup races.""" + try: + # Imported here to avoid a circular import at module load time. + from invokeai.app.api.dependencies import ApiDependencies + + return bool(ApiDependencies.invoker.services.configuration.multiuser) + except Exception: + # If dependencies are not initialized, fail closed (treat as multiuser) + # so we never accidentally admit an anonymous socket. + return True + async def _handle_disconnect(self, sid: str) -> None: """Handle socket disconnection and cleanup user info.""" if sid in self._socket_users: @@ -165,15 +210,20 @@ async def _handle_sub_queue(self, sid: str, data: Any) -> None: """Handle queue subscription and add socket to both queue and user-specific rooms.""" queue_id = QueueSubscriptionEvent(**data).queue_id - # Check if we have user info for this socket + # Check if we have user info for this socket. In multiuser mode _handle_connect + # will have already rejected any socket without a valid token, so missing user + # info here is a bug — refuse the subscription rather than silently falling back + # to an anonymous system user who could then receive queue item events. if sid not in self._socket_users: - logger.warning( - f"Socket {sid} subscribing to queue {queue_id} but has no user info - need to authenticate via connect event" - ) - # Store as system user temporarily - real auth should happen in connect + if self._is_multiuser_enabled(): + logger.warning( + f"Refusing queue subscription for socket {sid}: no user info (socket not authenticated via connect event)" + ) + return + # Single-user mode: safe to fall back to the system admin user. self._socket_users[sid] = { "user_id": "system", - "is_admin": False, + "is_admin": True, } user_id = self._socket_users[sid]["user_id"] @@ -198,6 +248,13 @@ async def _handle_unsub_queue(self, sid: str, data: Any) -> None: await self._sio.leave_room(sid, QueueSubscriptionEvent(**data).queue_id) async def _handle_sub_bulk_download(self, sid: str, data: Any) -> None: + # In multiuser mode, only allow authenticated sockets to subscribe. + # Bulk download events are routed to user-specific rooms, so the + # bulk_download_id room subscription is only kept for single-user + # backward compatibility. + if self._is_multiuser_enabled() and sid not in self._socket_users: + logger.warning(f"Refusing bulk download subscription for unknown socket {sid} in multiuser mode") + return await self._sio.enter_room(sid, BulkDownloadSubscriptionEvent(**data).bulk_download_id) async def _handle_unsub_bulk_download(self, sid: str, data: Any) -> None: @@ -206,9 +263,17 @@ async def _handle_unsub_bulk_download(self, sid: str, data: Any) -> None: async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]): """Handle queue events with user isolation. - Invocation events (progress, started, complete) are private - only emit to owner and admins. - Queue item status events are public - emit to all users (field values hidden via API). - Other queue events emit to all subscribers. + All queue item events (invocation events AND QueueItemStatusChangedEvent) are + private to the owning user and admins. They carry unsanitized user_id, batch_id, + session_id, origin, destination and error metadata, and must never be broadcast + to the whole queue room — otherwise any other authenticated subscriber could + observe cross-user queue activity. + + RecallParametersUpdatedEvent is also private to the owner + admins. + + BatchEnqueuedEvent carries the enqueuing user's batch_id/origin/counts and + is also routed privately. QueueClearedEvent is the only queue event that + is still broadcast to the whole queue room. IMPORTANT: Check InvocationEventBase BEFORE QueueItemEventBase since InvocationEventBase inherits from QueueItemEventBase. The order of isinstance checks matters! @@ -237,24 +302,40 @@ async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]): logger.debug(f"Emitted private invocation event {event_name} to user room {user_room} and admin room") - # Queue item status events are visible to all users (field values masked via API) - # This catches QueueItemStatusChangedEvent but NOT InvocationEvents (already handled above) + # Other queue item events (QueueItemStatusChangedEvent) carry unsanitized + # user_id, batch_id, session_id, origin, destination and error metadata. + # They are private to the owning user + admins — never broadcast to the + # full queue room. elif isinstance(event_data, QueueItemEventBase) and hasattr(event_data, "user_id"): - # Emit to all subscribers in the queue - await self._sio.emit( - event=event_name, data=event_data.model_dump(mode="json"), room=event_data.queue_id - ) + user_room = f"user:{event_data.user_id}" + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room) + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin") - logger.info( - f"Emitted public queue item event {event_name} to all subscribers in queue {event_data.queue_id}" - ) + logger.debug(f"Emitted private queue item event {event_name} to user room {user_room} and admin room") + + # RecallParametersUpdatedEvent is private - only emit to owner + admins + elif isinstance(event_data, RecallParametersUpdatedEvent): + user_room = f"user:{event_data.user_id}" + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room) + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin") + logger.debug(f"Emitted private recall_parameters_updated event to user room {user_room} and admin room") + + # BatchEnqueuedEvent carries the enqueuing user's batch_id, origin, and + # enqueued counts. Route it privately to the owner + admins so other + # users do not observe cross-user batch activity. + elif isinstance(event_data, BatchEnqueuedEvent): + user_room = f"user:{event_data.user_id}" + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room) + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin") + logger.debug(f"Emitted private batch_enqueued event to user room {user_room} and admin room") else: - # For other queue events (like QueueClearedEvent, BatchEnqueuedEvent), emit to all subscribers + # For remaining queue events (e.g. QueueClearedEvent) that do not + # carry user identity, emit to all subscribers in the queue room. await self._sio.emit( event=event_name, data=event_data.model_dump(mode="json"), room=event_data.queue_id ) - logger.info( + logger.debug( f"Emitted general queue event {event_name} to all subscribers in queue {event_data.queue_id}" ) except Exception as e: @@ -265,4 +346,17 @@ async def _handle_model_event(self, event: FastAPIEvent[ModelEventBase | Downloa await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json")) async def _handle_bulk_image_download_event(self, event: FastAPIEvent[BulkDownloadEventBase]) -> None: - await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json"), room=event[1].bulk_download_id) + event_name, event_data = event + # Route to user-specific + admin rooms so that other authenticated + # users cannot learn the bulk_download_item_name (the capability token + # needed to fetch the zip from the unauthenticated GET endpoint). + # In single-user mode (user_id="system"), fall back to the shared + # bulk_download_id room for backward compatibility. + if hasattr(event_data, "user_id") and event_data.user_id != "system": + user_room = f"user:{event_data.user_id}" + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room) + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin") + else: + await self._sio.emit( + event=event_name, data=event_data.model_dump(mode="json"), room=event_data.bulk_download_id + ) diff --git a/invokeai/app/api_app.py b/invokeai/app/api_app.py index 49894dba3cc..2ca6746b496 100644 --- a/invokeai/app/api_app.py +++ b/invokeai/app/api_app.py @@ -79,6 +79,50 @@ async def lifespan(app: FastAPI): ) +class SlidingWindowTokenMiddleware(BaseHTTPMiddleware): + """Refresh the JWT token on each authenticated response. + + When a request includes a valid Bearer token, the response includes a + X-Refreshed-Token header with a new token that has a fresh expiry. + This implements sliding-window session expiry: the session only expires + after a period of *inactivity*, not a fixed time after login. + """ + + async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): + response = await call_next(request) + + # Only refresh on mutating requests (POST/PUT/PATCH/DELETE) — these indicate + # genuine user activity. GET requests are often background fetches (RTK Query + # cache revalidation, refetch-on-focus, etc.) and should not reset the + # inactivity timer. + if response.status_code < 400 and request.method in ("POST", "PUT", "PATCH", "DELETE"): + auth_header = request.headers.get("authorization", "") + if auth_header.startswith("Bearer "): + token = auth_header[7:] + try: + from datetime import timedelta + + from invokeai.app.api.routers.auth import TOKEN_EXPIRATION_NORMAL, TOKEN_EXPIRATION_REMEMBER_ME + from invokeai.app.services.auth.token_service import create_access_token, verify_token + + token_data = verify_token(token) + if token_data is not None: + # Use the remember_me claim from the token to determine the + # correct refresh duration. This avoids the bug where a 7-day + # token with <24h remaining would be silently downgraded to 1 day. + if token_data.remember_me: + expires_delta = timedelta(days=TOKEN_EXPIRATION_REMEMBER_ME) + else: + expires_delta = timedelta(days=TOKEN_EXPIRATION_NORMAL) + + new_token = create_access_token(token_data, expires_delta) + response.headers["X-Refreshed-Token"] = new_token + except Exception: + pass # Don't fail the request if token refresh fails + + return response + + class RedirectRootWithQueryStringMiddleware(BaseHTTPMiddleware): """When a request is made to the root path with a query string, redirect to the root path without the query string. @@ -99,6 +143,7 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): # Add the middleware app.add_middleware(RedirectRootWithQueryStringMiddleware) +app.add_middleware(SlidingWindowTokenMiddleware) # Add event handler @@ -117,6 +162,7 @@ async def dispatch(self, request: Request, call_next: RequestResponseEndpoint): allow_credentials=app_config.allow_credentials, allow_methods=app_config.allow_methods, allow_headers=app_config.allow_headers, + expose_headers=["X-Refreshed-Token"], ) app.add_middleware(GZipMiddleware, minimum_size=1000) diff --git a/invokeai/app/invocations/anima_denoise.py b/invokeai/app/invocations/anima_denoise.py new file mode 100644 index 00000000000..530bf2918f1 --- /dev/null +++ b/invokeai/app/invocations/anima_denoise.py @@ -0,0 +1,715 @@ +"""Anima denoising invocation. + +Implements the rectified flow denoising loop for Anima models: +- Direct prediction: denoised = input - output * sigma +- Fixed shift=3.0 via loglinear_timestep_shift (Flux paper by Black Forest Labs) +- Timestep convention: timestep = sigma * 1.0 (raw sigma, NOT 1-sigma like Z-Image) +- NO v-prediction negation (unlike Z-Image) +- 3D latent space: [B, C, T, H, W] with T=1 for images +- 16 latent channels, 8x spatial compression + +Key differences from Z-Image denoise: +- Anima uses fixed shift=3.0, Z-Image uses dynamic shift based on resolution +- Anima: timestep = sigma (raw), Z-Image: model_t = 1.0 - sigma +- Anima: noise_pred = model_output (direct), Z-Image: noise_pred = -model_output (v-pred) +- Anima transformer takes (x, timesteps, context, t5xxl_ids, t5xxl_weights) +- Anima uses 3D latents directly, Z-Image converts 4D -> list of 5D +""" + +import inspect +import math +from contextlib import ExitStack +from typing import Callable, Iterator, Optional, Tuple + +import torch +import torchvision.transforms as tv_transforms +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from torchvision.transforms.functional import resize as tv_resize +from tqdm import tqdm + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + AnimaConditioningField, + DenoiseMaskField, + FieldDescriptions, + Input, + InputField, + LatentsField, +) +from invokeai.app.invocations.model import TransformerField +from invokeai.app.invocations.primitives import LatentsOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.anima.anima_transformer_patch import patch_anima_for_regional_prompting +from invokeai.backend.anima.conditioning_data import AnimaRegionalTextConditioning, AnimaTextConditioning +from invokeai.backend.anima.regional_prompting import AnimaRegionalPromptingExtension +from invokeai.backend.flux.schedulers import ANIMA_SCHEDULER_LABELS, ANIMA_SCHEDULER_MAP, ANIMA_SCHEDULER_NAME_VALUES +from invokeai.backend.model_manager.taxonomy import BaseModelType +from invokeai.backend.patches.layer_patcher import LayerPatcher +from invokeai.backend.patches.lora_conversions.anima_lora_constants import ANIMA_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import ( + RectifiedFlowInpaintExtension, + assert_broadcastable, +) +from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import AnimaConditioningInfo, Range +from invokeai.backend.util.devices import TorchDevice + +# Anima uses 8x spatial compression (VAE downsamples by 2^3) +ANIMA_LATENT_SCALE_FACTOR = 8 +# Anima uses 16 latent channels +ANIMA_LATENT_CHANNELS = 16 +# Anima uses fixed shift=3.0 for the rectified flow schedule +ANIMA_SHIFT = 3.0 +# Anima uses raw sigma values as timesteps (no rescaling) +ANIMA_MULTIPLIER = 1.0 + + +def loglinear_timestep_shift(alpha: float, t: float) -> float: + """Apply log-linear timestep shift to a noise schedule value. + + This shift biases the noise schedule toward higher noise levels, as described + in the Flux model (Black Forest Labs, 2024). With alpha > 1, the model spends + proportionally more denoising steps at higher noise levels. + + Formula: sigma = alpha * t / (1 + (alpha - 1) * t) + + Args: + alpha: Shift factor (3.0 for Anima, resolution-dependent for Flux). + t: Timestep value in [0, 1]. + + Returns: + Shifted timestep value. + """ + if alpha == 1.0: + return t + return alpha * t / (1 + (alpha - 1) * t) + + +def inverse_loglinear_timestep_shift(alpha: float, sigma: float) -> float: + """Recover linear t from a shifted sigma value. + + Inverse of loglinear_timestep_shift: given sigma = alpha * t / (1 + (alpha-1) * t), + solve for t = sigma / (alpha - (alpha-1) * sigma). + + This is needed for the inpainting extension, which expects linear t values + for gradient mask thresholding. With Anima's shift=3.0, the difference + between shifted sigma and linear t is large (e.g. at t=0.5, sigma=0.75), + causing overly aggressive mask thresholding if sigma is used directly. + + Args: + alpha: Shift factor (3.0 for Anima). + sigma: Shifted sigma value in [0, 1]. + + Returns: + Linear t value in [0, 1]. + """ + if alpha == 1.0: + return sigma + denominator = alpha - (alpha - 1) * sigma + if abs(denominator) < 1e-8: + return 1.0 + return sigma / denominator + + +class AnimaInpaintExtension(RectifiedFlowInpaintExtension): + """Inpaint extension for Anima that accounts for the time-SNR shift. + + Anima uses a fixed shift=3.0 which makes sigma values significantly larger + than the corresponding linear t values. The base RectifiedFlowInpaintExtension + uses t_prev for both gradient mask thresholding and noise mixing, which assumes + linear t values. + + This subclass: + - Uses the LINEAR t for gradient mask thresholding (correct progressive reveal) + - Uses the SHIFTED sigma for noise mixing (matches the denoiser's noise level) + """ + + def __init__( + self, + init_latents: torch.Tensor, + inpaint_mask: torch.Tensor, + noise: torch.Tensor, + shift: float = ANIMA_SHIFT, + ): + assert_broadcastable(init_latents.shape, inpaint_mask.shape, noise.shape) + self._init_latents = init_latents + self._inpaint_mask = inpaint_mask + self._noise = noise + self._shift = shift + + def merge_intermediate_latents_with_init_latents( + self, intermediate_latents: torch.Tensor, sigma_prev: float + ) -> torch.Tensor: + """Merge intermediate latents with init latents, correcting for Anima's shift. + + Args: + intermediate_latents: The denoised latents at the current step. + sigma_prev: The SHIFTED sigma value for the next step. + """ + # Recover linear t from shifted sigma for gradient mask thresholding. + # This ensures the gradient mask is revealed at the correct pace. + t_prev = inverse_loglinear_timestep_shift(self._shift, sigma_prev) + mask = self._apply_mask_gradient_adjustment(t_prev) + + # Use shifted sigma for noise mixing to match the denoiser's noise level. + # The Euler step produces latents at noise level sigma_prev, so the + # preserved regions must also be at sigma_prev noise level. + noised_init_latents = self._noise * sigma_prev + (1.0 - sigma_prev) * self._init_latents + + return intermediate_latents * mask + noised_init_latents * (1.0 - mask) + + +@invocation( + "anima_denoise", + title="Denoise - Anima", + tags=["image", "anima"], + category="image", + version="1.2.0", + classification=Classification.Prototype, +) +class AnimaDenoiseInvocation(BaseInvocation): + """Run the denoising process with an Anima model. + + Uses rectified flow sampling with shift=3.0 and the Cosmos Predict2 DiT + backbone with integrated LLM Adapter for text conditioning. + + Supports txt2img, img2img (via latents input), and inpainting (via denoise_mask). + """ + + # If latents is provided, this means we are doing image-to-image. + latents: Optional[LatentsField] = InputField( + default=None, description=FieldDescriptions.latents, input=Input.Connection + ) + # denoise_mask is used for inpainting. Only the masked region is modified. + denoise_mask: Optional[DenoiseMaskField] = InputField( + default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection + ) + denoising_start: float = InputField(default=0.0, ge=0, le=1, description=FieldDescriptions.denoising_start) + denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end) + add_noise: bool = InputField(default=True, description="Add noise based on denoising start.") + transformer: TransformerField = InputField( + description="Anima transformer model.", input=Input.Connection, title="Transformer" + ) + positive_conditioning: AnimaConditioningField | list[AnimaConditioningField] = InputField( + description=FieldDescriptions.positive_cond, input=Input.Connection + ) + negative_conditioning: AnimaConditioningField | list[AnimaConditioningField] | None = InputField( + default=None, description=FieldDescriptions.negative_cond, input=Input.Connection + ) + guidance_scale: float = InputField( + default=4.5, + ge=1.0, + description="Guidance scale for classifier-free guidance. Recommended: 4.0-5.0 for Anima.", + title="Guidance Scale", + ) + width: int = InputField(default=1024, multiple_of=8, description="Width of the generated image.") + height: int = InputField(default=1024, multiple_of=8, description="Height of the generated image.") + steps: int = InputField(default=30, gt=0, description="Number of denoising steps. 30 recommended for Anima.") + seed: int = InputField(default=0, description="Randomness seed for reproducibility.") + scheduler: ANIMA_SCHEDULER_NAME_VALUES = InputField( + default="euler", + description="Scheduler (sampler) for the denoising process.", + ui_choice_labels=ANIMA_SCHEDULER_LABELS, + ) + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> LatentsOutput: + latents = self._run_diffusion(context) + latents = latents.detach().to("cpu") + name = context.tensors.save(tensor=latents) + return LatentsOutput.build(latents_name=name, latents=latents, seed=None) + + def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> torch.Tensor | None: + """Prepare the inpaint mask for Anima. + + Anima uses 3D latents [B, C, T, H, W] internally but the mask operates + on the spatial dimensions [B, C, H, W] which match the squeezed output. + """ + if self.denoise_mask is None: + return None + mask = context.tensors.load(self.denoise_mask.mask_name) + + # Invert mask: 0.0 = regions to denoise, 1.0 = regions to preserve + mask = 1.0 - mask + + _, _, latent_height, latent_width = latents.shape + mask = tv_resize( + img=mask, + size=[latent_height, latent_width], + interpolation=tv_transforms.InterpolationMode.BILINEAR, + antialias=False, + ) + + mask = mask.to(device=latents.device, dtype=latents.dtype) + return mask + + def _get_noise( + self, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + seed: int, + ) -> torch.Tensor: + """Generate initial noise tensor in 3D latent space [B, C, T, H, W].""" + rand_device = "cpu" + return torch.randn( + 1, + ANIMA_LATENT_CHANNELS, + 1, # T=1 for single image + height // ANIMA_LATENT_SCALE_FACTOR, + width // ANIMA_LATENT_SCALE_FACTOR, + device=rand_device, + dtype=torch.float32, + generator=torch.Generator(device=rand_device).manual_seed(seed), + ).to(device=device, dtype=dtype) + + def _get_sigmas(self, num_steps: int) -> list[float]: + """Generate sigma schedule with fixed shift=3.0. + + Uses the log-linear timestep shift from the Flux model (Black Forest Labs) + with a fixed shift factor of 3.0 (no dynamic resolution-based shift). + + Returns: + List of num_steps + 1 sigma values from ~1.0 (noise) to 0.0 (clean). + """ + sigmas = [] + for i in range(num_steps + 1): + t = 1.0 - i / num_steps + sigma = loglinear_timestep_shift(ANIMA_SHIFT, t) + sigmas.append(sigma) + return sigmas + + def _load_conditioning( + self, + context: InvocationContext, + cond_field: AnimaConditioningField, + dtype: torch.dtype, + device: torch.device, + ) -> AnimaConditioningInfo: + """Load Anima conditioning data from storage.""" + cond_data = context.conditioning.load(cond_field.conditioning_name) + assert len(cond_data.conditionings) == 1 + cond_info = cond_data.conditionings[0] + assert isinstance(cond_info, AnimaConditioningInfo) + return cond_info.to(dtype=dtype, device=device) + + def _load_text_conditionings( + self, + context: InvocationContext, + cond_field: AnimaConditioningField | list[AnimaConditioningField], + img_token_height: int, + img_token_width: int, + dtype: torch.dtype, + device: torch.device, + ) -> list[AnimaTextConditioning]: + """Load Anima text conditioning with optional regional masks. + + Args: + context: The invocation context. + cond_field: Single conditioning field or list of fields. + img_token_height: Height of the image token grid (H // patch_size). + img_token_width: Width of the image token grid (W // patch_size). + dtype: Target dtype. + device: Target device. + + Returns: + List of AnimaTextConditioning objects with optional masks. + """ + cond_list = cond_field if isinstance(cond_field, list) else [cond_field] + + text_conditionings: list[AnimaTextConditioning] = [] + for cond in cond_list: + cond_info = self._load_conditioning(context, cond, dtype, device) + + # Load the mask, if provided + mask: torch.Tensor | None = None + if cond.mask is not None: + mask = context.tensors.load(cond.mask.tensor_name) + mask = mask.to(device=device) + mask = AnimaRegionalPromptingExtension.preprocess_regional_prompt_mask( + mask, img_token_height, img_token_width, dtype, device + ) + + text_conditionings.append( + AnimaTextConditioning( + qwen3_embeds=cond_info.qwen3_embeds, + t5xxl_ids=cond_info.t5xxl_ids, + t5xxl_weights=cond_info.t5xxl_weights, + mask=mask, + ) + ) + + return text_conditionings + + def _run_llm_adapter_for_regions( + self, + transformer, + text_conditionings: list[AnimaTextConditioning], + dtype: torch.dtype, + ) -> AnimaRegionalTextConditioning: + """Run the LLM Adapter separately for each regional conditioning and concatenate. + + Args: + transformer: The AnimaTransformer instance (must be on device). + text_conditionings: List of per-region conditioning data. + dtype: Inference dtype. + + Returns: + AnimaRegionalTextConditioning with concatenated context and masks. + """ + context_embeds_list: list[torch.Tensor] = [] + context_ranges: list[Range] = [] + image_masks: list[torch.Tensor | None] = [] + cur_len = 0 + + for tc in text_conditionings: + qwen3_embeds = tc.qwen3_embeds.unsqueeze(0) # (1, seq_len, 1024) + t5xxl_ids = tc.t5xxl_ids.unsqueeze(0) # (1, seq_len) + t5xxl_weights = None + if tc.t5xxl_weights is not None: + t5xxl_weights = tc.t5xxl_weights.unsqueeze(0).unsqueeze(-1) # (1, seq_len, 1) + + # Run the LLM Adapter to produce context for this region + context = transformer.preprocess_text_embeds( + qwen3_embeds.to(dtype=dtype), + t5xxl_ids, + t5xxl_weights=t5xxl_weights.to(dtype=dtype) if t5xxl_weights is not None else None, + ) + # context shape: (1, 512, 1024) — squeeze batch dim + context_2d = context.squeeze(0) # (512, 1024) + + context_embeds_list.append(context_2d) + context_ranges.append(Range(start=cur_len, end=cur_len + context_2d.shape[0])) + image_masks.append(tc.mask) + cur_len += context_2d.shape[0] + + concatenated_context = torch.cat(context_embeds_list, dim=0) + + return AnimaRegionalTextConditioning( + context_embeds=concatenated_context, + image_masks=image_masks, + context_ranges=context_ranges, + ) + + def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: + device = TorchDevice.choose_torch_device() + inference_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) + + if self.denoising_start >= self.denoising_end: + raise ValueError( + f"denoising_start ({self.denoising_start}) must be less than denoising_end ({self.denoising_end})." + ) + + transformer_info = context.models.load(self.transformer.transformer) + + # Compute image token grid dimensions for regional prompting + # Anima: 8x VAE compression, 2x patch size → 16x total + patch_size = 2 + latent_height = self.height // ANIMA_LATENT_SCALE_FACTOR + latent_width = self.width // ANIMA_LATENT_SCALE_FACTOR + img_token_height = latent_height // patch_size + img_token_width = latent_width // patch_size + img_seq_len = img_token_height * img_token_width + + # Load positive conditioning with optional regional masks + pos_text_conditionings = self._load_text_conditionings( + context=context, + cond_field=self.positive_conditioning, + img_token_height=img_token_height, + img_token_width=img_token_width, + dtype=inference_dtype, + device=device, + ) + has_regional = len(pos_text_conditionings) > 1 or any(tc.mask is not None for tc in pos_text_conditionings) + + # Load negative conditioning if CFG is enabled + do_cfg = not math.isclose(self.guidance_scale, 1.0) and self.negative_conditioning is not None + neg_text_conditionings: list[AnimaTextConditioning] | None = None + if do_cfg: + assert self.negative_conditioning is not None + neg_text_conditionings = self._load_text_conditionings( + context=context, + cond_field=self.negative_conditioning, + img_token_height=img_token_height, + img_token_width=img_token_width, + dtype=inference_dtype, + device=device, + ) + + # Generate sigma schedule + sigmas = self._get_sigmas(self.steps) + + # Apply denoising_start and denoising_end clipping (for img2img/inpaint) + if self.denoising_start > 0 or self.denoising_end < 1: + total_sigmas = len(sigmas) + start_idx = int(self.denoising_start * (total_sigmas - 1)) + end_idx = int(self.denoising_end * (total_sigmas - 1)) + 1 + sigmas = sigmas[start_idx:end_idx] + + total_steps = len(sigmas) - 1 + + # Load input latents if provided (image-to-image) + init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None + if init_latents is not None: + init_latents = init_latents.to(device=device, dtype=inference_dtype) + # Anima denoiser works in 3D: add temporal dim if needed + if init_latents.ndim == 4: + init_latents = init_latents.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + + # Generate initial noise (3D latent: [B, C, T, H, W]) + noise = self._get_noise(self.height, self.width, inference_dtype, device, self.seed) + + # Prepare input latents + if init_latents is not None: + if self.add_noise: + s_0 = sigmas[0] + latents = s_0 * noise + (1.0 - s_0) * init_latents + else: + latents = init_latents + else: + if self.denoising_start > 1e-5: + raise ValueError("denoising_start should be 0 when initial latents are not provided.") + latents = noise + + if total_steps <= 0: + return latents.squeeze(2) + + # Prepare inpaint extension + inpaint_mask = self._prep_inpaint_mask(context, latents.squeeze(2)) + inpaint_extension: AnimaInpaintExtension | None = None + if inpaint_mask is not None: + if init_latents is None: + raise ValueError("Initial latents are required when using an inpaint mask (image-to-image inpainting)") + inpaint_extension = AnimaInpaintExtension( + init_latents=init_latents.squeeze(2), + inpaint_mask=inpaint_mask, + noise=noise.squeeze(2), + shift=ANIMA_SHIFT, + ) + + step_callback = self._build_step_callback(context) + + # Initialize diffusers scheduler if not using built-in Euler + scheduler: SchedulerMixin | None = None + use_scheduler = self.scheduler != "euler" + + if use_scheduler: + scheduler_class = ANIMA_SCHEDULER_MAP[self.scheduler] + scheduler = scheduler_class(num_train_timesteps=1000, shift=1.0) + is_lcm = self.scheduler == "lcm" + set_timesteps_sig = inspect.signature(scheduler.set_timesteps) + if not is_lcm and "sigmas" in set_timesteps_sig.parameters: + scheduler.set_timesteps(sigmas=sigmas, device=device) + else: + scheduler.set_timesteps(num_inference_steps=total_steps, device=device) + num_scheduler_steps = len(scheduler.timesteps) + else: + num_scheduler_steps = total_steps + + with ExitStack() as exit_stack: + (cached_weights, transformer) = exit_stack.enter_context(transformer_info.model_on_device()) + + # Apply LoRA models to the transformer. + # Note: We apply the LoRA after the transformer has been moved to its target device for faster patching. + exit_stack.enter_context( + LayerPatcher.apply_smart_model_patches( + model=transformer, + patches=self._lora_iterator(context), + prefix=ANIMA_LORA_TRANSFORMER_PREFIX, + dtype=inference_dtype, + cached_weights=cached_weights, + ) + ) + + # Run LLM Adapter for each regional conditioning to produce context vectors. + # This must happen with the transformer on device since it uses the adapter weights. + if has_regional: + pos_regional = self._run_llm_adapter_for_regions(transformer, pos_text_conditionings, inference_dtype) + pos_context = pos_regional.context_embeds.unsqueeze(0) # (1, total_ctx_len, 1024) + + # Build regional prompting extension with cross-attention mask + regional_extension = AnimaRegionalPromptingExtension.from_regional_conditioning( + pos_regional, img_seq_len + ) + + # For negative, concatenate all regions without masking (matches Z-Image behavior) + neg_context = None + if do_cfg and neg_text_conditionings is not None: + neg_regional = self._run_llm_adapter_for_regions( + transformer, neg_text_conditionings, inference_dtype + ) + neg_context = neg_regional.context_embeds.unsqueeze(0) + else: + # Single conditioning — run LLM Adapter via normal forward path + tc = pos_text_conditionings[0] + pos_qwen3_embeds = tc.qwen3_embeds.unsqueeze(0) + pos_t5xxl_ids = tc.t5xxl_ids.unsqueeze(0) + pos_t5xxl_weights = None + if tc.t5xxl_weights is not None: + pos_t5xxl_weights = tc.t5xxl_weights.unsqueeze(0).unsqueeze(-1) + + # Pre-compute context via LLM Adapter + pos_context = transformer.preprocess_text_embeds( + pos_qwen3_embeds.to(dtype=inference_dtype), + pos_t5xxl_ids, + t5xxl_weights=pos_t5xxl_weights.to(dtype=inference_dtype) + if pos_t5xxl_weights is not None + else None, + ) + + neg_context = None + if do_cfg and neg_text_conditionings is not None: + ntc = neg_text_conditionings[0] + neg_qwen3 = ntc.qwen3_embeds.unsqueeze(0) + neg_ids = ntc.t5xxl_ids.unsqueeze(0) + neg_weights = None + if ntc.t5xxl_weights is not None: + neg_weights = ntc.t5xxl_weights.unsqueeze(0).unsqueeze(-1) + neg_context = transformer.preprocess_text_embeds( + neg_qwen3.to(dtype=inference_dtype), + neg_ids, + t5xxl_weights=neg_weights.to(dtype=inference_dtype) if neg_weights is not None else None, + ) + + regional_extension = None + + # Apply regional prompting patch if we have regional masks + exit_stack.enter_context(patch_anima_for_regional_prompting(transformer, regional_extension)) + + # Helper to run transformer with pre-computed context (bypasses LLM Adapter) + def _run_transformer(ctx: torch.Tensor, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + return transformer( + x=x.to(transformer.dtype if hasattr(transformer, "dtype") else inference_dtype), + timesteps=t, + context=ctx, + # t5xxl_ids=None skips the LLM Adapter — context is already pre-computed + ) + + if use_scheduler and scheduler is not None: + # Scheduler-based denoising + user_step = 0 + pbar = tqdm(total=total_steps, desc="Denoising (Anima)") + for step_index in range(num_scheduler_steps): + sched_timestep = scheduler.timesteps[step_index] + sigma_curr = sched_timestep.item() / scheduler.config.num_train_timesteps + + is_heun = hasattr(scheduler, "state_in_first_order") + in_first_order = scheduler.state_in_first_order if is_heun else True + + timestep = torch.tensor( + [sigma_curr * ANIMA_MULTIPLIER], device=device, dtype=inference_dtype + ).expand(latents.shape[0]) + + noise_pred_cond = _run_transformer(pos_context, latents, timestep).float() + + if do_cfg and neg_context is not None: + noise_pred_uncond = _run_transformer(neg_context, latents, timestep).float() + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond + + step_output = scheduler.step(model_output=noise_pred, timestep=sched_timestep, sample=latents) + latents = step_output.prev_sample + + if step_index + 1 < len(scheduler.sigmas): + sigma_prev = scheduler.sigmas[step_index + 1].item() + else: + sigma_prev = 0.0 + + if inpaint_extension is not None: + latents_4d = latents.squeeze(2) + latents_4d = inpaint_extension.merge_intermediate_latents_with_init_latents( + latents_4d, sigma_prev + ) + latents = latents_4d.unsqueeze(2) + + if is_heun: + if not in_first_order: + user_step += 1 + if user_step <= total_steps: + pbar.update(1) + step_callback( + PipelineIntermediateState( + step=user_step, + order=2, + total_steps=total_steps, + timestep=int(sigma_curr * 1000), + latents=latents.squeeze(2), + ) + ) + else: + user_step += 1 + if user_step <= total_steps: + pbar.update(1) + step_callback( + PipelineIntermediateState( + step=user_step, + order=1, + total_steps=total_steps, + timestep=int(sigma_curr * 1000), + latents=latents.squeeze(2), + ) + ) + pbar.close() + else: + # Built-in Euler implementation (default for Anima) + for step_idx in tqdm(range(total_steps), desc="Denoising (Anima)"): + sigma_curr = sigmas[step_idx] + sigma_prev = sigmas[step_idx + 1] + + timestep = torch.tensor( + [sigma_curr * ANIMA_MULTIPLIER], device=device, dtype=inference_dtype + ).expand(latents.shape[0]) + + noise_pred_cond = _run_transformer(pos_context, latents, timestep).float() + + if do_cfg and neg_context is not None: + noise_pred_uncond = _run_transformer(neg_context, latents, timestep).float() + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond + + latents_dtype = latents.dtype + latents = latents.to(dtype=torch.float32) + latents = latents + (sigma_prev - sigma_curr) * noise_pred + latents = latents.to(dtype=latents_dtype) + + if inpaint_extension is not None: + latents_4d = latents.squeeze(2) + latents_4d = inpaint_extension.merge_intermediate_latents_with_init_latents( + latents_4d, sigma_prev + ) + latents = latents_4d.unsqueeze(2) + + step_callback( + PipelineIntermediateState( + step=step_idx + 1, + order=1, + total_steps=total_steps, + timestep=int(sigma_curr * 1000), + latents=latents.squeeze(2), + ), + ) + + # Remove temporal dimension for output: [B, C, 1, H, W] -> [B, C, H, W] + return latents.squeeze(2) + + def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]: + def step_callback(state: PipelineIntermediateState) -> None: + context.util.sd_step_callback(state, BaseModelType.Anima) + + return step_callback + + def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]: + """Iterate over LoRA models to apply to the transformer.""" + for lora in self.transformer.loras: + lora_info = context.models.load(lora.lora) + if not isinstance(lora_info.model, ModelPatchRaw): + raise TypeError( + f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}. " + "The LoRA model may be corrupted or incompatible." + ) + yield (lora_info.model, lora.weight) + del lora_info diff --git a/invokeai/app/invocations/anima_image_to_latents.py b/invokeai/app/invocations/anima_image_to_latents.py new file mode 100644 index 00000000000..83073ab4a80 --- /dev/null +++ b/invokeai/app/invocations/anima_image_to_latents.py @@ -0,0 +1,119 @@ +"""Anima image-to-latents invocation. + +Encodes an image to latent space using the Anima VAE (AutoencoderKLWan or FLUX VAE). + +For Wan VAE (AutoencoderKLWan): +- Input image is converted to 5D tensor [B, C, T, H, W] with T=1 +- After encoding, latents are normalized: (latents - mean) / std + (inverse of the denormalization in anima_latents_to_image.py) + +For FLUX VAE (AutoEncoder): +- Encoding is handled internally by the FLUX VAE +""" + +from typing import Union + +import einops +import torch +from diffusers.models.autoencoders import AutoencoderKLWan + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + FieldDescriptions, + ImageField, + Input, + InputField, + WithBoard, + WithMetadata, +) +from invokeai.app.invocations.model import VAEField +from invokeai.app.invocations.primitives import LatentsOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder +from invokeai.backend.model_manager.load.load_base import LoadedModel +from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux + +AnimaVAE = Union[AutoencoderKLWan, FluxAutoEncoder] + + +@invocation( + "anima_i2l", + title="Image to Latents - Anima", + tags=["image", "latents", "vae", "i2l", "anima"], + category="image", + version="1.0.0", + classification=Classification.Prototype, +) +class AnimaImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard): + """Generates latents from an image using the Anima VAE (supports Wan 2.1 and FLUX VAE).""" + + image: ImageField = InputField(description="The image to encode.") + vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection) + + @staticmethod + def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor: + if not isinstance(vae_info.model, (AutoencoderKLWan, FluxAutoEncoder)): + raise TypeError( + f"Expected AutoencoderKLWan or FluxAutoEncoder for Anima VAE, got {type(vae_info.model).__name__}." + ) + + estimated_working_memory = estimate_vae_working_memory_flux( + operation="encode", + image_tensor=image_tensor, + vae=vae_info.model, + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + if not isinstance(vae, (AutoencoderKLWan, FluxAutoEncoder)): + raise TypeError(f"Expected AutoencoderKLWan or FluxAutoEncoder, got {type(vae).__name__}.") + + vae_dtype = next(iter(vae.parameters())).dtype + image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + + with torch.inference_mode(): + if isinstance(vae, FluxAutoEncoder): + # FLUX VAE handles scaling internally + generator = torch.Generator(device=TorchDevice.choose_torch_device()).manual_seed(0) + latents = vae.encode(image_tensor, sample=True, generator=generator) + else: + # AutoencoderKLWan expects 5D input [B, C, T, H, W] + if image_tensor.ndim == 4: + image_tensor = image_tensor.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + + encoded = vae.encode(image_tensor, return_dict=False)[0] + latents = encoded.sample().to(dtype=vae_dtype) + + # Normalize to denoiser space: (latents - mean) / std + # This is the inverse of the denormalization in anima_latents_to_image.py + latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents) + latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents) + latents = (latents - latents_mean) / latents_std + + # Remove temporal dimension: [B, C, 1, H, W] -> [B, C, H, W] + if latents.ndim == 5: + latents = latents.squeeze(2) + + return latents + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> LatentsOutput: + image = context.images.get_pil(self.image.image_name) + + image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB")) + if image_tensor.dim() == 3: + image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w") + + vae_info = context.models.load(self.vae.vae) + if not isinstance(vae_info.model, (AutoencoderKLWan, FluxAutoEncoder)): + raise TypeError( + f"Expected AutoencoderKLWan or FluxAutoEncoder for Anima VAE, got {type(vae_info.model).__name__}." + ) + + context.util.signal_progress("Running Anima VAE encode") + latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor) + + latents = latents.to("cpu") + name = context.tensors.save(tensor=latents) + return LatentsOutput.build(latents_name=name, latents=latents, seed=None) diff --git a/invokeai/app/invocations/anima_latents_to_image.py b/invokeai/app/invocations/anima_latents_to_image.py new file mode 100644 index 00000000000..080c101fa44 --- /dev/null +++ b/invokeai/app/invocations/anima_latents_to_image.py @@ -0,0 +1,108 @@ +"""Anima latents-to-image invocation. + +Decodes Anima latents using the QwenImage VAE (AutoencoderKLWan) or +compatible FLUX VAE as fallback. + +Latents from the denoiser are in normalized space (zero-centered). Before +VAE decode, they must be denormalized using the Wan 2.1 per-channel +mean/std: latents = latents * std + mean (matching diffusers WanPipeline). + +The VAE expects 5D latents [B, C, T, H, W] — for single images, T=1. +""" + +import torch +from diffusers.models.autoencoders import AutoencoderKLWan +from einops import rearrange +from PIL import Image + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + FieldDescriptions, + Input, + InputField, + LatentsField, + WithBoard, + WithMetadata, +) +from invokeai.app.invocations.model import VAEField +from invokeai.app.invocations.primitives import ImageOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux + + +@invocation( + "anima_l2i", + title="Latents to Image - Anima", + tags=["latents", "image", "vae", "l2i", "anima"], + category="latents", + version="1.0.2", + classification=Classification.Prototype, +) +class AnimaLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard): + """Generates an image from latents using the Anima VAE. + + Supports the Wan 2.1 QwenImage VAE (AutoencoderKLWan) with explicit + latent denormalization, and FLUX VAE as fallback. + """ + + latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection) + vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection) + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> ImageOutput: + latents = context.tensors.load(self.latents.latents_name) + + vae_info = context.models.load(self.vae.vae) + if not isinstance(vae_info.model, (AutoencoderKLWan, FluxAutoEncoder)): + raise TypeError( + f"Expected AutoencoderKLWan or FluxAutoEncoder for Anima VAE, got {type(vae_info.model).__name__}." + ) + + estimated_working_memory = estimate_vae_working_memory_flux( + operation="decode", + image_tensor=latents, + vae=vae_info.model, + ) + + with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae): + context.util.signal_progress("Running Anima VAE decode") + if not isinstance(vae, (AutoencoderKLWan, FluxAutoEncoder)): + raise TypeError(f"Expected AutoencoderKLWan or FluxAutoEncoder, got {type(vae).__name__}.") + + vae_dtype = next(iter(vae.parameters())).dtype + latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype) + + TorchDevice.empty_cache() + + with torch.inference_mode(): + if isinstance(vae, FluxAutoEncoder): + # FLUX VAE handles scaling internally, expects 4D [B, C, H, W] + img = vae.decode(latents) + else: + # Expects 5D latents [B, C, T, H, W] + if latents.ndim == 4: + latents = latents.unsqueeze(2) # [B, C, H, W] -> [B, C, 1, H, W] + + # Denormalize from denoiser space to raw VAE space + # (same as diffusers WanPipeline and ComfyUI Wan21.process_out) + latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latents) + latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1).to(latents) + latents = latents * latents_std + latents_mean + + decoded = vae.decode(latents, return_dict=False)[0] + + # Output is 5D [B, C, T, H, W] — squeeze temporal dim + if decoded.ndim == 5: + decoded = decoded.squeeze(2) + img = decoded + + img = img.clamp(-1, 1) + img = rearrange(img[0], "c h w -> h w c") + img_pil = Image.fromarray((127.5 * (img + 1.0)).byte().cpu().numpy()) + + TorchDevice.empty_cache() + + image_dto = context.images.save(image=img_pil) + return ImageOutput.build(image_dto) diff --git a/invokeai/app/invocations/anima_lora_loader.py b/invokeai/app/invocations/anima_lora_loader.py new file mode 100644 index 00000000000..6a035b55aa6 --- /dev/null +++ b/invokeai/app/invocations/anima_lora_loader.py @@ -0,0 +1,162 @@ +from typing import Optional + +from invokeai.app.invocations.baseinvocation import ( + BaseInvocation, + BaseInvocationOutput, + Classification, + invocation, + invocation_output, +) +from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField +from invokeai.app.invocations.model import LoRAField, ModelIdentifierField, Qwen3EncoderField, TransformerField +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType + + +@invocation_output("anima_lora_loader_output") +class AnimaLoRALoaderOutput(BaseInvocationOutput): + """Anima LoRA Loader Output""" + + transformer: Optional[TransformerField] = OutputField( + default=None, description=FieldDescriptions.transformer, title="Anima Transformer" + ) + qwen3_encoder: Optional[Qwen3EncoderField] = OutputField( + default=None, description=FieldDescriptions.qwen3_encoder, title="Qwen3 Encoder" + ) + + +@invocation( + "anima_lora_loader", + title="Apply LoRA - Anima", + tags=["lora", "model", "anima"], + category="model", + version="1.0.0", + classification=Classification.Prototype, +) +class AnimaLoRALoaderInvocation(BaseInvocation): + """Apply a LoRA model to an Anima transformer and/or Qwen3 text encoder.""" + + lora: ModelIdentifierField = InputField( + description=FieldDescriptions.lora_model, + title="LoRA", + ui_model_base=BaseModelType.Anima, + ui_model_type=ModelType.LoRA, + ) + weight: float = InputField(default=0.75, description=FieldDescriptions.lora_weight) + transformer: TransformerField | None = InputField( + default=None, + description=FieldDescriptions.transformer, + input=Input.Connection, + title="Anima Transformer", + ) + qwen3_encoder: Qwen3EncoderField | None = InputField( + default=None, + title="Qwen3 Encoder", + description=FieldDescriptions.qwen3_encoder, + input=Input.Connection, + ) + + def invoke(self, context: InvocationContext) -> AnimaLoRALoaderOutput: + lora_key = self.lora.key + + if not context.models.exists(lora_key): + raise ValueError(f"Unknown lora: {lora_key}!") + + if self.transformer and any(lora.lora.key == lora_key for lora in self.transformer.loras): + raise ValueError(f'LoRA "{lora_key}" already applied to transformer.') + if self.qwen3_encoder and any(lora.lora.key == lora_key for lora in self.qwen3_encoder.loras): + raise ValueError(f'LoRA "{lora_key}" already applied to Qwen3 encoder.') + + output = AnimaLoRALoaderOutput() + + if self.transformer is not None: + output.transformer = self.transformer.model_copy(deep=True) + output.transformer.loras.append( + LoRAField( + lora=self.lora, + weight=self.weight, + ) + ) + if self.qwen3_encoder is not None: + output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True) + output.qwen3_encoder.loras.append( + LoRAField( + lora=self.lora, + weight=self.weight, + ) + ) + + return output + + +@invocation( + "anima_lora_collection_loader", + title="Apply LoRA Collection - Anima", + tags=["lora", "model", "anima"], + category="model", + version="1.0.0", + classification=Classification.Prototype, +) +class AnimaLoRACollectionLoader(BaseInvocation): + """Applies a collection of LoRAs to an Anima transformer.""" + + loras: Optional[LoRAField | list[LoRAField]] = InputField( + default=None, description="LoRA models and weights. May be a single LoRA or collection.", title="LoRAs" + ) + + transformer: Optional[TransformerField] = InputField( + default=None, + description=FieldDescriptions.transformer, + input=Input.Connection, + title="Transformer", + ) + qwen3_encoder: Qwen3EncoderField | None = InputField( + default=None, + title="Qwen3 Encoder", + description=FieldDescriptions.qwen3_encoder, + input=Input.Connection, + ) + + def invoke(self, context: InvocationContext) -> AnimaLoRALoaderOutput: + output = AnimaLoRALoaderOutput() + + if self.loras is None: + if self.transformer is not None: + output.transformer = self.transformer.model_copy(deep=True) + if self.qwen3_encoder is not None: + output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True) + return output + + loras = self.loras if isinstance(self.loras, list) else [self.loras] + added_loras: list[str] = [] + + if self.transformer is not None: + output.transformer = self.transformer.model_copy(deep=True) + + if self.qwen3_encoder is not None: + output.qwen3_encoder = self.qwen3_encoder.model_copy(deep=True) + + for lora in loras: + if lora is None: + continue + if lora.lora.key in added_loras: + continue + + if not context.models.exists(lora.lora.key): + raise ValueError(f"Unknown lora: {lora.lora.key}!") + + if lora.lora.base is not BaseModelType.Anima: + raise ValueError( + f"LoRA '{lora.lora.key}' is for {lora.lora.base.value if lora.lora.base else 'unknown'} models, " + "not Anima models. Ensure you are using an Anima compatible LoRA." + ) + + added_loras.append(lora.lora.key) + + if self.transformer is not None and output.transformer is not None: + output.transformer.loras.append(lora) + + if self.qwen3_encoder is not None and output.qwen3_encoder is not None: + output.qwen3_encoder.loras.append(lora) + + return output diff --git a/invokeai/app/invocations/anima_model_loader.py b/invokeai/app/invocations/anima_model_loader.py new file mode 100644 index 00000000000..7051148cb12 --- /dev/null +++ b/invokeai/app/invocations/anima_model_loader.py @@ -0,0 +1,102 @@ +from invokeai.app.invocations.baseinvocation import ( + BaseInvocation, + BaseInvocationOutput, + Classification, + invocation, + invocation_output, +) +from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField +from invokeai.app.invocations.model import ( + ModelIdentifierField, + Qwen3EncoderField, + T5EncoderField, + TransformerField, + VAEField, +) +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.app.util.t5_model_identifier import ( + preprocess_t5_encoder_model_identifier, + preprocess_t5_tokenizer_model_identifier, +) +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType, SubModelType + + +@invocation_output("anima_model_loader_output") +class AnimaModelLoaderOutput(BaseInvocationOutput): + """Anima model loader output.""" + + transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer") + qwen3_encoder: Qwen3EncoderField = OutputField(description=FieldDescriptions.qwen3_encoder, title="Qwen3 Encoder") + vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE") + t5_encoder: T5EncoderField = OutputField(description=FieldDescriptions.t5_encoder, title="T5 Encoder") + + +@invocation( + "anima_model_loader", + title="Main Model - Anima", + tags=["model", "anima"], + category="model", + version="1.3.0", + classification=Classification.Prototype, +) +class AnimaModelLoaderInvocation(BaseInvocation): + """Loads an Anima model, outputting its submodels. + + Anima uses: + - Transformer: Cosmos Predict2 DiT + LLM Adapter (from single-file checkpoint) + - Qwen3 Encoder: Qwen3 0.6B (standalone single-file) + - VAE: AutoencoderKLQwenImage / Wan 2.1 VAE (standalone single-file or FLUX VAE) + - T5 Encoder: T5-XXL model (only the tokenizer submodel is used, for LLM Adapter token IDs) + """ + + model: ModelIdentifierField = InputField( + description="Anima main model (transformer + LLM adapter).", + input=Input.Direct, + ui_model_base=BaseModelType.Anima, + ui_model_type=ModelType.Main, + title="Transformer", + ) + + vae_model: ModelIdentifierField = InputField( + description="Standalone VAE model. Anima uses a Wan 2.1 / QwenImage VAE (16-channel). " + "A FLUX VAE can also be used as a compatible fallback.", + input=Input.Direct, + ui_model_type=ModelType.VAE, + title="VAE", + ) + + qwen3_encoder_model: ModelIdentifierField = InputField( + description="Standalone Qwen3 0.6B Encoder model.", + input=Input.Direct, + ui_model_type=ModelType.Qwen3Encoder, + title="Qwen3 Encoder", + ) + + t5_encoder_model: ModelIdentifierField = InputField( + description="T5-XXL encoder model. The tokenizer submodel is used for Anima text encoding.", + input=Input.Direct, + ui_model_type=ModelType.T5Encoder, + title="T5 Encoder", + ) + + def invoke(self, context: InvocationContext) -> AnimaModelLoaderOutput: + # Transformer always comes from the main model + transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer}) + + # VAE + vae = self.vae_model.model_copy(update={"submodel_type": SubModelType.VAE}) + + # Qwen3 Encoder + qwen3_tokenizer = self.qwen3_encoder_model.model_copy(update={"submodel_type": SubModelType.Tokenizer}) + qwen3_encoder = self.qwen3_encoder_model.model_copy(update={"submodel_type": SubModelType.TextEncoder}) + + # T5 Encoder (only tokenizer submodel is used by Anima) + t5_tokenizer = preprocess_t5_tokenizer_model_identifier(self.t5_encoder_model) + t5_encoder = preprocess_t5_encoder_model_identifier(self.t5_encoder_model) + + return AnimaModelLoaderOutput( + transformer=TransformerField(transformer=transformer, loras=[]), + qwen3_encoder=Qwen3EncoderField(tokenizer=qwen3_tokenizer, text_encoder=qwen3_encoder), + vae=VAEField(vae=vae), + t5_encoder=T5EncoderField(tokenizer=t5_tokenizer, text_encoder=t5_encoder, loras=[]), + ) diff --git a/invokeai/app/invocations/anima_text_encoder.py b/invokeai/app/invocations/anima_text_encoder.py new file mode 100644 index 00000000000..1856a69ae76 --- /dev/null +++ b/invokeai/app/invocations/anima_text_encoder.py @@ -0,0 +1,221 @@ +"""Anima text encoder invocation. + +Encodes text using the dual-conditioning pipeline: +1. Qwen3 0.6B: Produces hidden states (last layer) +2. T5-XXL Tokenizer: Produces token IDs only (no T5 model needed) + +Both outputs are stored together in AnimaConditioningInfo and used by +the LLM Adapter inside the transformer during denoising. + +Key differences from Z-Image text encoder: +- Anima uses Qwen3 0.6B (base model, NOT instruct) — no chat template +- Anima additionally tokenizes with T5-XXL tokenizer to get token IDs +- Qwen3 output uses all positions (including padding) for full context +""" + +from contextlib import ExitStack +from typing import Iterator, Tuple + +import torch +from transformers import PreTrainedModel, PreTrainedTokenizerBase + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + AnimaConditioningField, + FieldDescriptions, + Input, + InputField, + TensorField, + UIComponent, +) +from invokeai.app.invocations.model import Qwen3EncoderField, T5EncoderField +from invokeai.app.invocations.primitives import AnimaConditioningOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.patches.layer_patcher import LayerPatcher +from invokeai.backend.patches.lora_conversions.anima_lora_constants import ANIMA_LORA_QWEN3_PREFIX +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( + AnimaConditioningInfo, + ConditioningFieldData, +) +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger + +logger = InvokeAILogger.get_logger(__name__) + +# T5-XXL max sequence length for token IDs +T5_MAX_SEQ_LEN = 512 + +# Safety cap for Qwen3 sequence length to prevent GPU OOM on extremely long prompts. +# Qwen3 0.6B supports 32K context but the LLM Adapter doesn't need that much. +QWEN3_MAX_SEQ_LEN = 8192 + + +@invocation( + "anima_text_encoder", + title="Prompt - Anima", + tags=["prompt", "conditioning", "anima"], + category="conditioning", + version="1.3.0", + classification=Classification.Prototype, +) +class AnimaTextEncoderInvocation(BaseInvocation): + """Encodes and preps a prompt for an Anima image. + + Uses Qwen3 0.6B for hidden state extraction and T5-XXL tokenizer for + token IDs (no T5 model weights needed). Both are combined by the + LLM Adapter inside the Anima transformer during denoising. + """ + + prompt: str = InputField(description="Text prompt to encode.", ui_component=UIComponent.Textarea) + qwen3_encoder: Qwen3EncoderField = InputField( + title="Qwen3 Encoder", + description=FieldDescriptions.qwen3_encoder, + input=Input.Connection, + ) + t5_encoder: T5EncoderField = InputField( + title="T5 Encoder", + description=FieldDescriptions.t5_encoder, + input=Input.Connection, + ) + mask: TensorField | None = InputField( + default=None, + description="A mask defining the region that this conditioning prompt applies to.", + ) + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> AnimaConditioningOutput: + qwen3_embeds, t5xxl_ids, t5xxl_weights = self._encode_prompt(context) + + # Move to CPU for storage + qwen3_embeds = qwen3_embeds.detach().to("cpu") + t5xxl_ids = t5xxl_ids.detach().to("cpu") + t5xxl_weights = t5xxl_weights.detach().to("cpu") if t5xxl_weights is not None else None + + conditioning_data = ConditioningFieldData( + conditionings=[ + AnimaConditioningInfo( + qwen3_embeds=qwen3_embeds, + t5xxl_ids=t5xxl_ids, + t5xxl_weights=t5xxl_weights, + ) + ] + ) + conditioning_name = context.conditioning.save(conditioning_data) + return AnimaConditioningOutput( + conditioning=AnimaConditioningField(conditioning_name=conditioning_name, mask=self.mask) + ) + + def _encode_prompt( + self, + context: InvocationContext, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Encode prompt using Qwen3 0.6B and T5-XXL tokenizer. + + Returns: + Tuple of (qwen3_embeds, t5xxl_ids, t5xxl_weights). + - qwen3_embeds: Shape (max_seq_len, 1024) — includes all positions (including padding) + to preserve full sequence context for the LLM Adapter. + - t5xxl_ids: Shape (seq_len,) — T5-XXL token IDs (unpadded). + - t5xxl_weights: None (uniform weights for now). + """ + prompt = self.prompt + + # --- Step 1: Encode with Qwen3 0.6B --- + text_encoder_info = context.models.load(self.qwen3_encoder.text_encoder) + tokenizer_info = context.models.load(self.qwen3_encoder.tokenizer) + + with ExitStack() as exit_stack: + (_, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device()) + (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device()) + + device = text_encoder.device + + # Apply LoRA models to the text encoder + lora_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) + exit_stack.enter_context( + LayerPatcher.apply_smart_model_patches( + model=text_encoder, + patches=self._lora_iterator(context), + prefix=ANIMA_LORA_QWEN3_PREFIX, + dtype=lora_dtype, + ) + ) + + if not isinstance(text_encoder, PreTrainedModel): + raise TypeError(f"Expected PreTrainedModel for text encoder, got {type(text_encoder).__name__}.") + if not isinstance(tokenizer, PreTrainedTokenizerBase): + raise TypeError(f"Expected PreTrainedTokenizerBase for tokenizer, got {type(tokenizer).__name__}.") + + context.util.signal_progress("Running Qwen3 0.6B text encoder") + + # Anima uses base Qwen3 (not instruct) — tokenize directly, no chat template. + # A safety cap is applied to prevent GPU OOM on extremely long prompts. + text_inputs = tokenizer( + prompt, + padding=False, + truncation=True, + max_length=QWEN3_MAX_SEQ_LEN, + return_attention_mask=True, + return_tensors="pt", + ) + + text_input_ids = text_inputs.input_ids + attention_mask = text_inputs.attention_mask + if not isinstance(text_input_ids, torch.Tensor) or not isinstance(attention_mask, torch.Tensor): + raise TypeError("Tokenizer returned unexpected types.") + + if text_input_ids.shape[-1] == QWEN3_MAX_SEQ_LEN: + logger.warning( + f"Prompt was truncated to {QWEN3_MAX_SEQ_LEN} tokens. " + "Consider shortening the prompt for best results." + ) + + # Ensure at least 1 token (empty prompts produce 0 tokens with padding=False) + if text_input_ids.shape[-1] == 0: + pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + text_input_ids = torch.tensor([[pad_id]]) + attention_mask = torch.tensor([[1]]) + + # Get last hidden state from Qwen3 (final layer output) + prompt_mask = attention_mask.to(device).bool() + outputs = text_encoder( + text_input_ids.to(device), + attention_mask=prompt_mask, + output_hidden_states=True, + ) + + if not hasattr(outputs, "hidden_states") or outputs.hidden_states is None: + raise RuntimeError("Text encoder did not return hidden_states.") + if len(outputs.hidden_states) < 1: + raise RuntimeError(f"Expected at least 1 hidden state, got {len(outputs.hidden_states)}.") + + # Use last hidden state — only real tokens, no padding + qwen3_embeds = outputs.hidden_states[-1][0] # Shape: (seq_len, 1024) + + # --- Step 2: Tokenize with T5-XXL tokenizer (IDs only, no model) --- + context.util.signal_progress("Tokenizing with T5-XXL") + t5_tokenizer_info = context.models.load(self.t5_encoder.tokenizer) + with t5_tokenizer_info.model_on_device() as (_, t5_tokenizer): + t5_tokens = t5_tokenizer( + prompt, + padding=False, + truncation=True, + max_length=T5_MAX_SEQ_LEN, + return_tensors="pt", + ) + t5xxl_ids = t5_tokens.input_ids[0] # Shape: (seq_len,) + + return qwen3_embeds, t5xxl_ids, None + + def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]: + """Iterate over LoRA models to apply to the Qwen3 text encoder.""" + for lora in self.qwen3_encoder.loras: + lora_info = context.models.load(lora.lora) + if not isinstance(lora_info.model, ModelPatchRaw): + raise TypeError( + f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}. " + "The LoRA model may be corrupted or incompatible." + ) + yield (lora_info.model, lora.weight) + del lora_info diff --git a/invokeai/app/invocations/canvas.py b/invokeai/app/invocations/canvas.py new file mode 100644 index 00000000000..cf13c3334ff --- /dev/null +++ b/invokeai/app/invocations/canvas.py @@ -0,0 +1,27 @@ +from invokeai.app.invocations.baseinvocation import BaseInvocation, invocation +from invokeai.app.invocations.fields import FieldDescriptions, ImageField, InputField +from invokeai.app.invocations.primitives import ImageOutput +from invokeai.app.services.shared.invocation_context import InvocationContext + + +@invocation( + "canvas_output", + title="Canvas Output", + tags=["canvas", "output", "image"], + category="canvas", + version="1.0.0", + use_cache=False, +) +class CanvasOutputInvocation(BaseInvocation): + """Outputs an image to the canvas staging area. + + Use this node in workflows intended for canvas workflow integration. + Connect the final image of your workflow to this node to send it + to the canvas staging area when run via 'Run Workflow on Canvas'.""" + + image: ImageField = InputField(description=FieldDescriptions.image) + + def invoke(self, context: InvocationContext) -> ImageOutput: + image = context.images.get_pil(self.image.image_name) + image_dto = context.images.save(image=image) + return ImageOutput.build(image_dto) diff --git a/invokeai/app/invocations/compel.py b/invokeai/app/invocations/compel.py index 5ce88145ffd..0ff6be969f2 100644 --- a/invokeai/app/invocations/compel.py +++ b/invokeai/app/invocations/compel.py @@ -19,6 +19,7 @@ from invokeai.app.invocations.primitives import ConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.util.ti_utils import generate_ti_list +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.model_patcher import ModelPatcher from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.model_patch_raw import ModelPatchRaw @@ -103,7 +104,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]: textual_inversion_manager=ti_manager, dtype_for_device_getter=TorchDevice.choose_torch_dtype, truncate_long_prompts=False, - device=text_encoder.device, # Use the device the model is actually on + device=get_effective_device(text_encoder), split_long_text_mode=SplitLongTextMode.SENTENCES, ) @@ -212,7 +213,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]: truncate_long_prompts=False, # TODO: returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED, # TODO: clip skip requires_pooled=get_pooled, - device=text_encoder.device, # Use the device the model is actually on + device=get_effective_device(text_encoder), split_long_text_mode=SplitLongTextMode.SENTENCES, ) diff --git a/invokeai/app/invocations/fields.py b/invokeai/app/invocations/fields.py index cca09a059d5..fbe0e9a6153 100644 --- a/invokeai/app/invocations/fields.py +++ b/invokeai/app/invocations/fields.py @@ -171,6 +171,8 @@ class FieldDescriptions: sd3_model = "SD3 model (MMDiTX) to load" cogview4_model = "CogView4 model (Transformer) to load" z_image_model = "Z-Image model (Transformer) to load" + qwen_image_model = "Qwen Image Edit model (Transformer) to load" + qwen_vl_encoder = "Qwen2.5-VL tokenizer, processor and text/vision encoder" sdxl_main_model = "SDXL Main model (UNet, VAE, CLIP1, CLIP2) to load" sdxl_refiner_model = "SDXL Refiner Main Modde (UNet, VAE, CLIP2) to load" onnx_main_model = "ONNX Main model (UNet, VAE, CLIP) to load" @@ -340,6 +342,27 @@ class ZImageConditioningField(BaseModel): ) +class QwenImageConditioningField(BaseModel): + """A Qwen Image Edit conditioning tensor primitive value""" + + conditioning_name: str = Field(description="The name of conditioning tensor") + + +class AnimaConditioningField(BaseModel): + """An Anima conditioning tensor primitive value. + + Anima conditioning contains Qwen3 0.6B hidden states and T5-XXL token IDs, + which are combined by the LLM Adapter inside the transformer. + """ + + conditioning_name: str = Field(description="The name of conditioning tensor") + mask: Optional[TensorField] = Field( + default=None, + description="The mask associated with this conditioning tensor for regional prompting. " + "Excluded regions should be set to False, included regions should be set to True.", + ) + + class ConditioningField(BaseModel): """A conditioning tensor primitive value""" diff --git a/invokeai/app/invocations/logic.py b/invokeai/app/invocations/logic.py new file mode 100644 index 00000000000..3197427d4e6 --- /dev/null +++ b/invokeai/app/invocations/logic.py @@ -0,0 +1,34 @@ +from typing import Any, Optional + +from invokeai.app.invocations.baseinvocation import BaseInvocation, BaseInvocationOutput, invocation, invocation_output +from invokeai.app.invocations.fields import InputField, OutputField, UIType +from invokeai.app.services.shared.invocation_context import InvocationContext + + +@invocation_output("if_output") +class IfInvocationOutput(BaseInvocationOutput): + value: Optional[Any] = OutputField( + default=None, description="The selected value", title="Output", ui_type=UIType.Any + ) + + +@invocation("if", title="If", tags=["logic", "conditional"], category="logic", version="1.0.0") +class IfInvocation(BaseInvocation): + """Selects between two optional inputs based on a boolean condition.""" + + condition: bool = InputField(default=False, description="The condition used to select an input", title="Condition") + true_input: Optional[Any] = InputField( + default=None, + description="Selected when the condition is true", + title="True Input", + ui_type=UIType.Any, + ) + false_input: Optional[Any] = InputField( + default=None, + description="Selected when the condition is false", + title="False Input", + ui_type=UIType.Any, + ) + + def invoke(self, context: InvocationContext) -> IfInvocationOutput: + return IfInvocationOutput(value=self.true_input if self.condition else self.false_input) diff --git a/invokeai/app/invocations/metadata.py b/invokeai/app/invocations/metadata.py index bc13b72c7bb..da24d8802bb 100644 --- a/invokeai/app/invocations/metadata.py +++ b/invokeai/app/invocations/metadata.py @@ -166,6 +166,14 @@ def invoke(self, context: InvocationContext) -> MetadataOutput: "z_image_img2img", "z_image_inpaint", "z_image_outpaint", + "qwen_image_txt2img", + "qwen_image_img2img", + "qwen_image_inpaint", + "qwen_image_outpaint", + "anima_txt2img", + "anima_img2img", + "anima_inpaint", + "anima_outpaint", ] diff --git a/invokeai/app/invocations/model.py b/invokeai/app/invocations/model.py index 29fbe5100c6..6b5afb55290 100644 --- a/invokeai/app/invocations/model.py +++ b/invokeai/app/invocations/model.py @@ -72,6 +72,13 @@ class GlmEncoderField(BaseModel): text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel") +class QwenVLEncoderField(BaseModel): + """Field for Qwen2.5-VL encoder used by Qwen Image Edit models.""" + + tokenizer: ModelIdentifierField = Field(description="Info to load tokenizer submodel") + text_encoder: ModelIdentifierField = Field(description="Info to load text_encoder submodel") + + class Qwen3EncoderField(BaseModel): """Field for Qwen3 text encoder used by Z-Image models.""" diff --git a/invokeai/app/invocations/primitives.py b/invokeai/app/invocations/primitives.py index dcb1fc6a45f..7ec6c3dc149 100644 --- a/invokeai/app/invocations/primitives.py +++ b/invokeai/app/invocations/primitives.py @@ -12,6 +12,7 @@ ) from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR from invokeai.app.invocations.fields import ( + AnimaConditioningField, BoundingBoxField, CogView4ConditioningField, ColorField, @@ -24,6 +25,7 @@ InputField, LatentsField, OutputField, + QwenImageConditioningField, SD3ConditioningField, TensorField, UIComponent, @@ -473,6 +475,28 @@ def build(cls, conditioning_name: str) -> "ZImageConditioningOutput": return cls(conditioning=ZImageConditioningField(conditioning_name=conditioning_name)) +@invocation_output("qwen_image_conditioning_output") +class QwenImageConditioningOutput(BaseInvocationOutput): + """Base class for nodes that output a Qwen Image Edit conditioning tensor.""" + + conditioning: QwenImageConditioningField = OutputField(description=FieldDescriptions.cond) + + @classmethod + def build(cls, conditioning_name: str) -> "QwenImageConditioningOutput": + return cls(conditioning=QwenImageConditioningField(conditioning_name=conditioning_name)) + + +@invocation_output("anima_conditioning_output") +class AnimaConditioningOutput(BaseInvocationOutput): + """Base class for nodes that output an Anima text conditioning tensor.""" + + conditioning: AnimaConditioningField = OutputField(description=FieldDescriptions.cond) + + @classmethod + def build(cls, conditioning_name: str) -> "AnimaConditioningOutput": + return cls(conditioning=AnimaConditioningField(conditioning_name=conditioning_name)) + + @invocation_output("conditioning_output") class ConditioningOutput(BaseInvocationOutput): """Base class for nodes that output a single conditioning tensor""" diff --git a/invokeai/app/invocations/qwen_image_denoise.py b/invokeai/app/invocations/qwen_image_denoise.py new file mode 100644 index 00000000000..04e21a26c3f --- /dev/null +++ b/invokeai/app/invocations/qwen_image_denoise.py @@ -0,0 +1,490 @@ +from contextlib import ExitStack +from typing import Callable, Iterator, Optional, Tuple + +import torch +import torchvision.transforms as tv_transforms +from diffusers.models.transformers.transformer_qwenimage import QwenImageTransformer2DModel +from torchvision.transforms.functional import resize as tv_resize +from tqdm import tqdm + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.constants import LATENT_SCALE_FACTOR +from invokeai.app.invocations.fields import ( + DenoiseMaskField, + FieldDescriptions, + Input, + InputField, + LatentsField, + QwenImageConditioningField, + WithBoard, + WithMetadata, +) +from invokeai.app.invocations.model import TransformerField +from invokeai.app.invocations.primitives import LatentsOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat +from invokeai.backend.patches.layer_patcher import LayerPatcher +from invokeai.backend.patches.lora_conversions.qwen_image_lora_constants import ( + QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX, +) +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.rectified_flow.rectified_flow_inpaint_extension import RectifiedFlowInpaintExtension +from invokeai.backend.stable_diffusion.diffusers_pipeline import PipelineIntermediateState +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import QwenImageConditioningInfo +from invokeai.backend.util.devices import TorchDevice + + +@invocation( + "qwen_image_denoise", + title="Denoise - Qwen Image", + tags=["image", "qwen_image"], + category="image", + version="1.0.0", + classification=Classification.Prototype, +) +class QwenImageDenoiseInvocation(BaseInvocation, WithMetadata, WithBoard): + """Run the denoising process with a Qwen Image model.""" + + # If latents is provided, this means we are doing image-to-image. + latents: Optional[LatentsField] = InputField( + default=None, description=FieldDescriptions.latents, input=Input.Connection + ) + # Reference image latents (encoded through VAE) to concatenate with noisy latents. + reference_latents: Optional[LatentsField] = InputField( + default=None, + description="Reference image latents to guide generation. Encoded through the VAE.", + input=Input.Connection, + ) + # denoise_mask is used for image-to-image inpainting. Only the masked region is modified. + denoise_mask: Optional[DenoiseMaskField] = InputField( + default=None, description=FieldDescriptions.denoise_mask, input=Input.Connection + ) + denoising_start: float = InputField(default=0.0, ge=0, le=1, description=FieldDescriptions.denoising_start) + denoising_end: float = InputField(default=1.0, ge=0, le=1, description=FieldDescriptions.denoising_end) + transformer: TransformerField = InputField( + description=FieldDescriptions.qwen_image_model, input=Input.Connection, title="Transformer" + ) + positive_conditioning: QwenImageConditioningField = InputField( + description=FieldDescriptions.positive_cond, input=Input.Connection + ) + negative_conditioning: Optional[QwenImageConditioningField] = InputField( + default=None, description=FieldDescriptions.negative_cond, input=Input.Connection + ) + cfg_scale: float | list[float] = InputField(default=4.0, description=FieldDescriptions.cfg_scale, title="CFG Scale") + width: int = InputField(default=1024, multiple_of=16, description="Width of the generated image.") + height: int = InputField(default=1024, multiple_of=16, description="Height of the generated image.") + steps: int = InputField(default=40, gt=0, description=FieldDescriptions.steps) + seed: int = InputField(default=0, description="Randomness seed for reproducibility.") + shift: Optional[float] = InputField( + default=None, + description="Override the sigma schedule shift. " + "When set, uses a fixed shift (e.g. 3.0 for Lightning LoRAs) instead of the default dynamic shifting. " + "Leave unset for the base model's default schedule.", + ) + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> LatentsOutput: + latents = self._run_diffusion(context) + latents = latents.detach().to("cpu") + + name = context.tensors.save(tensor=latents) + return LatentsOutput.build(latents_name=name, latents=latents, seed=None) + + def _prep_inpaint_mask(self, context: InvocationContext, latents: torch.Tensor) -> torch.Tensor | None: + if self.denoise_mask is None: + return None + mask = context.tensors.load(self.denoise_mask.mask_name) + mask = 1.0 - mask + + _, _, latent_height, latent_width = latents.shape + mask = tv_resize( + img=mask, + size=[latent_height, latent_width], + interpolation=tv_transforms.InterpolationMode.BILINEAR, + antialias=False, + ) + + mask = mask.to(device=latents.device, dtype=latents.dtype) + return mask + + def _load_text_conditioning( + self, + context: InvocationContext, + conditioning_name: str, + dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + cond_data = context.conditioning.load(conditioning_name) + assert len(cond_data.conditionings) == 1 + conditioning = cond_data.conditionings[0] + assert isinstance(conditioning, QwenImageConditioningInfo) + conditioning = conditioning.to(dtype=dtype, device=device) + return conditioning.prompt_embeds, conditioning.prompt_embeds_mask + + def _get_noise( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + seed: int, + ) -> torch.Tensor: + rand_device = "cpu" + rand_dtype = torch.float32 + + return torch.randn( + batch_size, + num_channels_latents, + int(height) // LATENT_SCALE_FACTOR, + int(width) // LATENT_SCALE_FACTOR, + device=rand_device, + dtype=rand_dtype, + generator=torch.Generator(device=rand_device).manual_seed(seed), + ).to(device=device, dtype=dtype) + + def _prepare_cfg_scale(self, num_timesteps: int) -> list[float]: + if isinstance(self.cfg_scale, float): + cfg_scale = [self.cfg_scale] * num_timesteps + elif isinstance(self.cfg_scale, list): + assert len(self.cfg_scale) == num_timesteps + cfg_scale = self.cfg_scale + else: + raise ValueError(f"Invalid CFG scale type: {type(self.cfg_scale)}") + return cfg_scale + + @staticmethod + def _pack_latents( + latents: torch.Tensor, batch_size: int, num_channels: int, height: int, width: int + ) -> torch.Tensor: + """Pack 4D latents (B, C, H, W) into 2x2-patched 3D (B, H/2*W/2, C*4).""" + latents = latents.view(batch_size, num_channels, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4) + return latents + + @staticmethod + def _unpack_latents(latents: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Unpack 3D patched latents (B, seq, C*4) back to 4D (B, C, H, W).""" + batch_size, _num_patches, channels = latents.shape + # height/width are in latent space; they must be divisible by 2 for packing + h = 2 * (height // 2) + w = 2 * (width // 2) + latents = latents.view(batch_size, h // 2, w // 2, channels // 4, 2, 2) + latents = latents.permute(0, 3, 1, 4, 2, 5) + latents = latents.reshape(batch_size, channels // 4, h, w) + return latents + + def _run_diffusion(self, context: InvocationContext): + inference_dtype = torch.bfloat16 + device = TorchDevice.choose_torch_device() + + transformer_info = context.models.load(self.transformer.transformer) + assert isinstance(transformer_info.model, QwenImageTransformer2DModel) + + # Load conditioning + pos_prompt_embeds, pos_prompt_mask = self._load_text_conditioning( + context=context, + conditioning_name=self.positive_conditioning.conditioning_name, + dtype=inference_dtype, + device=device, + ) + + neg_prompt_embeds = None + neg_prompt_mask = None + # Match the diffusers pipeline: only enable CFG when cfg_scale > 1 AND negative conditioning is provided. + # With cfg_scale <= 1, the negative prediction is unused, so skip it entirely. + # For per-step arrays, enable CFG if any step has scale > 1. + if isinstance(self.cfg_scale, list): + any_cfg_above_one = any(v > 1.0 for v in self.cfg_scale) + else: + any_cfg_above_one = self.cfg_scale > 1.0 + do_classifier_free_guidance = self.negative_conditioning is not None and any_cfg_above_one + if do_classifier_free_guidance: + neg_prompt_embeds, neg_prompt_mask = self._load_text_conditioning( + context=context, + conditioning_name=self.negative_conditioning.conditioning_name, + dtype=inference_dtype, + device=device, + ) + + # Prepare the timestep / sigma schedule + patch_size = transformer_info.model.config.patch_size + assert isinstance(patch_size, int) + # Output channels is 16 (the actual latent channels) + out_channels = transformer_info.model.config.out_channels + assert isinstance(out_channels, int) + + latent_height = self.height // LATENT_SCALE_FACTOR + latent_width = self.width // LATENT_SCALE_FACTOR + image_seq_len = (latent_height * latent_width) // (patch_size**2) + + # Use the actual FlowMatchEulerDiscreteScheduler to compute sigmas/timesteps, + # exactly matching the diffusers pipeline. + import math + + import numpy as np + from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + # Try to load the scheduler config from the model's directory (Diffusers models + # have a scheduler/ subdir). For GGUF models this path doesn't exist, so fall + # back to instantiating the scheduler with the known Qwen Image defaults. + model_path = context.models.get_absolute_path(context.models.get_config(self.transformer.transformer)) + scheduler_path = model_path / "scheduler" + if scheduler_path.is_dir() and (scheduler_path / "scheduler_config.json").exists(): + scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(str(scheduler_path), local_files_only=True) + else: + scheduler = FlowMatchEulerDiscreteScheduler( + use_dynamic_shifting=True, + base_shift=0.5, + max_shift=0.9, + base_image_seq_len=256, + max_image_seq_len=8192, + shift_terminal=0.02, + num_train_timesteps=1000, + time_shift_type="exponential", + ) + + if self.shift is not None: + # Lightning LoRA: fixed shift + mu = math.log(self.shift) + else: + # Default dynamic shifting + # Linear interpolation matching diffusers' calculate_shift + base_shift = scheduler.config.get("base_shift", 0.5) + max_shift = scheduler.config.get("max_shift", 0.9) + base_seq = scheduler.config.get("base_image_seq_len", 256) + max_seq = scheduler.config.get("max_image_seq_len", 4096) + m = (max_shift - base_shift) / (max_seq - base_seq) + b = base_shift - m * base_seq + mu = image_seq_len * m + b + + init_sigmas = np.linspace(1.0, 1.0 / self.steps, self.steps).tolist() + scheduler.set_timesteps(sigmas=init_sigmas, mu=mu, device=device) + + # Clip the schedule based on denoising_start/denoising_end to support img2img strength. + # The scheduler's sigmas go from high (noisy) to 0 (clean). We clip to the fractional range. + sigmas_sched = scheduler.sigmas # (N+1,) including terminal 0 + if self.denoising_start > 0 or self.denoising_end < 1: + total_sigmas = len(sigmas_sched) - 1 # exclude terminal + start_idx = int(round(self.denoising_start * total_sigmas)) + end_idx = int(round(self.denoising_end * total_sigmas)) + sigmas_sched = sigmas_sched[start_idx : end_idx + 1] # +1 to include the next sigma for dt + # Rebuild timesteps from clipped sigmas (exclude terminal 0) + timesteps_sched = sigmas_sched[:-1] * scheduler.config.num_train_timesteps + else: + timesteps_sched = scheduler.timesteps + + total_steps = len(timesteps_sched) + + cfg_scale = self._prepare_cfg_scale(total_steps) + + # Load initial latents if provided (for img2img) + init_latents = context.tensors.load(self.latents.latents_name) if self.latents else None + if init_latents is not None: + init_latents = init_latents.to(device=device, dtype=inference_dtype) + if init_latents.dim() == 5: + init_latents = init_latents.squeeze(2) + + # Load reference image latents if provided + ref_latents = None + if self.reference_latents is not None: + ref_latents = context.tensors.load(self.reference_latents.latents_name) + ref_latents = ref_latents.to(device=device, dtype=inference_dtype) + # The VAE encoder produces 5D latents (B, C, 1, H, W); squeeze the frame dim + # so we have 4D (B, C, H, W) for packing. + if ref_latents.dim() == 5: + ref_latents = ref_latents.squeeze(2) + + # Generate noise (16 channels - the output latent channels) + noise = self._get_noise( + batch_size=1, + num_channels_latents=out_channels, + height=self.height, + width=self.width, + dtype=inference_dtype, + device=device, + seed=self.seed, + ) + + # Prepare input latent image + if init_latents is not None: + s_0 = sigmas_sched[0].item() + latents = s_0 * noise + (1.0 - s_0) * init_latents + else: + if self.denoising_start > 1e-5: + raise ValueError("denoising_start should be 0 when initial latents are not provided.") + latents = noise + + if total_steps <= 0: + return latents + + # Pack latents into 2x2 patches: (B, C, H, W) -> (B, H/2*W/2, C*4) + latents = self._pack_latents(latents, 1, out_channels, latent_height, latent_width) + + # Determine whether the model uses reference latent conditioning (zero_cond_t). + # Edit models (zero_cond_t=True) expect [noisy_patches ; ref_patches] in the sequence. + # Txt2img models (zero_cond_t=False) only take noisy patches. + has_zero_cond_t = getattr(transformer_info.model, "zero_cond_t", False) or getattr( + transformer_info.model.config, "zero_cond_t", False + ) + use_ref_latents = has_zero_cond_t + + ref_latents_packed = None + if use_ref_latents: + if ref_latents is not None: + _, ref_ch, rh, rw = ref_latents.shape + if rh != latent_height or rw != latent_width: + ref_latents = torch.nn.functional.interpolate( + ref_latents, size=(latent_height, latent_width), mode="bilinear" + ) + else: + # No reference image provided — use zeros so the model still gets the + # expected sequence layout. + ref_latents = torch.zeros( + 1, out_channels, latent_height, latent_width, device=device, dtype=inference_dtype + ) + ref_latents_packed = self._pack_latents(ref_latents, 1, out_channels, latent_height, latent_width) + + # img_shapes tells the transformer the spatial layout of patches. + if use_ref_latents: + img_shapes = [ + [ + (1, latent_height // 2, latent_width // 2), + (1, latent_height // 2, latent_width // 2), + ] + ] + else: + img_shapes = [ + [ + (1, latent_height // 2, latent_width // 2), + ] + ] + + # Prepare inpaint extension (operates in 4D space, so unpack/repack around it) + inpaint_mask = self._prep_inpaint_mask(context, noise) # noise has the right 4D shape + inpaint_extension: RectifiedFlowInpaintExtension | None = None + if inpaint_mask is not None: + assert init_latents is not None + inpaint_extension = RectifiedFlowInpaintExtension( + init_latents=init_latents, + inpaint_mask=inpaint_mask, + noise=noise, + ) + + step_callback = self._build_step_callback(context) + + step_callback( + PipelineIntermediateState( + step=0, + order=1, + total_steps=total_steps, + timestep=int(timesteps_sched[0].item()) if len(timesteps_sched) > 0 else 0, + latents=self._unpack_latents(latents, latent_height, latent_width), + ), + ) + + noisy_seq_len = latents.shape[1] + + # Determine if the model is quantized — GGUF models need sidecar patching for LoRAs + transformer_config = context.models.get_config(self.transformer.transformer) + model_is_quantized = transformer_config.format in (ModelFormat.GGUFQuantized,) + + with ExitStack() as exit_stack: + (cached_weights, transformer) = exit_stack.enter_context(transformer_info.model_on_device()) + assert isinstance(transformer, QwenImageTransformer2DModel) + + # Apply LoRA patches to the transformer + exit_stack.enter_context( + LayerPatcher.apply_smart_model_patches( + model=transformer, + patches=self._lora_iterator(context), + prefix=QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX, + dtype=inference_dtype, + cached_weights=cached_weights, + force_sidecar_patching=model_is_quantized, + ) + ) + + for step_idx, t in enumerate(tqdm(timesteps_sched)): + # The pipeline passes timestep / 1000 to the transformer + timestep = t.expand(latents.shape[0]).to(inference_dtype) + + # For edit models: concatenate noisy and reference patches along the sequence dim + # For txt2img models: just use noisy patches + if ref_latents_packed is not None: + model_input = torch.cat([latents, ref_latents_packed], dim=1) + else: + model_input = latents + + noise_pred_cond = transformer( + hidden_states=model_input, + encoder_hidden_states=pos_prompt_embeds, + encoder_hidden_states_mask=pos_prompt_mask, + timestep=timestep / 1000, + img_shapes=img_shapes, + return_dict=False, + )[0] + # Only keep the noisy-latent portion of the output + noise_pred_cond = noise_pred_cond[:, :noisy_seq_len] + + if do_classifier_free_guidance and neg_prompt_embeds is not None: + noise_pred_uncond = transformer( + hidden_states=model_input, + encoder_hidden_states=neg_prompt_embeds, + encoder_hidden_states_mask=neg_prompt_mask, + timestep=timestep / 1000, + img_shapes=img_shapes, + return_dict=False, + )[0] + noise_pred_uncond = noise_pred_uncond[:, :noisy_seq_len] + + noise_pred = noise_pred_uncond + cfg_scale[step_idx] * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond + + # Euler step using the (possibly clipped) sigma schedule + sigma_curr = sigmas_sched[step_idx] + sigma_next = sigmas_sched[step_idx + 1] + dt = sigma_next - sigma_curr + latents = latents.to(torch.float32) + dt * noise_pred.to(torch.float32) + latents = latents.to(inference_dtype) + + if inpaint_extension is not None: + sigma_next = sigmas_sched[step_idx + 1].item() + latents_4d = self._unpack_latents(latents, latent_height, latent_width) + latents_4d = inpaint_extension.merge_intermediate_latents_with_init_latents(latents_4d, sigma_next) + latents = self._pack_latents(latents_4d, 1, out_channels, latent_height, latent_width) + + step_callback( + PipelineIntermediateState( + step=step_idx + 1, + order=1, + total_steps=total_steps, + timestep=int(t.item()), + latents=self._unpack_latents(latents, latent_height, latent_width), + ), + ) + + # Unpack back to 4D then add frame dim for the video-style VAE: (B, C, 1, H, W) + latents = self._unpack_latents(latents, latent_height, latent_width) + latents = latents.unsqueeze(2) + return latents + + def _build_step_callback(self, context: InvocationContext) -> Callable[[PipelineIntermediateState], None]: + def step_callback(state: PipelineIntermediateState) -> None: + context.util.sd_step_callback(state, BaseModelType.QwenImage) + + return step_callback + + def _lora_iterator(self, context: InvocationContext) -> Iterator[Tuple[ModelPatchRaw, float]]: + """Iterate over LoRA models to apply to the transformer.""" + for lora in self.transformer.loras: + lora_info = context.models.load(lora.lora) + if not isinstance(lora_info.model, ModelPatchRaw): + raise TypeError( + f"Expected ModelPatchRaw for LoRA '{lora.lora.key}', got {type(lora_info.model).__name__}." + ) + yield (lora_info.model, lora.weight) + del lora_info diff --git a/invokeai/app/invocations/qwen_image_image_to_latents.py b/invokeai/app/invocations/qwen_image_image_to_latents.py new file mode 100644 index 00000000000..c5fe1b5d5c8 --- /dev/null +++ b/invokeai/app/invocations/qwen_image_image_to_latents.py @@ -0,0 +1,96 @@ +import einops +import torch +from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage +from PIL import Image as PILImage + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + FieldDescriptions, + ImageField, + Input, + InputField, + WithBoard, + WithMetadata, +) +from invokeai.app.invocations.model import VAEField +from invokeai.app.invocations.primitives import LatentsOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.load_base import LoadedModel +from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor +from invokeai.backend.util.devices import TorchDevice + + +@invocation( + "qwen_image_i2l", + title="Image to Latents - Qwen Image", + tags=["image", "latents", "vae", "i2l", "qwen_image"], + category="image", + version="1.0.0", + classification=Classification.Prototype, +) +class QwenImageImageToLatentsInvocation(BaseInvocation, WithMetadata, WithBoard): + """Generates latents from an image using the Qwen Image VAE.""" + + image: ImageField = InputField(description="The image to encode.") + vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection) + width: int | None = InputField( + default=None, + description="Resize the image to this width before encoding. If not set, encodes at the image's original size.", + ) + height: int | None = InputField( + default=None, + description="Resize the image to this height before encoding. If not set, encodes at the image's original size.", + ) + + @staticmethod + def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tensor: + with vae_info.model_on_device() as (_, vae): + assert isinstance(vae, AutoencoderKLQwenImage) + + vae.disable_tiling() + + image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype) + with torch.inference_mode(): + # The Qwen Image VAE expects 5D input: (B, C, num_frames, H, W) + if image_tensor.dim() == 4: + image_tensor = image_tensor.unsqueeze(2) + + posterior = vae.encode(image_tensor).latent_dist + # Use mode (argmax) for deterministic encoding, matching diffusers + latents: torch.Tensor = posterior.mode().to(dtype=vae.dtype) + + # Normalize with per-channel latents_mean / latents_std + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(vae.config.latents_std) + .view(1, vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = (latents - latents_mean) / latents_std + + return latents + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> LatentsOutput: + image = context.images.get_pil(self.image.image_name) + + # If target dimensions are specified, resize the image BEFORE encoding + # (matching the diffusers pipeline which resizes in pixel space, not latent space). + if self.width is not None and self.height is not None: + image = image.convert("RGB").resize((self.width, self.height), resample=PILImage.LANCZOS) + + image_tensor = image_resized_to_grid_as_tensor(image.convert("RGB")) + if image_tensor.dim() == 3: + image_tensor = einops.rearrange(image_tensor, "c h w -> 1 c h w") + + vae_info = context.models.load(self.vae.vae) + + latents = self.vae_encode(vae_info=vae_info, image_tensor=image_tensor) + + latents = latents.to("cpu") + name = context.tensors.save(tensor=latents) + return LatentsOutput.build(latents_name=name, latents=latents, seed=None) diff --git a/invokeai/app/invocations/qwen_image_latents_to_image.py b/invokeai/app/invocations/qwen_image_latents_to_image.py new file mode 100644 index 00000000000..b3ea39c4bbf --- /dev/null +++ b/invokeai/app/invocations/qwen_image_latents_to_image.py @@ -0,0 +1,85 @@ +from contextlib import nullcontext + +import torch +from diffusers.models.autoencoders.autoencoder_kl_qwenimage import AutoencoderKLQwenImage +from einops import rearrange +from PIL import Image + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + FieldDescriptions, + Input, + InputField, + LatentsField, + WithBoard, + WithMetadata, +) +from invokeai.app.invocations.model import VAEField +from invokeai.app.invocations.primitives import ImageOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.stable_diffusion.extensions.seamless import SeamlessExt +from invokeai.backend.util.devices import TorchDevice + + +@invocation( + "qwen_image_l2i", + title="Latents to Image - Qwen Image", + tags=["latents", "image", "vae", "l2i", "qwen_image"], + category="latents", + version="1.0.0", + classification=Classification.Prototype, +) +class QwenImageLatentsToImageInvocation(BaseInvocation, WithMetadata, WithBoard): + """Generates an image from latents using the Qwen Image VAE.""" + + latents: LatentsField = InputField(description=FieldDescriptions.latents, input=Input.Connection) + vae: VAEField = InputField(description=FieldDescriptions.vae, input=Input.Connection) + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> ImageOutput: + latents = context.tensors.load(self.latents.latents_name) + + vae_info = context.models.load(self.vae.vae) + assert isinstance(vae_info.model, AutoencoderKLQwenImage) + with ( + SeamlessExt.static_patch_model(vae_info.model, self.vae.seamless_axes), + vae_info.model_on_device() as (_, vae), + ): + context.util.signal_progress("Running VAE") + assert isinstance(vae, AutoencoderKLQwenImage) + latents = latents.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype) + + vae.disable_tiling() + + tiling_context = nullcontext() + + TorchDevice.empty_cache() + + with torch.inference_mode(), tiling_context: + # The Qwen Image VAE uses per-channel latents_mean / latents_std + # instead of a single scaling_factor. + # Latents are 5D: (B, C, num_frames, H, W) — the unpack from the + # denoise step already produces this shape. + latents_mean = ( + torch.tensor(vae.config.latents_mean) + .view(1, vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + + img = vae.decode(latents, return_dict=False)[0] + # Drop the temporal frame dimension: (B, C, 1, H, W) -> (B, C, H, W) + img = img[:, :, 0] + + img = img.clamp(-1, 1) + img = rearrange(img[0], "c h w -> h w c") + img_pil = Image.fromarray((127.5 * (img + 1.0)).byte().cpu().numpy()) + + TorchDevice.empty_cache() + + image_dto = context.images.save(image=img_pil) + + return ImageOutput.build(image_dto) diff --git a/invokeai/app/invocations/qwen_image_lora_loader.py b/invokeai/app/invocations/qwen_image_lora_loader.py new file mode 100644 index 00000000000..f670b2d8954 --- /dev/null +++ b/invokeai/app/invocations/qwen_image_lora_loader.py @@ -0,0 +1,115 @@ +from typing import Optional + +from invokeai.app.invocations.baseinvocation import ( + BaseInvocation, + BaseInvocationOutput, + Classification, + invocation, + invocation_output, +) +from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField +from invokeai.app.invocations.model import LoRAField, ModelIdentifierField, TransformerField +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelType + + +@invocation_output("qwen_image_lora_loader_output") +class QwenImageLoRALoaderOutput(BaseInvocationOutput): + """Qwen Image LoRA Loader Output""" + + transformer: Optional[TransformerField] = OutputField( + default=None, description=FieldDescriptions.transformer, title="Transformer" + ) + + +@invocation( + "qwen_image_lora_loader", + title="Apply LoRA - Qwen Image", + tags=["lora", "model", "qwen_image"], + category="model", + version="1.0.0", + classification=Classification.Prototype, +) +class QwenImageLoRALoaderInvocation(BaseInvocation): + """Apply a LoRA model to a Qwen Image transformer.""" + + lora: ModelIdentifierField = InputField( + description=FieldDescriptions.lora_model, + title="LoRA", + ui_model_base=BaseModelType.QwenImage, + ui_model_type=ModelType.LoRA, + ) + weight: float = InputField(default=1.0, description=FieldDescriptions.lora_weight) + transformer: TransformerField | None = InputField( + default=None, + description=FieldDescriptions.transformer, + input=Input.Connection, + title="Transformer", + ) + + def invoke(self, context: InvocationContext) -> QwenImageLoRALoaderOutput: + lora_key = self.lora.key + + if not context.models.exists(lora_key): + raise ValueError(f"Unknown lora: {lora_key}!") + + if self.transformer and any(lora.lora.key == lora_key for lora in self.transformer.loras): + raise ValueError(f'LoRA "{lora_key}" already applied to transformer.') + + output = QwenImageLoRALoaderOutput() + + if self.transformer is not None: + output.transformer = self.transformer.model_copy(deep=True) + output.transformer.loras.append( + LoRAField( + lora=self.lora, + weight=self.weight, + ) + ) + + return output + + +@invocation( + "qwen_image_lora_collection_loader", + title="Apply LoRA Collection - Qwen Image", + tags=["lora", "model", "qwen_image"], + category="model", + version="1.0.0", + classification=Classification.Prototype, +) +class QwenImageLoRACollectionLoader(BaseInvocation): + """Applies a collection of LoRAs to a Qwen Image transformer.""" + + loras: Optional[LoRAField | list[LoRAField]] = InputField( + default=None, description="LoRA models and weights. May be a single LoRA or collection.", title="LoRAs" + ) + transformer: Optional[TransformerField] = InputField( + default=None, + description=FieldDescriptions.transformer, + input=Input.Connection, + title="Transformer", + ) + + def invoke(self, context: InvocationContext) -> QwenImageLoRALoaderOutput: + output = QwenImageLoRALoaderOutput() + loras = self.loras if isinstance(self.loras, list) else [self.loras] + added_loras: list[str] = [] + + if self.transformer is not None: + output.transformer = self.transformer.model_copy(deep=True) + + for lora in loras: + if lora is None: + continue + if lora.lora.key in added_loras: + continue + if not context.models.exists(lora.lora.key): + raise Exception(f"Unknown lora: {lora.lora.key}!") + + added_loras.append(lora.lora.key) + + if self.transformer is not None and output.transformer is not None: + output.transformer.loras.append(lora) + + return output diff --git a/invokeai/app/invocations/qwen_image_model_loader.py b/invokeai/app/invocations/qwen_image_model_loader.py new file mode 100644 index 00000000000..fd96067f561 --- /dev/null +++ b/invokeai/app/invocations/qwen_image_model_loader.py @@ -0,0 +1,107 @@ +from typing import Optional + +from invokeai.app.invocations.baseinvocation import ( + BaseInvocation, + BaseInvocationOutput, + Classification, + invocation, + invocation_output, +) +from invokeai.app.invocations.fields import FieldDescriptions, Input, InputField, OutputField +from invokeai.app.invocations.model import ( + ModelIdentifierField, + QwenVLEncoderField, + TransformerField, + VAEField, +) +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType + + +@invocation_output("qwen_image_model_loader_output") +class QwenImageModelLoaderOutput(BaseInvocationOutput): + """Qwen Image model loader output.""" + + transformer: TransformerField = OutputField(description=FieldDescriptions.transformer, title="Transformer") + qwen_vl_encoder: QwenVLEncoderField = OutputField( + description=FieldDescriptions.qwen_vl_encoder, title="Qwen VL Encoder" + ) + vae: VAEField = OutputField(description=FieldDescriptions.vae, title="VAE") + + +@invocation( + "qwen_image_model_loader", + title="Main Model - Qwen Image", + tags=["model", "qwen_image"], + category="model", + version="1.1.0", + classification=Classification.Prototype, +) +class QwenImageModelLoaderInvocation(BaseInvocation): + """Loads a Qwen Image model, outputting its submodels. + + The transformer is always loaded from the main model (Diffusers or GGUF). + + For GGUF quantized models, the VAE and Qwen VL encoder must come from a + separate Diffusers model specified in the "Component Source" field. + + For Diffusers models, all components are extracted from the main model + automatically. The "Component Source" field is ignored. + """ + + model: ModelIdentifierField = InputField( + description=FieldDescriptions.qwen_image_model, + input=Input.Direct, + ui_model_base=BaseModelType.QwenImage, + ui_model_type=ModelType.Main, + title="Transformer", + ) + + component_source: Optional[ModelIdentifierField] = InputField( + default=None, + description="Diffusers Qwen Image model to extract the VAE and Qwen VL encoder from. " + "Required when using a GGUF quantized transformer. " + "Ignored when the main model is already in Diffusers format.", + input=Input.Direct, + ui_model_base=BaseModelType.QwenImage, + ui_model_type=ModelType.Main, + ui_model_format=ModelFormat.Diffusers, + title="Component Source (Diffusers)", + ) + + def invoke(self, context: InvocationContext) -> QwenImageModelLoaderOutput: + main_config = context.models.get_config(self.model) + main_is_diffusers = main_config.format == ModelFormat.Diffusers + + # Transformer always comes from the main model + transformer = self.model.model_copy(update={"submodel_type": SubModelType.Transformer}) + + if main_is_diffusers: + # Diffusers model: extract all components directly + vae = self.model.model_copy(update={"submodel_type": SubModelType.VAE}) + tokenizer = self.model.model_copy(update={"submodel_type": SubModelType.Tokenizer}) + text_encoder = self.model.model_copy(update={"submodel_type": SubModelType.TextEncoder}) + elif self.component_source is not None: + # GGUF/checkpoint transformer: get VAE + encoder from the component source + source_config = context.models.get_config(self.component_source) + if source_config.format != ModelFormat.Diffusers: + raise ValueError( + f"The Component Source model must be in Diffusers format. " + f"The selected model '{source_config.name}' is in {source_config.format.value} format." + ) + vae = self.component_source.model_copy(update={"submodel_type": SubModelType.VAE}) + tokenizer = self.component_source.model_copy(update={"submodel_type": SubModelType.Tokenizer}) + text_encoder = self.component_source.model_copy(update={"submodel_type": SubModelType.TextEncoder}) + else: + raise ValueError( + "No source for VAE and Qwen VL encoder. " + "GGUF quantized models only contain the transformer — " + "please set 'Component Source' to a Diffusers Qwen Image model " + "to provide the VAE and text encoder." + ) + + return QwenImageModelLoaderOutput( + transformer=TransformerField(transformer=transformer, loras=[]), + qwen_vl_encoder=QwenVLEncoderField(tokenizer=tokenizer, text_encoder=text_encoder), + vae=VAEField(vae=vae), + ) diff --git a/invokeai/app/invocations/qwen_image_text_encoder.py b/invokeai/app/invocations/qwen_image_text_encoder.py new file mode 100644 index 00000000000..a067421452e --- /dev/null +++ b/invokeai/app/invocations/qwen_image_text_encoder.py @@ -0,0 +1,298 @@ +from typing import Literal + +import torch +from PIL import Image as PILImage + +from invokeai.app.invocations.baseinvocation import BaseInvocation, Classification, invocation +from invokeai.app.invocations.fields import ( + FieldDescriptions, + ImageField, + Input, + InputField, + UIComponent, +) +from invokeai.app.invocations.model import QwenVLEncoderField +from invokeai.app.invocations.primitives import QwenImageConditioningOutput +from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ( + ConditioningFieldData, + QwenImageConditioningInfo, +) + +# Prompt templates and drop indices for the two Qwen Image model modes. +# These are taken directly from the diffusers pipelines. + +# Image editing mode (QwenImagePipeline) +_EDIT_SYSTEM_PROMPT = ( + "Describe the key features of the input image (color, shape, size, texture, objects, background), " + "then explain how the user's text instruction should alter or modify the image. " + "Generate a new image that meets the user's requirements while maintaining consistency " + "with the original input where appropriate." +) +_EDIT_DROP_IDX = 64 + +# Text-to-image mode (QwenImagePipeline) +_GENERATE_SYSTEM_PROMPT = ( + "Describe the image by detailing the color, shape, size, texture, quantity, " + "text, spatial relationships of the objects and background:" +) +_GENERATE_DROP_IDX = 34 + +_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>" + + +def _build_prompt(user_prompt: str, num_images: int) -> str: + """Build the full prompt with the appropriate template based on whether reference images are provided.""" + if num_images > 0: + # Edit mode: include vision placeholders for reference images + image_tokens = _IMAGE_PLACEHOLDER * num_images + return ( + f"<|im_start|>system\n{_EDIT_SYSTEM_PROMPT}<|im_end|>\n" + f"<|im_start|>user\n{image_tokens}{user_prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + else: + # Generate mode: text-only prompt + return ( + f"<|im_start|>system\n{_GENERATE_SYSTEM_PROMPT}<|im_end|>\n" + f"<|im_start|>user\n{user_prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + + +@invocation( + "qwen_image_text_encoder", + title="Prompt - Qwen Image", + tags=["prompt", "conditioning", "qwen_image"], + category="conditioning", + version="1.2.0", + classification=Classification.Prototype, +) +class QwenImageTextEncoderInvocation(BaseInvocation): + """Encodes text and reference images for Qwen Image using Qwen2.5-VL.""" + + prompt: str = InputField(description="Text prompt describing the desired edit.", ui_component=UIComponent.Textarea) + reference_images: list[ImageField] = InputField( + default=[], + description="Reference images to guide the edit. The model can use multiple reference images.", + ) + qwen_vl_encoder: QwenVLEncoderField = InputField( + title="Qwen VL Encoder", + description=FieldDescriptions.qwen_vl_encoder, + input=Input.Connection, + ) + quantization: Literal["none", "int8", "nf4"] = InputField( + default="none", + description="Quantize the Qwen VL encoder to reduce VRAM usage. " + "'nf4' (4-bit) saves the most memory, 'int8' (8-bit) is a middle ground.", + ) + + @staticmethod + def _resize_for_vl_encoder(image: PILImage.Image, target_pixels: int = 512 * 512) -> PILImage.Image: + """Resize image to fit within target_pixels while preserving aspect ratio. + + Matches the diffusers pipeline's calculate_dimensions logic: the image is resized + so its total pixel count is approximately target_pixels, with dimensions rounded to + multiples of 32. This prevents large images from producing too many vision tokens + which can overwhelm the text prompt. + """ + w, h = image.size + aspect = w / h + # Compute dimensions that preserve aspect ratio at ~target_pixels total + new_w = int((target_pixels * aspect) ** 0.5) + new_h = int(target_pixels / new_w) + # Round to multiples of 32 + new_w = max(32, (new_w // 32) * 32) + new_h = max(32, (new_h // 32) * 32) + if new_w != w or new_h != h: + image = image.resize((new_w, new_h), resample=PILImage.LANCZOS) + return image + + @torch.no_grad() + def invoke(self, context: InvocationContext) -> QwenImageConditioningOutput: + # Load and resize reference images to ~1M pixels (matching diffusers pipeline) + pil_images: list[PILImage.Image] = [] + for img_field in self.reference_images: + pil_img = context.images.get_pil(img_field.image_name) + pil_img = self._resize_for_vl_encoder(pil_img.convert("RGB")) + pil_images.append(pil_img) + + prompt_embeds, prompt_mask = self._encode(context, pil_images) + prompt_embeds = prompt_embeds.detach().to("cpu") + prompt_mask = prompt_mask.detach().to("cpu") if prompt_mask is not None else None + + conditioning_data = ConditioningFieldData( + conditionings=[QwenImageConditioningInfo(prompt_embeds=prompt_embeds, prompt_embeds_mask=prompt_mask)] + ) + conditioning_name = context.conditioning.save(conditioning_data) + return QwenImageConditioningOutput.build(conditioning_name) + + def _encode( + self, context: InvocationContext, images: list[PILImage.Image] + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Encode text prompt and reference images using Qwen2.5-VL. + + Matches the diffusers QwenImagePipeline._get_qwen_prompt_embeds logic: + 1. Format prompt with the edit-specific system template + 2. Run through Qwen2.5-VL to get hidden states + 3. Extract valid (non-padding) tokens and drop the system prefix + 4. Return padded embeddings + attention mask + """ + from transformers import AutoTokenizer, Qwen2_5_VLProcessor + + try: + from transformers import Qwen2_5_VLImageProcessor as _ImageProcessorCls + except ImportError: + from transformers.models.qwen2_vl.image_processing_qwen2_vl import ( # type: ignore[no-redef] + Qwen2VLImageProcessor as _ImageProcessorCls, + ) + + try: + from transformers import Qwen2_5_VLVideoProcessor as _VideoProcessorCls + except ImportError: + from transformers.models.qwen2_vl.video_processing_qwen2_vl import ( # type: ignore[no-redef] + Qwen2VLVideoProcessor as _VideoProcessorCls, + ) + + # Format the prompt with one vision placeholder per reference image + text = _build_prompt(self.prompt, len(images)) + + # Build the processor + tokenizer_config = context.models.get_config(self.qwen_vl_encoder.tokenizer) + model_root = context.models.get_absolute_path(tokenizer_config) + tokenizer_dir = model_root / "tokenizer" + + tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_dir), local_files_only=True) + + image_processor = None + for search_dir in [model_root / "processor", tokenizer_dir, model_root, model_root / "image_processor"]: + if (search_dir / "preprocessor_config.json").exists(): + image_processor = _ImageProcessorCls.from_pretrained(str(search_dir), local_files_only=True) + break + if image_processor is None: + image_processor = _ImageProcessorCls() + + processor = Qwen2_5_VLProcessor( + tokenizer=tokenizer, + image_processor=image_processor, + video_processor=_VideoProcessorCls(), + ) + + context.util.signal_progress("Running Qwen2.5-VL text/vision encoder") + + if self.quantization != "none": + text_encoder, device, cleanup = self._load_quantized_encoder(context) + else: + text_encoder, device, cleanup = self._load_cached_encoder(context) + + try: + model_inputs = processor( + text=[text], + images=images if images else None, + padding=True, + return_tensors="pt", + ).to(device=device) + + outputs = text_encoder( + input_ids=model_inputs.input_ids, + attention_mask=model_inputs.attention_mask, + pixel_values=getattr(model_inputs, "pixel_values", None), + image_grid_thw=getattr(model_inputs, "image_grid_thw", None), + output_hidden_states=True, + ) + + # Use last hidden state (matching diffusers pipeline) + hidden_states = outputs.hidden_states[-1] + + # Extract valid (non-padding) tokens using the attention mask, + # then drop the system prompt prefix tokens. + # The drop index differs between edit mode (64) and generate mode (34). + drop_idx = _EDIT_DROP_IDX if images else _GENERATE_DROP_IDX + + attn_mask = model_inputs.attention_mask + bool_mask = attn_mask.bool() + valid_lengths = bool_mask.sum(dim=1) + selected = hidden_states[bool_mask] + split_hidden = torch.split(selected, valid_lengths.tolist(), dim=0) + + # Drop system prefix tokens and build padded output + trimmed = [h[drop_idx:] for h in split_hidden] + attn_mask_list = [torch.ones(h.size(0), dtype=torch.long, device=device) for h in trimmed] + max_seq_len = max(h.size(0) for h in trimmed) + + prompt_embeds = torch.stack( + [torch.cat([h, h.new_zeros(max_seq_len - h.size(0), h.size(1))]) for h in trimmed] + ) + encoder_attention_mask = torch.stack( + [torch.cat([m, m.new_zeros(max_seq_len - m.size(0))]) for m in attn_mask_list] + ) + + prompt_embeds = prompt_embeds.to(dtype=torch.bfloat16) + finally: + if cleanup is not None: + cleanup() + + # If all tokens are valid (no padding), mask is not needed + if encoder_attention_mask.all(): + encoder_attention_mask = None + + return prompt_embeds, encoder_attention_mask + + def _load_cached_encoder(self, context: InvocationContext): + """Load the text encoder through the model cache (no quantization).""" + from transformers import Qwen2_5_VLForConditionalGeneration + + text_encoder_info = context.models.load(self.qwen_vl_encoder.text_encoder) + ctx = text_encoder_info.model_on_device() + _, text_encoder = ctx.__enter__() + device = get_effective_device(text_encoder) + assert isinstance(text_encoder, Qwen2_5_VLForConditionalGeneration) + return text_encoder, device, lambda: ctx.__exit__(None, None, None) + + def _load_quantized_encoder(self, context: InvocationContext): + """Load the text encoder with BitsAndBytes quantization, bypassing the model cache. + + BnB-quantized models are pinned to GPU and can't be moved between devices, + so they can't go through the standard model cache. The model is loaded fresh + each time and freed after use via the cleanup callback. + """ + import gc + import warnings + + from transformers import BitsAndBytesConfig, Qwen2_5_VLForConditionalGeneration + + encoder_config = context.models.get_config(self.qwen_vl_encoder.text_encoder) + model_root = context.models.get_absolute_path(encoder_config) + encoder_path = model_root / "text_encoder" + + if self.quantization == "nf4": + bnb_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=torch.bfloat16, + bnb_4bit_quant_type="nf4", + ) + else: # int8 + bnb_config = BitsAndBytesConfig(load_in_8bit=True) + + context.util.signal_progress("Loading Qwen2.5-VL encoder (quantized)") + with warnings.catch_warnings(): + # BnB int8 internally casts bfloat16→float16; the warning is harmless + warnings.filterwarnings("ignore", message="MatMul8bitLt.*cast.*float16") + text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained( + str(encoder_path), + quantization_config=bnb_config, + device_map="auto", + torch_dtype=torch.bfloat16, + local_files_only=True, + ) + + device = next(text_encoder.parameters()).device + + def cleanup(): + nonlocal text_encoder + del text_encoder + gc.collect() + torch.cuda.empty_cache() + + return text_encoder, device, cleanup diff --git a/invokeai/app/invocations/sd3_text_encoder.py b/invokeai/app/invocations/sd3_text_encoder.py index 24647c9cfc7..58880f9a289 100644 --- a/invokeai/app/invocations/sd3_text_encoder.py +++ b/invokeai/app/invocations/sd3_text_encoder.py @@ -16,6 +16,7 @@ from invokeai.app.invocations.model import CLIPField, T5EncoderField from invokeai.app.invocations.primitives import SD3ConditioningOutput from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device from invokeai.backend.model_manager.taxonomy import ModelFormat from invokeai.backend.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX @@ -103,6 +104,7 @@ def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tens context.util.signal_progress("Running T5 encoder") assert isinstance(t5_text_encoder, T5EncoderModel) assert isinstance(t5_tokenizer, (T5Tokenizer, T5TokenizerFast)) + t5_device = get_effective_device(t5_text_encoder) text_inputs = t5_tokenizer( prompt, @@ -125,7 +127,7 @@ def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tens f" {max_seq_len} tokens: {removed_text}" ) - prompt_embeds = t5_text_encoder(text_input_ids.to(t5_text_encoder.device))[0] + prompt_embeds = t5_text_encoder(text_input_ids.to(t5_device))[0] assert isinstance(prompt_embeds, torch.Tensor) return prompt_embeds @@ -144,6 +146,7 @@ def _clip_encode( context.util.signal_progress("Running CLIP encoder") assert isinstance(clip_text_encoder, (CLIPTextModel, CLIPTextModelWithProjection)) assert isinstance(clip_tokenizer, CLIPTokenizer) + clip_device = get_effective_device(clip_text_encoder) clip_text_encoder_config = clip_text_encoder_info.config assert clip_text_encoder_config is not None @@ -187,9 +190,7 @@ def _clip_encode( "The following part of your input was truncated because CLIP can only handle sequences up to" f" {tokenizer_max_length} tokens: {removed_text}" ) - prompt_embeds = clip_text_encoder( - input_ids=text_input_ids.to(clip_text_encoder.device), output_hidden_states=True - ) + prompt_embeds = clip_text_encoder(input_ids=text_input_ids.to(clip_device), output_hidden_states=True) pooled_prompt_embeds = prompt_embeds[0] prompt_embeds = prompt_embeds.hidden_states[-2] diff --git a/invokeai/app/invocations/z_image_denoise.py b/invokeai/app/invocations/z_image_denoise.py index 6ab43d66574..24b135e4479 100644 --- a/invokeai/app/invocations/z_image_denoise.py +++ b/invokeai/app/invocations/z_image_denoise.py @@ -50,7 +50,7 @@ title="Denoise - Z-Image", tags=["image", "z-image"], category="image", - version="1.4.0", + version="1.5.0", classification=Classification.Prototype, ) class ZImageDenoiseInvocation(BaseInvocation): @@ -104,6 +104,15 @@ class ZImageDenoiseInvocation(BaseInvocation): description=FieldDescriptions.vae + " Required for control conditioning.", input=Input.Connection, ) + # Shift override for the sigma schedule. If None, shift is auto-calculated from image dimensions. + shift: Optional[float] = InputField( + default=None, + ge=0.0, + description="Override the timestep shift (mu) for the sigma schedule. " + "Leave blank to auto-calculate based on image dimensions (recommended). " + "Lower values (~0.5) produce less noise shifting, higher values (~1.15) produce more.", + title="Shift", + ) # Scheduler selection for the denoising process scheduler: ZIMAGE_SCHEDULER_NAME_VALUES = InputField( default="euler", @@ -225,34 +234,36 @@ def _calculate_shift( """Calculate timestep shift based on image sequence length. Based on diffusers ZImagePipeline.calculate_shift method. + Returns a linear shift value (exp(mu) from the original formula). """ + import math + m = (max_shift - base_shift) / (max_image_seq_len - base_image_seq_len) b = base_shift - m * base_image_seq_len mu = image_seq_len * m + b - return mu + # Convert from exponential mu to linear shift value + return math.exp(mu) - def _get_sigmas(self, mu: float, num_steps: int) -> list[float]: - """Generate sigma schedule with time shift. + def _get_sigmas(self, shift: float, num_steps: int) -> list[float]: + """Generate sigma schedule with linear time shift. - Based on FlowMatchEulerDiscreteScheduler with shift. + Uses linear time shift: shift / (shift + (1/t - 1)). + The shift value is used directly as a multiplier. Generates num_steps + 1 sigma values (including terminal 0.0). """ - import math - def time_shift(mu: float, sigma: float, t: float) -> float: - """Apply time shift to a single timestep value.""" + def time_shift(shift: float, t: float) -> float: + """Apply linear time shift to a single timestep value.""" if t <= 0: return 0.0 if t >= 1: return 1.0 - return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + return shift / (shift + (1 / t - 1)) - # Generate linearly spaced values from 1 to 0 (excluding endpoints for safety) - # then apply time shift sigmas = [] for i in range(num_steps + 1): t = 1.0 - i / num_steps # Goes from 1.0 to 0.0 - sigma = time_shift(mu, 1.0, t) + sigma = time_shift(shift, t) sigmas.append(sigma) return sigmas @@ -313,11 +324,14 @@ def _run_diffusion(self, context: InvocationContext) -> torch.Tensor: # Concatenate all negative embeddings neg_prompt_embeds = torch.cat([tc.prompt_embeds for tc in neg_text_conditionings], dim=0) - # Calculate shift based on image sequence length - mu = self._calculate_shift(img_seq_len) + # Calculate shift based on image sequence length, or use override + if self.shift is not None: + shift = self.shift + else: + shift = self._calculate_shift(img_seq_len) # Generate sigma schedule with time shift - sigmas = self._get_sigmas(mu, self.steps) + sigmas = self._get_sigmas(shift, self.steps) # Apply denoising_start and denoising_end clipping if self.denoising_start > 0 or self.denoising_end < 1: diff --git a/invokeai/app/services/auth/token_service.py b/invokeai/app/services/auth/token_service.py index 9c35261c380..2d766bb90aa 100644 --- a/invokeai/app/services/auth/token_service.py +++ b/invokeai/app/services/auth/token_service.py @@ -21,6 +21,7 @@ class TokenData(BaseModel): user_id: str email: str is_admin: bool + remember_me: bool = False def set_jwt_secret(secret: str) -> None: diff --git a/invokeai/app/services/board_records/board_records_common.py b/invokeai/app/services/board_records/board_records_common.py index ab6355a3930..b263f264cb8 100644 --- a/invokeai/app/services/board_records/board_records_common.py +++ b/invokeai/app/services/board_records/board_records_common.py @@ -9,6 +9,17 @@ from invokeai.app.util.model_exclude_null import BaseModelExcludeNull +class BoardVisibility(str, Enum, metaclass=MetaEnum): + """The visibility options for a board.""" + + Private = "private" + """Only the board owner (and admins) can see and modify this board.""" + Shared = "shared" + """All users can view this board, but only the owner (and admins) can modify it.""" + Public = "public" + """All users can view this board; only the owner (and admins) can modify its structure.""" + + class BoardRecord(BaseModelExcludeNull): """Deserialized board record.""" @@ -28,6 +39,10 @@ class BoardRecord(BaseModelExcludeNull): """The name of the cover image of the board.""" archived: bool = Field(description="Whether or not the board is archived.") """Whether or not the board is archived.""" + board_visibility: BoardVisibility = Field( + default=BoardVisibility.Private, description="The visibility of the board." + ) + """The visibility of the board (private, shared, or public).""" def deserialize_board_record(board_dict: dict) -> BoardRecord: @@ -44,6 +59,11 @@ def deserialize_board_record(board_dict: dict) -> BoardRecord: updated_at = board_dict.get("updated_at", get_iso_timestamp()) deleted_at = board_dict.get("deleted_at", get_iso_timestamp()) archived = board_dict.get("archived", False) + board_visibility_raw = board_dict.get("board_visibility", BoardVisibility.Private.value) + try: + board_visibility = BoardVisibility(board_visibility_raw) + except ValueError: + board_visibility = BoardVisibility.Private return BoardRecord( board_id=board_id, @@ -54,6 +74,7 @@ def deserialize_board_record(board_dict: dict) -> BoardRecord: updated_at=updated_at, deleted_at=deleted_at, archived=archived, + board_visibility=board_visibility, ) @@ -61,6 +82,7 @@ class BoardChanges(BaseModel, extra="forbid"): board_name: Optional[str] = Field(default=None, description="The board's new name.", max_length=300) cover_image_name: Optional[str] = Field(default=None, description="The name of the board's new cover image.") archived: Optional[bool] = Field(default=None, description="Whether or not the board is archived") + board_visibility: Optional[BoardVisibility] = Field(default=None, description="The visibility of the board.") class BoardRecordOrderBy(str, Enum, metaclass=MetaEnum): diff --git a/invokeai/app/services/board_records/board_records_sqlite.py b/invokeai/app/services/board_records/board_records_sqlite.py index a54f65686fd..1e3e11c8a36 100644 --- a/invokeai/app/services/board_records/board_records_sqlite.py +++ b/invokeai/app/services/board_records/board_records_sqlite.py @@ -116,6 +116,17 @@ def update( (changes.archived, board_id), ) + # Change the visibility of a board + if changes.board_visibility is not None: + cursor.execute( + """--sql + UPDATE boards + SET board_visibility = ? + WHERE board_id = ?; + """, + (changes.board_visibility.value, board_id), + ) + except sqlite3.Error as e: raise BoardRecordSaveException from e return self.get(board_id) @@ -155,7 +166,7 @@ def get_many( SELECT DISTINCT boards.* FROM boards LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.is_public = 1) + WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) {archived_filter} ORDER BY {order_by} {direction} LIMIT ? OFFSET ?; @@ -194,14 +205,14 @@ def get_many( SELECT COUNT(DISTINCT boards.board_id) FROM boards LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.is_public = 1); + WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')); """ else: count_query = """ SELECT COUNT(DISTINCT boards.board_id) FROM boards LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.is_public = 1) + WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) AND boards.archived = 0; """ @@ -251,7 +262,7 @@ def get_all( SELECT DISTINCT boards.* FROM boards LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.is_public = 1) + WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) {archived_filter} ORDER BY LOWER(boards.board_name) {direction} """ @@ -260,7 +271,7 @@ def get_all( SELECT DISTINCT boards.* FROM boards LEFT JOIN shared_boards ON boards.board_id = shared_boards.board_id - WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.is_public = 1) + WHERE (boards.user_id = ? OR shared_boards.user_id = ? OR boards.board_visibility IN ('shared', 'public')) {archived_filter} ORDER BY {order_by} {direction} """ diff --git a/invokeai/app/services/bulk_download/bulk_download_base.py b/invokeai/app/services/bulk_download/bulk_download_base.py index 617b611f566..6cd4ed0cbaf 100644 --- a/invokeai/app/services/bulk_download/bulk_download_base.py +++ b/invokeai/app/services/bulk_download/bulk_download_base.py @@ -7,7 +7,11 @@ class BulkDownloadBase(ABC): @abstractmethod def handler( - self, image_names: Optional[list[str]], board_id: Optional[str], bulk_download_item_id: Optional[str] + self, + image_names: Optional[list[str]], + board_id: Optional[str], + bulk_download_item_id: Optional[str], + user_id: str = "system", ) -> None: """ Create a zip file containing the images specified by the given image names or board id. @@ -15,6 +19,7 @@ def handler( :param image_names: A list of image names to include in the zip file. :param board_id: The ID of the board. If provided, all images associated with the board will be included in the zip file. :param bulk_download_item_id: The bulk_download_item_id that will be used to retrieve the bulk download item when it is prepared, if none is provided a uuid will be generated. + :param user_id: The ID of the user who initiated the download. """ @abstractmethod @@ -42,3 +47,12 @@ def delete(self, bulk_download_item_name: str) -> None: :param bulk_download_item_name: The name of the bulk download item. """ + + @abstractmethod + def get_owner(self, bulk_download_item_name: str) -> Optional[str]: + """ + Get the user_id of the user who initiated the download. + + :param bulk_download_item_name: The name of the bulk download item. + :return: The user_id of the owner, or None if not tracked. + """ diff --git a/invokeai/app/services/bulk_download/bulk_download_default.py b/invokeai/app/services/bulk_download/bulk_download_default.py index dc4f8b1d81b..c037e9c5c15 100644 --- a/invokeai/app/services/bulk_download/bulk_download_default.py +++ b/invokeai/app/services/bulk_download/bulk_download_default.py @@ -25,15 +25,24 @@ def __init__(self): self._temp_directory = TemporaryDirectory() self._bulk_downloads_folder = Path(self._temp_directory.name) / "bulk_downloads" self._bulk_downloads_folder.mkdir(parents=True, exist_ok=True) + # Track which user owns each download so the fetch endpoint can enforce ownership + self._download_owners: dict[str, str] = {} def handler( - self, image_names: Optional[list[str]], board_id: Optional[str], bulk_download_item_id: Optional[str] + self, + image_names: Optional[list[str]], + board_id: Optional[str], + bulk_download_item_id: Optional[str], + user_id: str = "system", ) -> None: bulk_download_id: str = DEFAULT_BULK_DOWNLOAD_ID bulk_download_item_id = bulk_download_item_id or uuid_string() bulk_download_item_name = bulk_download_item_id + ".zip" - self._signal_job_started(bulk_download_id, bulk_download_item_id, bulk_download_item_name) + # Record ownership so the fetch endpoint can verify the caller + self._download_owners[bulk_download_item_name] = user_id + + self._signal_job_started(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id) try: image_dtos: list[ImageDTO] = [] @@ -46,16 +55,16 @@ def handler( raise BulkDownloadParametersException() bulk_download_item_name: str = self._create_zip_file(image_dtos, bulk_download_item_id) - self._signal_job_completed(bulk_download_id, bulk_download_item_id, bulk_download_item_name) + self._signal_job_completed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id) except ( ImageRecordNotFoundException, BoardRecordNotFoundException, BulkDownloadException, BulkDownloadParametersException, ) as e: - self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e) + self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e, user_id) except Exception as e: - self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e) + self._signal_job_failed(bulk_download_id, bulk_download_item_id, bulk_download_item_name, e, user_id) self._invoker.services.logger.error("Problem bulk downloading images.") raise e @@ -103,43 +112,60 @@ def _clean_string_to_path_safe(self, s: str) -> str: return "".join([c for c in s if c.isalpha() or c.isdigit() or c == " " or c == "_" or c == "-"]).rstrip() def _signal_job_started( - self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str + self, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + user_id: str = "system", ) -> None: """Signal that a bulk download job has started.""" if self._invoker: assert bulk_download_id is not None self._invoker.services.events.emit_bulk_download_started( - bulk_download_id, bulk_download_item_id, bulk_download_item_name + bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id=user_id ) def _signal_job_completed( - self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str + self, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + user_id: str = "system", ) -> None: """Signal that a bulk download job has completed.""" if self._invoker: assert bulk_download_id is not None assert bulk_download_item_name is not None self._invoker.services.events.emit_bulk_download_complete( - bulk_download_id, bulk_download_item_id, bulk_download_item_name + bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id=user_id ) def _signal_job_failed( - self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, exception: Exception + self, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + exception: Exception, + user_id: str = "system", ) -> None: """Signal that a bulk download job has failed.""" if self._invoker: assert bulk_download_id is not None assert exception is not None self._invoker.services.events.emit_bulk_download_error( - bulk_download_id, bulk_download_item_id, bulk_download_item_name, str(exception) + bulk_download_id, bulk_download_item_id, bulk_download_item_name, str(exception), user_id=user_id ) def stop(self, *args, **kwargs): self._temp_directory.cleanup() + def get_owner(self, bulk_download_item_name: str) -> Optional[str]: + return self._download_owners.get(bulk_download_item_name) + def delete(self, bulk_download_item_name: str) -> None: path = self.get_path(bulk_download_item_name) Path(path).unlink() + self._download_owners.pop(bulk_download_item_name, None) def get_path(self, bulk_download_item_name: str) -> str: path = str(self._bulk_downloads_folder / bulk_download_item_name) diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index aa1cbb5e0ee..935b422a732 100644 --- a/invokeai/app/services/events/events_base.py +++ b/invokeai/app/services/events/events_base.py @@ -100,9 +100,9 @@ def emit_queue_item_status_changed( """Emitted when a queue item's status changes""" self.dispatch(QueueItemStatusChangedEvent.build(queue_item, batch_status, queue_status)) - def emit_batch_enqueued(self, enqueue_result: "EnqueueBatchResult") -> None: + def emit_batch_enqueued(self, enqueue_result: "EnqueueBatchResult", user_id: str = "system") -> None: """Emitted when a batch is enqueued""" - self.dispatch(BatchEnqueuedEvent.build(enqueue_result)) + self.dispatch(BatchEnqueuedEvent.build(enqueue_result, user_id)) def emit_queue_items_retried(self, retry_result: "RetryItemsResult") -> None: """Emitted when a list of queue items are retried""" @@ -112,9 +112,9 @@ def emit_queue_cleared(self, queue_id: str) -> None: """Emitted when a queue is cleared""" self.dispatch(QueueClearedEvent.build(queue_id)) - def emit_recall_parameters_updated(self, queue_id: str, parameters: dict) -> None: + def emit_recall_parameters_updated(self, queue_id: str, user_id: str, parameters: dict) -> None: """Emitted when recall parameters are updated""" - self.dispatch(RecallParametersUpdatedEvent.build(queue_id, parameters)) + self.dispatch(RecallParametersUpdatedEvent.build(queue_id, user_id, parameters)) # endregion @@ -194,23 +194,42 @@ def emit_model_install_error(self, job: "ModelInstallJob") -> None: # region Bulk image download def emit_bulk_download_started( - self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str + self, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + user_id: str = "system", ) -> None: """Emitted when a bulk image download is started""" - self.dispatch(BulkDownloadStartedEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name)) + self.dispatch( + BulkDownloadStartedEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id) + ) def emit_bulk_download_complete( - self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str + self, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + user_id: str = "system", ) -> None: """Emitted when a bulk image download is complete""" - self.dispatch(BulkDownloadCompleteEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name)) + self.dispatch( + BulkDownloadCompleteEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name, user_id) + ) def emit_bulk_download_error( - self, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, error: str + self, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + error: str, + user_id: str = "system", ) -> None: """Emitted when a bulk image download has an error""" self.dispatch( - BulkDownloadErrorEvent.build(bulk_download_id, bulk_download_item_id, bulk_download_item_name, error) + BulkDownloadErrorEvent.build( + bulk_download_id, bulk_download_item_id, bulk_download_item_name, error, user_id + ) ) # endregion diff --git a/invokeai/app/services/events/events_common.py b/invokeai/app/services/events/events_common.py index bfb44eb48e8..998fe4f5309 100644 --- a/invokeai/app/services/events/events_common.py +++ b/invokeai/app/services/events/events_common.py @@ -281,9 +281,10 @@ class BatchEnqueuedEvent(QueueEventBase): ) priority: int = Field(description="The priority of the batch") origin: str | None = Field(default=None, description="The origin of the batch") + user_id: str = Field(default="system", description="The ID of the user who enqueued the batch") @classmethod - def build(cls, enqueue_result: EnqueueBatchResult) -> "BatchEnqueuedEvent": + def build(cls, enqueue_result: EnqueueBatchResult, user_id: str = "system") -> "BatchEnqueuedEvent": return cls( queue_id=enqueue_result.queue_id, batch_id=enqueue_result.batch.batch_id, @@ -291,6 +292,7 @@ def build(cls, enqueue_result: EnqueueBatchResult) -> "BatchEnqueuedEvent": enqueued=enqueue_result.enqueued, requested=enqueue_result.requested, priority=enqueue_result.priority, + user_id=user_id, ) @@ -609,6 +611,7 @@ class BulkDownloadEventBase(EventBase): bulk_download_id: str = Field(description="The ID of the bulk image download") bulk_download_item_id: str = Field(description="The ID of the bulk image download item") bulk_download_item_name: str = Field(description="The name of the bulk image download item") + user_id: str = Field(default="system", description="The ID of the user who initiated the download") @payload_schema.register @@ -619,12 +622,17 @@ class BulkDownloadStartedEvent(BulkDownloadEventBase): @classmethod def build( - cls, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str + cls, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + user_id: str = "system", ) -> "BulkDownloadStartedEvent": return cls( bulk_download_id=bulk_download_id, bulk_download_item_id=bulk_download_item_id, bulk_download_item_name=bulk_download_item_name, + user_id=user_id, ) @@ -636,12 +644,17 @@ class BulkDownloadCompleteEvent(BulkDownloadEventBase): @classmethod def build( - cls, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str + cls, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + user_id: str = "system", ) -> "BulkDownloadCompleteEvent": return cls( bulk_download_id=bulk_download_id, bulk_download_item_id=bulk_download_item_id, bulk_download_item_name=bulk_download_item_name, + user_id=user_id, ) @@ -655,13 +668,19 @@ class BulkDownloadErrorEvent(BulkDownloadEventBase): @classmethod def build( - cls, bulk_download_id: str, bulk_download_item_id: str, bulk_download_item_name: str, error: str + cls, + bulk_download_id: str, + bulk_download_item_id: str, + bulk_download_item_name: str, + error: str, + user_id: str = "system", ) -> "BulkDownloadErrorEvent": return cls( bulk_download_id=bulk_download_id, bulk_download_item_id=bulk_download_item_id, bulk_download_item_name=bulk_download_item_name, error=error, + user_id=user_id, ) @@ -671,8 +690,9 @@ class RecallParametersUpdatedEvent(QueueEventBase): __event_name__ = "recall_parameters_updated" + user_id: str = Field(description="The ID of the user whose recall parameters were updated") parameters: dict[str, Any] = Field(description="The recall parameters that were updated") @classmethod - def build(cls, queue_id: str, parameters: dict[str, Any]) -> "RecallParametersUpdatedEvent": - return cls(queue_id=queue_id, parameters=parameters) + def build(cls, queue_id: str, user_id: str, parameters: dict[str, Any]) -> "RecallParametersUpdatedEvent": + return cls(queue_id=queue_id, user_id=user_id, parameters=parameters) diff --git a/invokeai/app/services/events/events_fastapievents.py b/invokeai/app/services/events/events_fastapievents.py index f44eecc5559..90e1402773d 100644 --- a/invokeai/app/services/events/events_fastapievents.py +++ b/invokeai/app/services/events/events_fastapievents.py @@ -46,3 +46,9 @@ async def _dispatch_from_queue(self, stop_event: threading.Event): except asyncio.CancelledError as e: raise e # Raise a proper error + except Exception: + import logging + + logging.getLogger("InvokeAI").error( + f"Error dispatching event {getattr(event, '__event_name__', event)}", exc_info=True + ) diff --git a/invokeai/app/services/image_records/image_records_base.py b/invokeai/app/services/image_records/image_records_base.py index 16405c52708..457cf2f4686 100644 --- a/invokeai/app/services/image_records/image_records_base.py +++ b/invokeai/app/services/image_records/image_records_base.py @@ -74,8 +74,8 @@ def delete_intermediates(self) -> list[str]: pass @abstractmethod - def get_intermediates_count(self) -> int: - """Gets a count of all intermediate images.""" + def get_intermediates_count(self, user_id: Optional[str] = None) -> int: + """Gets a count of intermediate images. If user_id is provided, only counts that user's intermediates.""" pass @abstractmethod @@ -97,6 +97,11 @@ def save( """Saves an image record.""" pass + @abstractmethod + def get_user_id(self, image_name: str) -> Optional[str]: + """Gets the user_id of the image owner. Returns None if image not found.""" + pass + @abstractmethod def get_most_recent_image_for_board(self, board_id: str) -> Optional[ImageRecord]: """Gets the most recent image for a board.""" diff --git a/invokeai/app/services/image_records/image_records_sqlite.py b/invokeai/app/services/image_records/image_records_sqlite.py index c6c237fc1e7..07126d53a9f 100644 --- a/invokeai/app/services/image_records/image_records_sqlite.py +++ b/invokeai/app/services/image_records/image_records_sqlite.py @@ -46,6 +46,20 @@ def get(self, image_name: str) -> ImageRecord: return deserialize_image_record(dict(result)) + def get_user_id(self, image_name: str) -> Optional[str]: + with self._db.transaction() as cursor: + cursor.execute( + """--sql + SELECT user_id FROM images + WHERE image_name = ?; + """, + (image_name,), + ) + result = cast(Optional[sqlite3.Row], cursor.fetchone()) + if not result: + return None + return cast(Optional[str], dict(result).get("user_id")) + def get_metadata(self, image_name: str) -> Optional[MetadataField]: with self._db.transaction() as cursor: try: @@ -269,14 +283,14 @@ def delete_many(self, image_names: list[str]) -> None: except sqlite3.Error as e: raise ImageRecordDeleteException from e - def get_intermediates_count(self) -> int: + def get_intermediates_count(self, user_id: Optional[str] = None) -> int: with self._db.transaction() as cursor: - cursor.execute( - """--sql - SELECT COUNT(*) FROM images - WHERE is_intermediate = TRUE; - """ - ) + query = "SELECT COUNT(*) FROM images WHERE is_intermediate = TRUE" + params: list[str] = [] + if user_id is not None: + query += " AND user_id = ?" + params.append(user_id) + cursor.execute(query, params) count = cast(int, cursor.fetchone()[0]) return count diff --git a/invokeai/app/services/images/images_base.py b/invokeai/app/services/images/images_base.py index d11d75b3c1d..aebbead2f35 100644 --- a/invokeai/app/services/images/images_base.py +++ b/invokeai/app/services/images/images_base.py @@ -143,8 +143,8 @@ def delete_intermediates(self) -> int: pass @abstractmethod - def get_intermediates_count(self) -> int: - """Gets the number of intermediate images.""" + def get_intermediates_count(self, user_id: Optional[str] = None) -> int: + """Gets the number of intermediate images. If user_id is provided, only counts that user's intermediates.""" pass @abstractmethod diff --git a/invokeai/app/services/images/images_default.py b/invokeai/app/services/images/images_default.py index e82bd7f4de1..0f03f7c400e 100644 --- a/invokeai/app/services/images/images_default.py +++ b/invokeai/app/services/images/images_default.py @@ -310,9 +310,9 @@ def delete_intermediates(self) -> int: self.__invoker.services.logger.error("Problem deleting image records and files") raise e - def get_intermediates_count(self) -> int: + def get_intermediates_count(self, user_id: Optional[str] = None) -> int: try: - return self.__invoker.services.image_records.get_intermediates_count() + return self.__invoker.services.image_records.get_intermediates_count(user_id=user_id) except Exception as e: self.__invoker.services.logger.error("Problem getting intermediates count") raise e diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 8503811bcd1..91acc22d03f 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -17,7 +17,7 @@ import torch import yaml -from huggingface_hub import HfFolder +from huggingface_hub import get_token as hf_get_token from pydantic.networks import AnyHttpUrl from pydantic_core import Url from requests import Session @@ -1024,7 +1024,7 @@ def _import_from_hf( ) -> ModelInstallJob: # Add user's cached access token to HuggingFace requests if source.access_token is None: - source.access_token = HfFolder.get_token() + source.access_token = hf_get_token() remote_files, metadata = self._remote_files_from_source(source) return self._import_remote_model( source=source, diff --git a/invokeai/app/services/model_records/model_records_base.py b/invokeai/app/services/model_records/model_records_base.py index 96e12d3b0a3..dcdc0ce5956 100644 --- a/invokeai/app/services/model_records/model_records_base.py +++ b/invokeai/app/services/model_records/model_records_base.py @@ -26,6 +26,7 @@ ModelType, ModelVariantType, Qwen3VariantType, + QwenImageVariantType, SchedulerPredictionType, ZImageVariantType, ) @@ -94,7 +95,13 @@ class ModelRecordChanges(BaseModelExcludeNull): # Checkpoint-specific changes # TODO(MM2): Should we expose these? Feels footgun-y... variant: Optional[ - ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType + ModelVariantType + | ClipVariantType + | FluxVariantType + | Flux2VariantType + | ZImageVariantType + | QwenImageVariantType + | Qwen3VariantType ] = Field(description="The variant of the model.", default=None) prediction_type: Optional[SchedulerPredictionType] = Field( description="The prediction type of the model.", default=None diff --git a/invokeai/app/services/session_queue/session_queue_base.py b/invokeai/app/services/session_queue/session_queue_base.py index 3c037dc77ab..14b93d97fc7 100644 --- a/invokeai/app/services/session_queue/session_queue_base.py +++ b/invokeai/app/services/session_queue/session_queue_base.py @@ -78,13 +78,15 @@ def get_queue_status(self, queue_id: str, user_id: Optional[str] = None) -> Sess pass @abstractmethod - def get_counts_by_destination(self, queue_id: str, destination: str) -> SessionQueueCountsByDestination: - """Gets the counts of queue items by destination""" + def get_counts_by_destination( + self, queue_id: str, destination: str, user_id: Optional[str] = None + ) -> SessionQueueCountsByDestination: + """Gets the counts of queue items by destination. If user_id is provided, only counts that user's items.""" pass @abstractmethod - def get_batch_status(self, queue_id: str, batch_id: str) -> BatchStatus: - """Gets the status of a batch""" + def get_batch_status(self, queue_id: str, batch_id: str, user_id: Optional[str] = None) -> BatchStatus: + """Gets the status of a batch. If user_id is provided, only counts that user's items.""" pass @abstractmethod @@ -172,8 +174,9 @@ def get_queue_item_ids( self, queue_id: str, order_dir: SQLiteDirection = SQLiteDirection.Descending, + user_id: Optional[str] = None, ) -> ItemIdsResult: - """Gets all queue item ids that match the given parameters""" + """Gets all queue item ids that match the given parameters. If user_id is provided, only returns items for that user.""" pass @abstractmethod diff --git a/invokeai/app/services/session_queue/session_queue_common.py b/invokeai/app/services/session_queue/session_queue_common.py index 58544422119..09820fe6217 100644 --- a/invokeai/app/services/session_queue/session_queue_common.py +++ b/invokeai/app/services/session_queue/session_queue_common.py @@ -304,12 +304,6 @@ class SessionQueueStatus(BaseModel): failed: int = Field(..., description="Number of queue items with status 'error'") canceled: int = Field(..., description="Number of queue items with status 'canceled'") total: int = Field(..., description="Total number of queue items") - user_pending: Optional[int] = Field( - default=None, description="Number of queue items with status 'pending' for the current user" - ) - user_in_progress: Optional[int] = Field( - default=None, description="Number of queue items with status 'in_progress' for the current user" - ) class SessionQueueCountsByDestination(BaseModel): diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py index 4f46136fd79..070a7cef293 100644 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ b/invokeai/app/services/session_queue/session_queue_sqlite.py @@ -151,7 +151,7 @@ async def enqueue_batch( priority=priority, item_ids=item_ids, ) - self.__invoker.services.events.emit_batch_enqueued(enqueue_result) + self.__invoker.services.events.emit_batch_enqueued(enqueue_result, user_id=user_id) return enqueue_result def dequeue(self) -> Optional[SessionQueueItem]: @@ -765,15 +765,21 @@ def get_queue_item_ids( self, queue_id: str, order_dir: SQLiteDirection = SQLiteDirection.Descending, + user_id: Optional[str] = None, ) -> ItemIdsResult: with self._db.transaction() as cursor_: - query = f"""--sql + query = """--sql SELECT item_id FROM session_queue WHERE queue_id = ? - ORDER BY created_at {order_dir.value} """ - query_params = [queue_id] + query_params: list[str] = [queue_id] + + if user_id is not None: + query += " AND user_id = ?" + query_params.append(user_id) + + query += f" ORDER BY created_at {order_dir.value}" cursor_.execute(query, query_params) result = cast(list[sqlite3.Row], cursor_.fetchall()) @@ -783,20 +789,7 @@ def get_queue_item_ids( def get_queue_status(self, queue_id: str, user_id: Optional[str] = None) -> SessionQueueStatus: with self._db.transaction() as cursor: - # Get total counts - cursor.execute( - """--sql - SELECT status, count(*) - FROM session_queue - WHERE queue_id = ? - GROUP BY status - """, - (queue_id,), - ) - counts_result = cast(list[sqlite3.Row], cursor.fetchall()) - - # Get user-specific counts if user_id is provided (using a single query with CASE) - user_counts_result = [] + # When user_id is provided (non-admin), only count that user's items if user_id is not None: cursor.execute( """--sql @@ -807,48 +800,51 @@ def get_queue_status(self, queue_id: str, user_id: Optional[str] = None) -> Sess """, (queue_id, user_id), ) - user_counts_result = cast(list[sqlite3.Row], cursor.fetchall()) + else: + cursor.execute( + """--sql + SELECT status, count(*) + FROM session_queue + WHERE queue_id = ? + GROUP BY status + """, + (queue_id,), + ) + counts_result = cast(list[sqlite3.Row], cursor.fetchall()) current_item = self.get_current(queue_id=queue_id) total = sum(row[1] or 0 for row in counts_result) counts: dict[str, int] = {row[0]: row[1] for row in counts_result} - # Process user-specific counts if available - user_pending = None - user_in_progress = None - if user_id is not None: - user_counts: dict[str, int] = {row[0]: row[1] for row in user_counts_result} - user_pending = user_counts.get("pending", 0) - user_in_progress = user_counts.get("in_progress", 0) + # For non-admin users, hide current item details if they don't own it + show_current_item = current_item is not None and (user_id is None or current_item.user_id == user_id) return SessionQueueStatus( queue_id=queue_id, - item_id=current_item.item_id if current_item else None, - session_id=current_item.session_id if current_item else None, - batch_id=current_item.batch_id if current_item else None, + item_id=current_item.item_id if show_current_item else None, + session_id=current_item.session_id if show_current_item else None, + batch_id=current_item.batch_id if show_current_item else None, pending=counts.get("pending", 0), in_progress=counts.get("in_progress", 0), completed=counts.get("completed", 0), failed=counts.get("failed", 0), canceled=counts.get("canceled", 0), total=total, - user_pending=user_pending, - user_in_progress=user_in_progress, ) - def get_batch_status(self, queue_id: str, batch_id: str) -> BatchStatus: + def get_batch_status(self, queue_id: str, batch_id: str, user_id: Optional[str] = None) -> BatchStatus: with self._db.transaction() as cursor: - cursor.execute( - """--sql + query = """--sql SELECT status, count(*), origin, destination FROM session_queue - WHERE - queue_id = ? - AND batch_id = ? - GROUP BY status - """, - (queue_id, batch_id), - ) + WHERE queue_id = ? AND batch_id = ? + """ + params: list[str] = [queue_id, batch_id] + if user_id is not None: + query += " AND user_id = ?" + params.append(user_id) + query += " GROUP BY status" + cursor.execute(query, params) result = cast(list[sqlite3.Row], cursor.fetchall()) total = sum(row[1] or 0 for row in result) counts: dict[str, int] = {row[0]: row[1] for row in result} @@ -868,18 +864,21 @@ def get_batch_status(self, queue_id: str, batch_id: str) -> BatchStatus: total=total, ) - def get_counts_by_destination(self, queue_id: str, destination: str) -> SessionQueueCountsByDestination: + def get_counts_by_destination( + self, queue_id: str, destination: str, user_id: Optional[str] = None + ) -> SessionQueueCountsByDestination: with self._db.transaction() as cursor: - cursor.execute( - """--sql + query = """--sql SELECT status, count(*) FROM session_queue - WHERE queue_id = ? - AND destination = ? - GROUP BY status - """, - (queue_id, destination), - ) + WHERE queue_id = ? AND destination = ? + """ + params: list[str] = [queue_id, destination] + if user_id is not None: + query += " AND user_id = ?" + params.append(user_id) + query += " GROUP BY status" + cursor.execute(query, params) counts_result = cast(list[sqlite3.Row], cursor.fetchall()) total = sum(row[1] or 0 for row in counts_result) diff --git a/invokeai/app/services/shared/README.md b/invokeai/app/services/shared/README.md index a65a8ebe495..113b7a41e54 100644 --- a/invokeai/app/services/shared/README.md +++ b/invokeai/app/services/shared/README.md @@ -6,32 +6,32 @@ High-level design for the graph module. Focuses on responsibilities, data flow, Provide a typed, acyclic workflow model (**Graph**) plus a runtime scheduler (**GraphExecutionState**) that expands iterator patterns, tracks readiness via indegree (the number of incoming edges to a node in the directed graph), and -executes nodes in class-grouped batches. Source graphs remain immutable during a run; runtime expansion happens in a -separate execution graph. +executes nodes in class-grouped batches. In normal execution, runtime expansion happens in a separate execution graph +instead of mutating the source graph. ## 2) Major Data Types ### EdgeConnection -* Fields: `node_id: str`, `field: str`. -* Hashable; printed as `node.field` for readable diagnostics. +- Fields: `node_id: str`, `field: str`. +- Hashable; printed as `node.field` for readable diagnostics. ### Edge -* Fields: `source: EdgeConnection`, `destination: EdgeConnection`. -* One directed connection from a specific output port to a specific input port. +- Fields: `source: EdgeConnection`, `destination: EdgeConnection`. +- One directed connection from a specific output port to a specific input port. ### AnyInvocation / AnyInvocationOutput -* Pydantic wrappers that carry concrete invocation models and outputs. -* No registry logic in this file; they are permissive containers for heterogeneous nodes. +- Pydantic wrappers that carry concrete invocation models and outputs. +- No registry logic in this file; they are permissive containers for heterogeneous nodes. ### IterateInvocation / CollectInvocation -* Control nodes used by validation and execution: +- Control nodes used by validation and execution: - * **IterateInvocation**: input `collection`, outputs include `item` (and index/total). - * **CollectInvocation**: many `item` inputs aggregated to one `collection` output. + - **IterateInvocation**: input `collection`, outputs include `item` (and index/total). + - **CollectInvocation**: many `item` inputs aggregated to one `collection` output. ## 3) Graph (author-time model) @@ -39,156 +39,209 @@ A container for declared nodes and edges. Does **not** perform iteration expansi ### 3.1 Data -* `nodes: dict[str, AnyInvocation]` - key must equal `node.id`. -* `edges: list[Edge]` - zero or more. -* Utility: `_get_input_edges(node_id, field?)`, `_get_output_edges(node_id, field?)` - These scan `self.edges` (no adjacency indices in the current code). +- `nodes: dict[str, AnyInvocation]` - key must equal `node.id`. +- `edges: list[Edge]` - zero or more. +- Utility: `_get_input_edges(node_id, field?)`, `_get_output_edges(node_id, field?)` These scan `self.edges` (no + adjacency indices in the current code). ### 3.2 Validation (`validate_self`) Runs a sequence of checks: -1. **Node ID uniqueness** - No duplicate IDs; map key equals `node.id`. -2. **Endpoint existence** - Source and destination node IDs must exist. -3. **Port existence** - Input ports must exist on the node class; output ports on the node's output model. -4. **Type compatibility** - `get_output_field_type` vs `get_input_field_type` and `are_connection_types_compatible`. -5. **DAG constraint** - Build a *flat* `DiGraph` (no runtime expansion) and assert acyclicity. -6. **Iterator / collector structure** - Enforce special rules: - - * Iterator's input must be `collection`; its outgoing edges use `item`. - * Collector accepts many `item` inputs; outputs a single `collection`. - * Edge fan-in to a non-collector input is rejected. +1. **Node ID uniqueness** No duplicate IDs; map key equals `node.id`. + +1. **Endpoint existence** Source and destination node IDs must exist. + +1. **Port existence** Input ports must exist on the node class; output ports on the node's output model. + +1. **DAG constraint** Build a *flat* `DiGraph` (no runtime expansion) and assert acyclicity. + +1. **Type compatibility** `get_output_field_type` vs `get_input_field_type` and `are_connection_types_compatible`. + +1. **Iterator / collector structure** Enforce special rules: + + - Iterator's input must be `collection`; its outgoing edges use `item`. + - Collector accepts many `item` inputs; outputs a single `collection`. + - Edge fan-in to a non-collector input is rejected. ### 3.3 Edge admission (`_validate_edge`) Checks a single prospective edge before insertion: -* Endpoints/ports exist. -* Destination port is not already occupied unless it's a collector `item`. -* Adding the edge to the flat DAG must keep it acyclic. -* Iterator/collector constraints re-checked when the edge creates relevant patterns. +- Endpoints/ports exist. +- Destination port is not already occupied unless it's a collector `item`. +- Adding the edge to the flat DAG must keep it acyclic. +- Iterator/collector constraints re-checked when the edge creates relevant patterns. ### 3.4 Topology utilities -* `nx_graph()` - DiGraph of declared nodes and edges. -* `nx_graph_with_data()` - includes node/edge attributes. -* `nx_graph_flat()` - "flattened" DAG (still author-time; no runtime copies). - Used in validation and in `_prepare()` during execution planning. +- `nx_graph()` - DiGraph of declared nodes and edges. +- `nx_graph_flat()` - "flattened" DAG (still author-time; no runtime copies). Used in validation and in `_prepare()` + during execution planning. ### 3.5 Mutation helpers -* `add_node`, `update_node` (preserve edges, rewrite endpoints if id changes), `delete_node`. -* `add_edge`, `delete_edge` (with validation). +- `add_node`, `update_node` (preserve edges, rewrite endpoints if id changes), `delete_node`. +- `add_edge`, `delete_edge` (with validation). ## 4) GraphExecutionState (runtime) -Holds the state for a single run. Keeps the source graph intact; materializes a separate execution graph. +Holds the state for a single run. Keeps the source graph intact and materializes a separate execution graph. +`GraphExecutionState` is still the public runtime entry point, but most execution behavior is now delegated to a small +set of internal helper classes. + +The source graph is treated as stable during normal execution, but the runtime object still exposes guarded graph +mutation helpers. Those helpers reject changes once the affected nodes have already been prepared or executed. ### 4.1 Data -* `graph: Graph` - immutable source during a run. -* `execution_graph: Graph` - materialized runtime nodes/edges. -* `executed: set[str]`, `executed_history: list[str]`. -* `results: dict[str, AnyInvocationOutput]`, `errors: dict[str, str]`. -* `prepared_source_mapping: dict[str, str]` - exec id → source id. -* `source_prepared_mapping: dict[str, set[str]]` - source id → exec ids. -* `indegree: dict[str, int]` - unmet inputs per exec node. -* **Ready queues grouped by class** (private attrs): - `_ready_queues: dict[class_name, deque[str]]`, `_active_class: Optional[str]`. Optional `ready_order: list[str]` to - prioritize classes. +- `graph: Graph` - source graph for the run; treated as stable during normal execution. +- `execution_graph: Graph` - materialized runtime nodes/edges. +- `executed: set[str]`, `executed_history: list[str]`. +- `results: dict[str, AnyInvocationOutput]`, `errors: dict[str, str]`. +- `prepared_source_mapping: dict[str, str]` - exec id -> source id. +- `source_prepared_mapping: dict[str, set[str]]` - source id -> exec ids. +- `indegree: dict[str, int]` - unmet inputs per exec node. +- Prepared exec metadata caches: + - source node id + - iteration path + - runtime state such as pending, ready, executed, or skipped +- **Ready queues grouped by class** (private attrs): `_ready_queues: dict[class_name, deque[str]]`, + `_active_class: Optional[str]`. Optional `ready_order: list[str]` to prioritize classes. ### 4.2 Core methods -* `next()` - Returns the next ready exec node. If none, calls `_prepare()` to materialize more, then retries. Before returning a - node, `_prepare_inputs()` deep-copies inbound values into the node fields. -* `complete(node_id, output)` - Record result; mark exec node executed; if all exec copies of the same **source** are done, mark the source executed. - For each outgoing exec edge, decrement child indegree and enqueue when it reaches zero. +- `next()` Returns the next ready exec node. If none are ready, it asks the materializer to expand more source nodes and + then retries. Before returning a node, the runtime helper deep-copies inbound values into the node fields. +- `complete(node_id, output)` Records the result, marks the exec node executed, marks the source node executed once all + of its prepared exec copies are done, then decrements downstream indegrees and enqueues newly ready nodes. + +### 4.3 Runtime helper classes + +`GraphExecutionState` now delegates most runtime behavior to internal helpers: + +- `_PreparedExecRegistry` Owns the relationship between source graph nodes and prepared execution graph nodes, plus + cached metadata such as iteration path and runtime state. +- `_ExecutionMaterializer` Expands source graph nodes into concrete execution graph nodes when the scheduler runs out of + ready work. +- `_ExecutionScheduler` Owns indegree transitions, ready queues, class batching, and downstream release on completion. +- `_ExecutionRuntime` Owns iteration-path lookup and input hydration for prepared exec nodes. +- `_IfBranchScheduler` Applies lazy `If` semantics by deferring branch-local work until the condition is known, then + releasing the selected branch and skipping the unselected branch. + +### 4.4 Preparation (`_prepare()`) -### 4.3 Preparation (`_prepare()`) +- Build a flat DAG from the **source** graph. -* Build a flat DAG from the **source** graph. -* Choose the **next source node** in topological order that: +- Choose the **next source node** in topological order that: 1. has not been prepared, - 2. if it is an iterator, *its inputs are already executed*, - 3. it has *no unexecuted iterator ancestors*. -* If the node is a **CollectInvocation**: collapse all prepared parents into one mapping and create **one** exec node. -* Otherwise: compute all combinations of prepared iterator ancestors. For each combination, pick the matching prepared parent per upstream and create **one** exec node. -* For each new exec node: + 1. if it is an iterator, *its inputs are already executed*, + 1. it has *no unexecuted iterator ancestors*. - * Deep-copy the source node; assign a fresh ID (and `index` for iterators). - * Wire edges from chosen prepared parents. - * Set `indegree = number of unmet inputs` (i.e., parents not yet executed). - * If `indegree == 0`, enqueue into its class queue. +- If the node is a **CollectInvocation**: collapse all prepared parents into one mapping and create **one** exec node. -### 4.4 Readiness and batching +- Otherwise: compute all combinations of prepared iterator ancestors. For each combination, choose the prepared parent + for each upstream by matching iterator ancestry, then create **one** exec node. -* `_enqueue_if_ready(nid)` enqueues by class name only when `indegree == 0` and not executed. -* `_get_next_node()` drains the `_active_class` queue FIFO; when empty, selects the next nonempty class queue (by `ready_order` if set, else alphabetical), and continues. Optional fairness knobs can limit batch size per class; default is drain fully. +- For each new exec node: -#### 4.4.1 Indegree (what it is and how it's used) + - Deep-copy the source node; assign a fresh ID (and `index` for iterators). + - Wire edges from chosen prepared parents. + - Set `indegree = number of unmet inputs` (i.e., parents not yet executed). + - Try to resolve any `If`-specific scheduling state. + - If the node is ready and not deferred by an unresolved `If`, enqueue it into its class queue. + +### 4.5 Readiness and batching + +- `_enqueue_if_ready(nid)` enqueues by class name only when `indegree == 0`, the node has not already executed, and the + node is not deferred by an unresolved `If`. +- `_get_next_node()` drains the `_active_class` queue FIFO; when empty, selects the next nonempty class queue (by + `ready_order` if set, else alphabetical), and continues. Optional fairness knobs can limit batch size per class; + default is drain fully. + +#### 4.5.1 Indegree (what it is and how it's used) **Indegree** is the number of incoming edges to a node in the execution graph that are still unmet. In this engine: -* For every materialized exec node, `indegree[node]` equals the count of its prerequisite parents that have **not** finished yet. -* A node is "ready" exactly when `indegree[node] == 0`; only then is it enqueued. -* When a node completes, the scheduler decrements `indegree[child]` for each outgoing edge. Any child that reaches 0 is enqueued. -Example: edges `A→C`, `B→C`, `C→D`. Start: `A:0, B:0, C:2, D:1`. Run `A` → `C:1`. Run `B` → `C:0` → enqueue `C`. Run `C` -→ `D:0` → enqueue `D`. Run `D` → done. +- For every materialized exec node, `indegree[node]` equals the count of its prerequisite parents that have **not** + finished yet. +- A node is "ready" exactly when `indegree[node] == 0`; only then is it enqueued. +- When a node completes, the scheduler decrements `indegree[child]` for each outgoing edge. Any child that reaches 0 is + enqueued. + +Example: edges `A->C`, `B->C`, `C->D`. Start: `A:0, B:0, C:2, D:1`. Run `A` -> `C:1`. Run `B` -> `C:0` -> enqueue `C`. +Run `C` -> `D:0` -> enqueue `D`. Run `D` -> done. + +### 4.6 Input hydration (`_prepare_inputs()`) + +- For **CollectInvocation**: gather all incoming `item` values into `collection`, sorting inputs by iteration path so + collected results are stable across expanded iterations. Incoming `collection` values are merged first, then incoming + `item` values are appended. +- For **IfInvocation**: hydrate only `condition` and the selected branch input. +- For all others: deep-copy each incoming edge's value into the destination field. This prevents cross-node mutation + through shared references. -### 4.5 Input hydration (`_prepare_inputs()`) +### 4.7 Lazy `If` semantics -* For **CollectInvocation**: gather all incoming `item` values into `collection`. -* For all others: deep-copy each incoming edge's value into the destination field. - This prevents cross-node mutation through shared references. +`IfInvocation` now acts as a lazy branch boundary rather than a simple value multiplexer. + +- The `condition` input must resolve first. +- Nodes that are exclusive to the true or false branch can remain deferred even when their indegree is zero. +- Once the prepared `If` node resolves its condition: + - the selected branch is released + - the unselected branch is marked skipped + - branch-exclusive ancestors of the unselected branch are never executed +- Shared ancestors still execute if they are required by the selected branch or by any other live path in the graph. + +This behavior is implemented in the runtime scheduler, not in the invocation body itself. ## 5) Traversal Summary 1. Author builds a valid **Graph**. -2. Create **GraphExecutionState** with that graph. -3. Loop: - * `node = state.next()` → may trigger `_prepare()` expansion. - * Execute node externally → `output`. - * `state.complete(node.id, output)` → updates indegrees and queues. -4. Finish when `next()` returns `None`. +1. Create **GraphExecutionState** with that graph. + +1. Loop: + + - `node = state.next()` -> may trigger `_prepare()` expansion. + - Execute node externally -> `output`. + - `state.complete(node.id, output)` -> updates indegrees, `If` state, and ready queues. + +1. Finish when `next()` returns `None`. -The source graph is never mutated; all expansion occurs in `execution_graph` with traceability back to source nodes. +In normal execution, all runtime expansion occurs in `execution_graph` with traceability back to source nodes. ## 6) Invariants -* Source **Graph** remains a DAG and type-consistent. -* `execution_graph` remains a DAG. -* Nodes are enqueued only when `indegree == 0`. -* `results` and `errors` are keyed by **exec node id**. -* Collectors only aggregate `item` inputs; other inputs behave one-to-one. +- Source **Graph** remains a DAG and type-consistent. +- `execution_graph` remains a DAG. +- Nodes are enqueued only when `indegree == 0` and they are not deferred by an unresolved `If`. +- `results` and `errors` are keyed by **exec node id**. +- Collectors aggregate `item` inputs and may also merge incoming `collection` inputs during runtime hydration. +- Branch-exclusive nodes behind an unselected `If` branch are skipped, not failed. ## 7) Extensibility -* **New node types**: implement as Pydantic models with typed fields and outputs. Register per your invocation system; this file accepts them as `AnyInvocation`. -* **Scheduling policy**: adjust `ready_order` to batch by class; add a batch cap for fairness without changing complexity. -* **Dynamic behaviors** (future): can be added in `GraphExecutionState` by creating exec nodes and edges at `complete()` time, as long as the DAG invariant holds. +- **New node types**: implement as Pydantic models with typed fields and outputs. Register per your invocation system; + this file accepts them as `AnyInvocation`. +- **Scheduling policy**: adjust `ready_order` to batch by class; add a batch cap for fairness without changing + complexity. +- **Dynamic behaviors** (future): can be added in `GraphExecutionState` by creating exec nodes and edges at `complete()` + time, as long as the DAG invariant holds. ## 8) Error Model (selected) -* `DuplicateNodeIdError`, `NodeAlreadyInGraphError` -* `NodeNotFoundError`, `NodeFieldNotFoundError` -* `InvalidEdgeError`, `CyclicalGraphError` -* `NodeInputError` (raised when preparing inputs for execution) +- `DuplicateNodeIdError`, `NodeAlreadyInGraphError` +- `NodeNotFoundError`, `NodeFieldNotFoundError` +- `InvalidEdgeError`, `CyclicalGraphError` +- `NodeInputError` (raised when preparing inputs for execution) Messages favor short, precise diagnostics (node id, field, and failing condition). ## 9) Rationale -* **Two-graph approach** isolates authoring from execution expansion and keeps validation simple. -* **Indegree + queues** gives O(1) scheduling decisions with clear batching semantics. -* **Iterator/collector separation** keeps fan-out/fan-in explicit and testable. -* **Deep-copy hydration** avoids incidental aliasing bugs between nodes. +- **Two-graph approach** isolates authoring from execution expansion and keeps validation simple. +- **Indegree + queues** gives O(1) scheduling decisions with clear batching semantics. +- **Iterator/collector separation** keeps fan-out/fan-in explicit and testable. +- **Deep-copy hydration** avoids incidental aliasing bugs between nodes. diff --git a/invokeai/app/services/shared/graph.py b/invokeai/app/services/shared/graph.py index 3ea67a52912..24c1dd1fe4f 100644 --- a/invokeai/app/services/shared/graph.py +++ b/invokeai/app/services/shared/graph.py @@ -3,7 +3,8 @@ import copy import itertools from collections import deque -from typing import Any, Deque, Iterable, Optional, Type, TypeVar, Union, get_args, get_origin +from dataclasses import dataclass +from typing import Any, Deque, Iterable, Literal, Optional, Type, TypeVar, Union, get_args, get_origin import networkx as nx from pydantic import ( @@ -29,6 +30,7 @@ invocation_output, ) from invokeai.app.invocations.fields import Input, InputField, OutputField, UIType +from invokeai.app.invocations.logic import IfInvocation from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.util.misc import uuid_string @@ -63,6 +65,731 @@ def __str__(self): return f"{self.source.node_id}.{self.source.field} -> {self.destination.node_id}.{self.destination.field}" +PreparedExecState = Literal["pending", "ready", "executed", "skipped"] + + +@dataclass +class _PreparedExecNodeMetadata: + """Cached metadata for a materialized execution node.""" + + source_node_id: str + iteration_path: Optional[tuple[int, ...]] = None + state: PreparedExecState = "pending" + + +class _PreparedExecRegistry: + """Tracks prepared execution nodes and their relationship to source graph nodes.""" + + def __init__( + self, + prepared_source_mapping: dict[str, str], + source_prepared_mapping: dict[str, set[str]], + metadata: dict[str, _PreparedExecNodeMetadata], + ) -> None: + self._prepared_source_mapping = prepared_source_mapping + self._source_prepared_mapping = source_prepared_mapping + self._metadata = metadata + + def register(self, exec_node_id: str, source_node_id: str) -> None: + self._prepared_source_mapping[exec_node_id] = source_node_id + self._metadata[exec_node_id] = _PreparedExecNodeMetadata(source_node_id=source_node_id) + if source_node_id not in self._source_prepared_mapping: + self._source_prepared_mapping[source_node_id] = set() + self._source_prepared_mapping[source_node_id].add(exec_node_id) + + def get_metadata(self, exec_node_id: str) -> _PreparedExecNodeMetadata: + metadata = self._metadata.get(exec_node_id) + if metadata is None: + metadata = _PreparedExecNodeMetadata(source_node_id=self._prepared_source_mapping[exec_node_id]) + self._metadata[exec_node_id] = metadata + return metadata + + def get_source_node_id(self, exec_node_id: str) -> str: + metadata = self._metadata.get(exec_node_id) + if metadata is not None: + return metadata.source_node_id + return self._prepared_source_mapping[exec_node_id] + + def get_prepared_ids(self, source_node_id: str) -> set[str]: + return self._source_prepared_mapping.get(source_node_id, set()) + + def set_state(self, exec_node_id: str, state: PreparedExecState) -> None: + self.get_metadata(exec_node_id).state = state + + def get_iteration_path(self, exec_node_id: str) -> Optional[tuple[int, ...]]: + metadata = self._metadata.get(exec_node_id) + return metadata.iteration_path if metadata is not None else None + + def set_iteration_path(self, exec_node_id: str, iteration_path: tuple[int, ...]) -> None: + self.get_metadata(exec_node_id).iteration_path = iteration_path + + +class _IfBranchScheduler: + """Applies lazy `If` semantics by deferring, releasing, and skipping branch-local exec nodes.""" + + def __init__(self, state: "GraphExecutionState") -> None: + self._state = state + + def _get_branch_input_sources(self, if_node_id: str, branch_field: str) -> set[str]: + return {e.source.node_id for e in self._state.graph._get_input_edges(if_node_id, branch_field)} + + def _expand_with_ancestors(self, node_ids: set[str]) -> set[str]: + expanded = set(node_ids) + source_graph = self._state.graph.nx_graph_flat() + for node_id in list(expanded): + expanded.update(nx.ancestors(source_graph, node_id)) + return expanded + + def _node_outputs_stay_in_branch( + self, node_id: str, if_node_id: str, branch_field: str, branch_nodes: set[str] + ) -> bool: + output_edges = self._state.graph._get_output_edges(node_id) + return all( + edge.destination.node_id in branch_nodes + or (edge.destination.node_id == if_node_id and edge.destination.field == branch_field) + for edge in output_edges + ) + + def _prune_nonexclusive_branch_nodes( + self, if_node_id: str, branch_field: str, candidate_nodes: set[str] + ) -> set[str]: + exclusive_nodes = set(candidate_nodes) + changed = True + while changed: + changed = False + for node_id in list(exclusive_nodes): + if self._node_outputs_stay_in_branch(node_id, if_node_id, branch_field, exclusive_nodes): + continue + exclusive_nodes.remove(node_id) + changed = True + return exclusive_nodes + + def _get_matching_prepared_if_ids(self, if_node_id: str, iteration_path: tuple[int, ...]) -> list[str]: + prepared_if_ids = self._state._prepared_registry().get_prepared_ids(if_node_id) + return [pid for pid in prepared_if_ids if self._state._get_iteration_path(pid) == iteration_path] + + def _has_unresolved_matching_if(self, if_node_id: str, iteration_path: tuple[int, ...]) -> bool: + matching_prepared_if_ids = self._get_matching_prepared_if_ids(if_node_id, iteration_path) + if not matching_prepared_if_ids: + return True + return not all(pid in self._state._resolved_if_exec_branches for pid in matching_prepared_if_ids) + + def _apply_condition_inputs(self, exec_node_id: str, node: IfInvocation) -> bool: + condition_edges = self._state.execution_graph._get_input_edges(exec_node_id, "condition") + if any(edge.source.node_id not in self._state.executed for edge in condition_edges): + return False + + for edge in condition_edges: + setattr( + node, + edge.destination.field, + copydeep(getattr(self._state.results[edge.source.node_id], edge.source.field)), + ) + return True + + def _get_selected_branch_fields(self, node: IfInvocation) -> tuple[str, str]: + selected_field = "true_input" if node.condition else "false_input" + unselected_field = "false_input" if node.condition else "true_input" + return selected_field, unselected_field + + def _prune_unselected_if_inputs(self, exec_node_id: str, unselected_field: str) -> None: + for edge in self._state.execution_graph._get_input_edges(exec_node_id, unselected_field): + if edge.source.node_id in self._state.executed: + continue + if self._state.indegree[exec_node_id] == 0: + raise RuntimeError(f"indegree underflow for {exec_node_id} when pruning {unselected_field}") + self._state.indegree[exec_node_id] -= 1 + + def _apply_branch_resolution( + self, + exec_node_id: str, + iteration_path: tuple[int, ...], + exclusive_sources: dict[str, set[str]], + selected_field: str, + unselected_field: str, + ) -> None: + # This iterates over the stable prepared-source mapping while mutating per-exec runtime state such as ready + # queues, execution state, and prepared metadata. Branch resolution never adds or removes prepared exec nodes. + for prepared_id, prepared_source in self._state.prepared_source_mapping.items(): + if prepared_id in self._state.executed: + continue + if self._state._get_iteration_path(prepared_id) != iteration_path: + continue + if prepared_source in exclusive_sources[selected_field]: + self._state._enqueue_if_ready(prepared_id) + elif prepared_source in exclusive_sources[unselected_field]: + self.mark_exec_node_skipped(prepared_id) + + def get_branch_exclusive_sources(self, if_node_id: str) -> dict[str, set[str]]: + cached = self._state._if_branch_exclusive_sources.get(if_node_id) + if cached is not None: + return cached + + branch_sources: dict[str, set[str]] = {} + for branch_field in ("true_input", "false_input"): + direct_inputs = self._get_branch_input_sources(if_node_id, branch_field) + candidate_nodes = self._expand_with_ancestors(direct_inputs) + branch_sources[branch_field] = self._prune_nonexclusive_branch_nodes( + if_node_id, branch_field, candidate_nodes + ) + + self._state._if_branch_exclusive_sources[if_node_id] = branch_sources + return branch_sources + + def is_deferred_by_unresolved_if(self, exec_node_id: str) -> bool: + source_node_id = self._state._prepared_registry().get_source_node_id(exec_node_id) + iteration_path = self._state._get_iteration_path(exec_node_id) + + for source_if_id, source_if_node in self._state.graph.nodes.items(): + if not isinstance(source_if_node, IfInvocation): + continue + + branches = self.get_branch_exclusive_sources(source_if_id) + if source_node_id not in branches["true_input"] and source_node_id not in branches["false_input"]: + continue + + if self._has_unresolved_matching_if(source_if_id, iteration_path): + return True + return False + + def mark_exec_node_skipped(self, exec_node_id: str) -> None: + self._state._remove_from_ready_queues(exec_node_id) + self._state._set_prepared_exec_state(exec_node_id, "skipped") + self._state.executed.add(exec_node_id) + + registry = self._state._prepared_registry() + source_node_id = registry.get_source_node_id(exec_node_id) + prepared_nodes = registry.get_prepared_ids(source_node_id) + if all(n in self._state.executed for n in prepared_nodes): + if source_node_id not in self._state.executed: + self._state.executed.add(source_node_id) + self._state.executed_history.append(source_node_id) + + def try_resolve_if_node(self, exec_node_id: str) -> None: + if exec_node_id in self._state._resolved_if_exec_branches: + return + node = self._state.execution_graph.get_node(exec_node_id) + if not isinstance(node, IfInvocation): + return + + if not self._apply_condition_inputs(exec_node_id, node): + return + + selected_field, unselected_field = self._get_selected_branch_fields(node) + self._state._resolved_if_exec_branches[exec_node_id] = selected_field + + source_if_node_id = self._state._prepared_registry().get_source_node_id(exec_node_id) + exclusive_sources = self.get_branch_exclusive_sources(source_if_node_id) + + iteration_path = self._state._get_iteration_path(exec_node_id) + self._prune_unselected_if_inputs(exec_node_id, unselected_field) + self._apply_branch_resolution(exec_node_id, iteration_path, exclusive_sources, selected_field, unselected_field) + self._state._enqueue_if_ready(exec_node_id) + + +class _ExecutionMaterializer: + """Expands source-graph nodes into concrete execution-graph nodes for the current runtime state. + + `GraphExecutionState.next()` calls into this helper when no prepared exec node is ready. The materializer chooses + the next source node that can be expanded, creates the corresponding exec nodes in the execution graph, wires their + inputs, and initializes their scheduler state. + """ + + def __init__(self, state: "GraphExecutionState") -> None: + self._state = state + + def _get_iterator_iteration_count(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> int: + input_collection_edge = next(iter(self._state.graph._get_input_edges(node_id, COLLECTION_FIELD))) + input_collection_prepared_node_id = next( + prepared_id + for source_id, prepared_id in iteration_node_map + if source_id == input_collection_edge.source.node_id + ) + input_collection_output = self._state.results[input_collection_prepared_node_id] + input_collection = getattr(input_collection_output, input_collection_edge.source.field) + return len(input_collection) + + def _get_new_node_iterations( + self, node: BaseInvocation, node_id: str, iteration_node_map: list[tuple[str, str]] + ) -> list[int]: + if not isinstance(node, IterateInvocation): + return [-1] + + iteration_count = self._get_iterator_iteration_count(node_id, iteration_node_map) + if iteration_count == 0: + return [] + return list(range(iteration_count)) + + def _build_execution_edges(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> list[Edge]: + input_edges = self._state.graph._get_input_edges(node_id) + new_edges: list[Edge] = [] + for edge in input_edges: + matching_inputs = [ + prepared_id for source_id, prepared_id in iteration_node_map if source_id == edge.source.node_id + ] + for input_node_id in matching_inputs: + new_edges.append( + Edge( + source=EdgeConnection(node_id=input_node_id, field=edge.source.field), + destination=EdgeConnection(node_id="", field=edge.destination.field), + ) + ) + return new_edges + + def _create_execution_node_copy(self, node: BaseInvocation, node_id: str, iteration_index: int) -> BaseInvocation: + new_node = node.model_copy(deep=True) + new_node.id = uuid_string() + + if isinstance(new_node, IterateInvocation): + new_node.index = iteration_index + + self._state.execution_graph.add_node(new_node) + self._state._register_prepared_exec_node(new_node.id, node_id) + return new_node + + def _attach_execution_edges(self, exec_node_id: str, new_edges: list[Edge]) -> None: + for edge in new_edges: + self._state.execution_graph.add_edge( + Edge( + source=edge.source, + destination=EdgeConnection(node_id=exec_node_id, field=edge.destination.field), + ) + ) + + def _initialize_execution_node(self, exec_node_id: str) -> None: + inputs = self._state.execution_graph._get_input_edges(exec_node_id) + unmet = sum(1 for edge in inputs if edge.source.node_id not in self._state.executed) + self._state.indegree[exec_node_id] = unmet + self._state._try_resolve_if_node(exec_node_id) + self._state._enqueue_if_ready(exec_node_id) + + def _get_collect_iteration_mappings(self, parent_node_ids: list[str]) -> list[tuple[str, str]]: + all_iteration_mappings: list[tuple[str, str]] = [] + for source_node_id in parent_node_ids: + prepared_nodes = self._state.source_prepared_mapping[source_node_id] + all_iteration_mappings.extend((source_node_id, prepared_id) for prepared_id in prepared_nodes) + return all_iteration_mappings + + def _get_parent_iteration_mappings(self, next_node_id: str, graph: nx.DiGraph) -> list[list[tuple[str, str]]]: + parent_node_ids = [source_id for source_id, _ in graph.in_edges(next_node_id)] + iterator_graph = self.iterator_graph(graph) + iterator_nodes = self.get_node_iterators(next_node_id, iterator_graph) + iterator_nodes_prepared = [list(self._state.source_prepared_mapping[node_id]) for node_id in iterator_nodes] + iterator_node_prepared_combinations = list(itertools.product(*iterator_nodes_prepared)) + + execution_graph = self._state.execution_graph.nx_graph_flat() + prepared_parent_mappings = [ + [ + (node_id, self.get_iteration_node(node_id, graph, execution_graph, prepared_iterators)) + for node_id in parent_node_ids + ] + for prepared_iterators in iterator_node_prepared_combinations + ] + return [ + mapping + for mapping in prepared_parent_mappings + if all(prepared_id is not None for _, prepared_id in mapping) + ] + + def create_execution_node(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> list[str]: + """Prepares an iteration node and connects all edges, returning the new node id""" + + node = self._state.graph.get_node(node_id) + iteration_indexes = self._get_new_node_iterations(node, node_id, iteration_node_map) + if not iteration_indexes: + return [] + + new_edges = self._build_execution_edges(node_id, iteration_node_map) + new_nodes: list[str] = [] + for iteration_index in iteration_indexes: + new_node = self._create_execution_node_copy(node, node_id, iteration_index) + self._attach_execution_edges(new_node.id, new_edges) + self._initialize_execution_node(new_node.id) + new_nodes.append(new_node.id) + + return new_nodes + + def iterator_graph(self, base: Optional[nx.DiGraph] = None) -> nx.DiGraph: + """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node""" + g = base.copy() if base is not None else self._state.graph.nx_graph_flat() + collectors = ( + n for n in self._state.graph.nodes if isinstance(self._state.graph.get_node(n), CollectInvocation) + ) + for c in collectors: + g.remove_edges_from(list(g.in_edges(c))) + return g + + def get_node_iterators(self, node_id: str, it_graph: Optional[nx.DiGraph] = None) -> list[str]: + g = it_graph or self.iterator_graph() + return [n for n in nx.ancestors(g, node_id) if isinstance(self._state.graph.get_node(n), IterateInvocation)] + + def _get_prepared_nodes_for_source(self, source_node_id: str) -> set[str]: + return self._state.source_prepared_mapping[source_node_id] + + def _get_parent_iterator_exec_nodes( + self, source_node_id: str, graph: nx.DiGraph, prepared_iterator_nodes: list[str] + ) -> list[tuple[str, str]]: + iterator_source_node_mapping = [ + (prepared_exec_node_id, self._state.prepared_source_mapping[prepared_exec_node_id]) + for prepared_exec_node_id in prepared_iterator_nodes + ] + return [ + iterator_mapping + for iterator_mapping in iterator_source_node_mapping + if nx.has_path(graph, iterator_mapping[1], source_node_id) + ] + + def _matches_parent_iterators( + self, candidate_exec_node_id: str, parent_iterators: list[tuple[str, str]], execution_graph: nx.DiGraph + ) -> bool: + return all( + nx.has_path(execution_graph, parent_iterator_exec_id, candidate_exec_node_id) + for parent_iterator_exec_id, _ in parent_iterators + ) + + def _get_direct_prepared_iterator_match( + self, + prepared_nodes: set[str], + prepared_iterator_nodes: list[str], + parent_iterators: list[tuple[str, str]], + execution_graph: nx.DiGraph, + ) -> Optional[str]: + prepared_iterator = next((node_id for node_id in prepared_nodes if node_id in prepared_iterator_nodes), None) + if prepared_iterator is None: + return None + if self._matches_parent_iterators(prepared_iterator, parent_iterators, execution_graph): + return prepared_iterator + return None + + def _find_prepared_node_matching_iterators( + self, prepared_nodes: set[str], parent_iterators: list[tuple[str, str]], execution_graph: nx.DiGraph + ) -> Optional[str]: + return next( + ( + node_id + for node_id in prepared_nodes + if self._matches_parent_iterators(node_id, parent_iterators, execution_graph) + ), + None, + ) + + def get_iteration_node( + self, + source_node_id: str, + graph: nx.DiGraph, + execution_graph: nx.DiGraph, + prepared_iterator_nodes: list[str], + ) -> Optional[str]: + prepared_nodes = self._get_prepared_nodes_for_source(source_node_id) + if len(prepared_nodes) == 1: + return next(iter(prepared_nodes)) + + parent_iterators = self._get_parent_iterator_exec_nodes(source_node_id, graph, prepared_iterator_nodes) + + direct_iterator_match = self._get_direct_prepared_iterator_match( + prepared_nodes, prepared_iterator_nodes, parent_iterators, execution_graph + ) + if direct_iterator_match is not None: + return direct_iterator_match + + return self._find_prepared_node_matching_iterators(prepared_nodes, parent_iterators, execution_graph) + + def prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]: + g = base_g or self._state.graph.nx_graph_flat() + next_node_id = next( + ( + node_id + for node_id in nx.topological_sort(g) + if node_id not in self._state.source_prepared_mapping + and ( + not isinstance(self._state.graph.get_node(node_id), IterateInvocation) + or all(source_id in self._state.executed for source_id, _ in g.in_edges(node_id)) + ) + and not any( + isinstance(self._state.graph.get_node(ancestor_id), IterateInvocation) + and ancestor_id not in self._state.executed + for ancestor_id in nx.ancestors(g, node_id) + ) + ), + None, + ) + + if next_node_id is None: + return None + + next_node = self._state.graph.get_node(next_node_id) + new_node_ids: list[str] = [] + + if isinstance(next_node, CollectInvocation): + next_node_parents = [source_id for source_id, _ in g.in_edges(next_node_id)] + create_results = self.create_execution_node( + next_node_id, self._get_collect_iteration_mappings(next_node_parents) + ) + if create_results is not None: + new_node_ids.extend(create_results) + else: + for iteration_mappings in self._get_parent_iteration_mappings(next_node_id, g): + create_results = self.create_execution_node(next_node_id, iteration_mappings) + if create_results is not None: + new_node_ids.extend(create_results) + + return next(iter(new_node_ids), None) + + +class _ExecutionScheduler: + """Owns ready-queue ordering and indegree-driven execution transitions.""" + + def __init__(self, state: "GraphExecutionState") -> None: + self._state = state + + def _validate_exec_node_ready_state(self, exec_node_id: str) -> None: + if exec_node_id not in self._state.execution_graph.nodes: + raise KeyError(f"exec node {exec_node_id} missing from execution_graph") + if exec_node_id not in self._state.indegree: + raise KeyError(f"indegree missing for exec node {exec_node_id}") + + def _should_skip_ready_enqueue(self, exec_node_id: str) -> bool: + return ( + self._state.indegree[exec_node_id] != 0 + or exec_node_id in self._state.executed + or self._state._is_deferred_by_unresolved_if(exec_node_id) + ) + + def _get_ready_queue(self, exec_node_id: str) -> Deque[str]: + node_obj = self._state.execution_graph.nodes[exec_node_id] + return self.queue_for(self._state._type_key(node_obj)) + + def _insert_ready_node(self, queue: Deque[str], exec_node_id: str) -> None: + exec_node_path = self._state._get_iteration_path(exec_node_id) + for i, existing in enumerate(queue): + if self._state._get_iteration_path(existing) > exec_node_path: + queue.insert(i, exec_node_id) + return + queue.append(exec_node_id) + + def _record_completed_node(self, exec_node_id: str, output: BaseInvocationOutput) -> None: + self._state._set_prepared_exec_state(exec_node_id, "executed") + self._state.executed.add(exec_node_id) + self._state.results[exec_node_id] = output + + def _mark_source_node_complete(self, exec_node_id: str) -> None: + registry = self._state._prepared_registry() + source_node_id = registry.get_source_node_id(exec_node_id) + prepared_nodes = registry.get_prepared_ids(source_node_id) + if all(node_id in self._state.executed for node_id in prepared_nodes): + self._state.executed.add(source_node_id) + self._state.executed_history.append(source_node_id) + + def _decrement_child_indegree(self, child_exec_node_id: str, parent_exec_node_id: str) -> None: + if child_exec_node_id not in self._state.indegree: + raise KeyError(f"indegree missing for exec node {child_exec_node_id}") + if self._state.indegree[child_exec_node_id] == 0: + raise RuntimeError(f"indegree underflow for {child_exec_node_id} from parent {parent_exec_node_id}") + self._state.indegree[child_exec_node_id] -= 1 + + def _release_downstream_nodes(self, exec_node_id: str) -> None: + for edge in self._state.execution_graph._get_output_edges(exec_node_id): + child = edge.destination.node_id + self._decrement_child_indegree(child, exec_node_id) + self._state._try_resolve_if_node(child) + if self._state.indegree[child] == 0: + self.enqueue_if_ready(child) + + def queue_for(self, cls_name: str) -> Deque[str]: + q = self._state._ready_queues.get(cls_name) + if q is None: + q = deque() + self._state._ready_queues[cls_name] = q + return q + + def remove_from_ready_queues(self, exec_node_id: str) -> None: + for q in self._state._ready_queues.values(): + try: + q.remove(exec_node_id) + except ValueError: + continue + + def enqueue_if_ready(self, exec_node_id: str) -> None: + """Push exec_node_id to its class queue if unmet inputs == 0.""" + self._validate_exec_node_ready_state(exec_node_id) + if self._should_skip_ready_enqueue(exec_node_id): + return + queue = self._get_ready_queue(exec_node_id) + if exec_node_id in queue: + return + self._state._set_prepared_exec_state(exec_node_id, "ready") + self._insert_ready_node(queue, exec_node_id) + + def get_next_node(self) -> Optional[BaseInvocation]: + """Gets the next ready node: FIFO within class, drain class before switching.""" + while True: + if self._state._active_class: + q = self._state._ready_queues.get(self._state._active_class) + while q: + exec_node_id = q.popleft() + if exec_node_id not in self._state.executed: + return self._state.execution_graph.nodes[exec_node_id] + self._state._active_class = None + continue + + seen = set(self._state.ready_order) + next_class = next( + (cls_name for cls_name in self._state.ready_order if self._state._ready_queues.get(cls_name)), + None, + ) + if next_class is None: + next_class = next( + ( + cls_name + for cls_name in sorted(k for k in self._state._ready_queues.keys() if k not in seen) + if self._state._ready_queues[cls_name] + ), + None, + ) + if next_class is None: + return None + + self._state._active_class = next_class + + def complete(self, exec_node_id: str, output: BaseInvocationOutput) -> None: + if exec_node_id not in self._state.execution_graph.nodes: + return + + self._record_completed_node(exec_node_id, output) + self._mark_source_node_complete(exec_node_id) + self._release_downstream_nodes(exec_node_id) + + +class _ExecutionRuntime: + """Provides runtime-only helpers such as iteration-path lookup and input hydration.""" + + def __init__(self, state: "GraphExecutionState") -> None: + self._state = state + + def _get_cached_iteration_path(self, exec_node_id: str) -> Optional[tuple[int, ...]]: + registry = self._state._prepared_registry() + metadata_iteration_path = registry.get_iteration_path(exec_node_id) + if metadata_iteration_path is not None: + return metadata_iteration_path + + return self._state._iteration_path_cache.get(exec_node_id) + + def _get_iteration_source_node_id(self, exec_node_id: str) -> Optional[str]: + if exec_node_id not in self._state.prepared_source_mapping: + return None + return self._state._prepared_registry().get_source_node_id(exec_node_id) + + def _get_ordered_iterator_sources(self, source_node_id: str) -> list[str]: + iterator_graph = self._state._iterator_graph(self._state.graph.nx_graph()) + iterator_sources = [ + node_id + for node_id in nx.ancestors(iterator_graph, source_node_id) + if isinstance(self._state.graph.get_node(node_id), IterateInvocation) + ] + + topo = list(nx.topological_sort(iterator_graph)) + topo_index = {node_id: i for i, node_id in enumerate(topo)} + iterator_sources.sort(key=lambda node_id: topo_index.get(node_id, 0)) + return iterator_sources + + def _get_iterator_exec_id( + self, iterator_source_id: str, exec_node_id: str, execution_graph: nx.DiGraph + ) -> Optional[str]: + prepared = self._state.source_prepared_mapping.get(iterator_source_id) + if not prepared: + return None + return next((pid for pid in prepared if nx.has_path(execution_graph, pid, exec_node_id)), None) + + def _build_iteration_path(self, exec_node_id: str, source_node_id: str) -> tuple[int, ...]: + iterator_sources = self._get_ordered_iterator_sources(source_node_id) + execution_graph = self._state.execution_graph.nx_graph() + path: list[int] = [] + for iterator_source_id in iterator_sources: + iterator_exec_id = self._get_iterator_exec_id(iterator_source_id, exec_node_id, execution_graph) + if iterator_exec_id is None: + continue + iterator_node = self._state.execution_graph.nodes.get(iterator_exec_id) + if isinstance(iterator_node, IterateInvocation): + path.append(iterator_node.index) + + node_obj = self._state.execution_graph.nodes.get(exec_node_id) + if isinstance(node_obj, IterateInvocation): + path.append(node_obj.index) + + return tuple(path) + + def _cache_iteration_path(self, exec_node_id: str, iteration_path: tuple[int, ...]) -> tuple[int, ...]: + self._state._iteration_path_cache[exec_node_id] = iteration_path + self._state._prepared_registry().set_iteration_path(exec_node_id, iteration_path) + return iteration_path + + def get_iteration_path(self, exec_node_id: str) -> tuple[int, ...]: + """Best-effort outer->inner iteration indices for an execution node, stopping at collectors.""" + cached = self._get_cached_iteration_path(exec_node_id) + if cached is not None: + return cached + + source_node_id = self._get_iteration_source_node_id(exec_node_id) + if source_node_id is None: + return self._cache_iteration_path(exec_node_id, ()) + + return self._cache_iteration_path(exec_node_id, self._build_iteration_path(exec_node_id, source_node_id)) + + def _sort_collect_input_edges(self, input_edges: list[Edge], field_name: str) -> list[Edge]: + matching_edges = [edge for edge in input_edges if edge.destination.field == field_name] + matching_edges.sort(key=lambda edge: (self.get_iteration_path(edge.source.node_id), edge.source.node_id)) + return matching_edges + + def _get_copied_result_value(self, edge: Edge) -> Any: + return copydeep(getattr(self._state.results[edge.source.node_id], edge.source.field)) + + def _build_collect_collection(self, input_edges: list[Edge]) -> list[Any]: + item_edges = self._sort_collect_input_edges(input_edges, ITEM_FIELD) + collection_edges = self._sort_collect_input_edges(input_edges, COLLECTION_FIELD) + + output_collection = [] + for edge in collection_edges: + source_value = self._get_copied_result_value(edge) + if isinstance(source_value, list): + output_collection.extend(source_value) + else: + output_collection.append(source_value) + output_collection.extend(self._get_copied_result_value(edge) for edge in item_edges) + return output_collection + + def _set_node_inputs( + self, node: BaseInvocation, input_edges: list[Edge], allowed_fields: Optional[set[str]] = None + ) -> None: + for edge in input_edges: + if allowed_fields is not None and edge.destination.field not in allowed_fields: + continue + setattr(node, edge.destination.field, self._get_copied_result_value(edge)) + + def _prepare_collect_inputs(self, node: "CollectInvocation", input_edges: list[Edge]) -> None: + node.collection = self._build_collect_collection(input_edges) + + def _prepare_if_inputs(self, node: IfInvocation, input_edges: list[Edge]) -> None: + selected_field = self._state._resolved_if_exec_branches.get(node.id) + allowed_fields = {"condition", selected_field} if selected_field is not None else {"condition"} + self._set_node_inputs(node, input_edges, allowed_fields) + + def _prepare_default_inputs(self, node: BaseInvocation, input_edges: list[Edge]) -> None: + self._set_node_inputs(node, input_edges) + + def prepare_inputs(self, node: BaseInvocation) -> None: + input_edges = self._state.execution_graph._get_input_edges(node.id) + + if isinstance(node, CollectInvocation): + self._prepare_collect_inputs(node, input_edges) + return + + if isinstance(node, IfInvocation): + self._prepare_if_inputs(node, input_edges) + return + + self._prepare_default_inputs(node, input_edges) + + def get_output_field_type(node: BaseInvocation, field: str) -> Any: # TODO(psyche): This is awkward - if field_info is None, it means the field is not defined in the output, which # really should raise. The consumers of this utility expect it to never raise, and return None instead. Fixing this @@ -364,6 +1091,8 @@ def __get_pydantic_json_schema__( class Graph(BaseModel): + """A validated invocation graph made of nodes and typed edges.""" + id: str = Field(description="The id of this graph", default_factory=uuid_string) # TODO: use a list (and never use dict in a BaseModel) because pydantic/fastapi hates me nodes: dict[str, AnyInvocation] = Field(description="The nodes in this graph", default_factory=dict) @@ -422,33 +1151,19 @@ def delete_edge(self, edge: Edge) -> None: except ValueError: pass - def validate_self(self) -> None: - """ - Validates the graph. - - Raises an exception if the graph is invalid: - - `DuplicateNodeIdError` - - `NodeIdMismatchError` - - `InvalidSubGraphError` - - `NodeNotFoundError` - - `NodeFieldNotFoundError` - - `CyclicalGraphError` - - `InvalidEdgeError` - """ - - # Validate that all node ids are unique + def _validate_unique_node_ids(self) -> None: node_ids = [n.id for n in self.nodes.values()] seen = set() duplicate_node_ids = {nid for nid in node_ids if (nid in seen) or seen.add(nid)} if duplicate_node_ids: raise DuplicateNodeIdError(f"Node ids must be unique, found duplicates {duplicate_node_ids}") - # Validate that all node ids match the keys in the nodes dict - for k, v in self.nodes.items(): - if k != v.id: - raise NodeIdMismatchError(f"Node ids must match, got {k} and {v.id}") + def _validate_node_id_mapping(self) -> None: + for node_dict_id, node in self.nodes.items(): + if node_dict_id != node.id: + raise NodeIdMismatchError(f"Node ids must match, got {node_dict_id} and {node.id}") - # Validate that all edges match nodes and fields in the graph + def _validate_edge_nodes_and_fields(self) -> None: for edge in self.edges: source_node = self.nodes.get(edge.source.node_id, None) if source_node is None: @@ -458,24 +1173,22 @@ def validate_self(self) -> None: if destination_node is None: raise NodeNotFoundError(f"Edge destination node {edge.destination.node_id} does not exist in the graph") - # output fields are not on the node object directly, they are on the output type if edge.source.field not in source_node.get_output_annotation().model_fields: raise NodeFieldNotFoundError( f"Edge source field {edge.source.field} does not exist in node {edge.source.node_id}" ) - # input fields are on the node if edge.destination.field not in type(destination_node).model_fields: raise NodeFieldNotFoundError( f"Edge destination field {edge.destination.field} does not exist in node {edge.destination.node_id}" ) - # Validate there are no cycles - g = self.nx_graph_flat() - if not nx.is_directed_acyclic_graph(g): + def _validate_graph_is_acyclic(self) -> None: + graph = self.nx_graph_flat() + if not nx.is_directed_acyclic_graph(graph): raise CyclicalGraphError("Graph contains cycles") - # Validate all edge connections are valid + def _validate_edge_type_compatibility(self) -> None: for edge in self.edges: if not are_connections_compatible( self.get_node(edge.source.node_id), @@ -485,7 +1198,7 @@ def validate_self(self) -> None: ): raise InvalidEdgeError(f"Edge source and target types do not match ({edge})") - # Validate all iterators & collectors + def _validate_special_nodes(self) -> None: # TODO: may need to validate all iterators & collectors in subgraphs so edge connections in parent graphs will be available for node in self.nodes.values(): if isinstance(node, IterateInvocation): @@ -497,6 +1210,26 @@ def validate_self(self) -> None: if err is not None: raise InvalidEdgeError(f"Invalid collector node ({node.id}): {err}") + def validate_self(self) -> None: + """ + Validates the graph. + + Raises an exception if the graph is invalid: + - `DuplicateNodeIdError` + - `NodeIdMismatchError` + - `InvalidSubGraphError` + - `NodeNotFoundError` + - `NodeFieldNotFoundError` + - `CyclicalGraphError` + - `InvalidEdgeError` + """ + + self._validate_unique_node_ids() + self._validate_node_id_mapping() + self._validate_edge_nodes_and_fields() + self._validate_graph_is_acyclic() + self._validate_edge_type_compatibility() + self._validate_special_nodes() return None def is_valid(self) -> bool: @@ -528,56 +1261,56 @@ def _is_destination_field_list_of_Any(self, edge: Edge) -> bool: """Checks if the destination field for an edge is of type typing.Any""" return get_input_field_type(self.get_node(edge.destination.node_id), edge.destination.field) == list[Any] - def _validate_edge(self, edge: Edge): - """Validates that a new edge doesn't create a cycle in the graph""" - - # Validate that the nodes exist + def _get_edge_nodes(self, edge: Edge) -> tuple[BaseInvocation, BaseInvocation]: try: - from_node = self.get_node(edge.source.node_id) - to_node = self.get_node(edge.destination.node_id) + return self.get_node(edge.source.node_id), self.get_node(edge.destination.node_id) except NodeNotFoundError: raise InvalidEdgeError(f"One or both nodes don't exist ({edge})") - # Validate that an edge to this node+field doesn't already exist + def _validate_edge_destination_uniqueness(self, edge: Edge, destination_node: BaseInvocation) -> None: input_edges = self._get_input_edges(edge.destination.node_id, edge.destination.field) if len(input_edges) > 0 and ( - not isinstance(to_node, CollectInvocation) or edge.destination.field != ITEM_FIELD + not isinstance(destination_node, CollectInvocation) or edge.destination.field != ITEM_FIELD ): raise InvalidEdgeError(f"Edge already exists ({edge})") - # Validate that no cycles would be created - g = self.nx_graph_flat() - g.add_edge(edge.source.node_id, edge.destination.node_id) - if not nx.is_directed_acyclic_graph(g): + def _validate_edge_would_not_create_cycle(self, edge: Edge) -> None: + graph = self.nx_graph_flat() + graph.add_edge(edge.source.node_id, edge.destination.node_id) + if not nx.is_directed_acyclic_graph(graph): raise InvalidEdgeError(f"Edge creates a cycle in the graph ({edge})") - # Validate that the field types are compatible - if not are_connections_compatible(from_node, edge.source.field, to_node, edge.destination.field): + def _validate_edge_field_compatibility( + self, edge: Edge, source_node: BaseInvocation, destination_node: BaseInvocation + ) -> None: + if not are_connections_compatible(source_node, edge.source.field, destination_node, edge.destination.field): raise InvalidEdgeError(f"Field types are incompatible ({edge})") - # Validate if iterator output type matches iterator input type (if this edge results in both being set) - if isinstance(to_node, IterateInvocation) and edge.destination.field == COLLECTION_FIELD: + def _validate_iterator_edge_rules( + self, edge: Edge, source_node: BaseInvocation, destination_node: BaseInvocation + ) -> None: + if isinstance(destination_node, IterateInvocation) and edge.destination.field == COLLECTION_FIELD: err = self._is_iterator_connection_valid(edge.destination.node_id, new_input=edge.source) if err is not None: raise InvalidEdgeError(f"Iterator input type does not match iterator output type ({edge}): {err}") - # Validate if iterator input type matches output type (if this edge results in both being set) - if isinstance(from_node, IterateInvocation) and edge.source.field == ITEM_FIELD: + if isinstance(source_node, IterateInvocation) and edge.source.field == ITEM_FIELD: err = self._is_iterator_connection_valid(edge.source.node_id, new_output=edge.destination) if err is not None: raise InvalidEdgeError(f"Iterator output type does not match iterator input type ({edge}): {err}") - # Validate if collector input type matches output type (if this edge results in both being set) - if isinstance(to_node, CollectInvocation) and edge.destination.field in (ITEM_FIELD, COLLECTION_FIELD): + def _validate_collector_edge_rules( + self, edge: Edge, source_node: BaseInvocation, destination_node: BaseInvocation + ) -> None: + if isinstance(destination_node, CollectInvocation) and edge.destination.field in (ITEM_FIELD, COLLECTION_FIELD): err = self._is_collector_connection_valid( edge.destination.node_id, new_input=edge.source, new_input_field=edge.destination.field ) if err is not None: raise InvalidEdgeError(f"Collector output type does not match collector input type ({edge}): {err}") - # Validate if collector output type matches input type (if this edge results in both being set) - skip if the destination field is not Any or list[Any] if ( - isinstance(from_node, CollectInvocation) + isinstance(source_node, CollectInvocation) and edge.source.field == COLLECTION_FIELD and not self._is_destination_field_list_of_Any(edge) and not self._is_destination_field_Any(edge) @@ -586,6 +1319,15 @@ def _validate_edge(self, edge: Edge): if err is not None: raise InvalidEdgeError(f"Collector input type does not match collector output type ({edge}): {err}") + def _validate_edge(self, edge: Edge): + """Validates that a new edge doesn't create a cycle in the graph""" + source_node, destination_node = self._get_edge_nodes(edge) + self._validate_edge_destination_uniqueness(edge, destination_node) + self._validate_edge_would_not_create_cycle(edge) + self._validate_edge_field_compatibility(edge, source_node, destination_node) + self._validate_iterator_edge_rules(edge, source_node, destination_node) + self._validate_collector_edge_rules(edge, source_node, destination_node) + def has_node(self, node_id: str) -> bool: """Determines whether or not a node exists in the graph.""" try: @@ -676,36 +1418,59 @@ def _is_iterator_connection_valid( if new_output is not None: outputs.append(new_output) + return self._validate_iterator_connections(inputs, outputs) + + def _validate_iterator_connections(self, inputs: list[EdgeConnection], outputs: list[EdgeConnection]) -> str | None: + presence_error = self._validate_iterator_input_presence(inputs) + if presence_error is not None: + return presence_error + + input_node = self.get_node(inputs[0].node_id) + input_field_type = get_output_field_type(input_node, inputs[0].field) + output_field_types = self._get_iterator_output_field_types(outputs) + + input_type_error = self._validate_iterator_input_type(input_field_type) + if input_type_error is not None: + return input_type_error + + output_type_error = self._validate_iterator_output_types(input_field_type, output_field_types) + if output_type_error is not None: + return output_type_error + + return self._validate_iterator_collector_input(input_node, output_field_types) + + def _validate_iterator_input_presence(self, inputs: list[EdgeConnection]) -> str | None: if len(inputs) == 0: return "Iterator must have a collection input edge" - - # Only one input is allowed for iterators if len(inputs) > 1: return "Iterator may only have one input edge" + return None - input_node = self.get_node(inputs[0].node_id) - - # Get input and output fields (the fields linked to the iterator's input/output) - input_field_type = get_output_field_type(input_node, inputs[0].field) - output_field_types = [get_input_field_type(self.get_node(e.node_id), e.field) for e in outputs] + def _get_iterator_output_field_types(self, outputs: list[EdgeConnection]) -> list[Any]: + return [get_input_field_type(self.get_node(e.node_id), e.field) for e in outputs] - # Input type must be a list + def _validate_iterator_input_type(self, input_field_type: Any) -> str | None: if get_origin(input_field_type) is not list: return "Iterator input must be a collection" + return None - # Validate that all outputs match the input type + def _validate_iterator_output_types(self, input_field_type: Any, output_field_types: list[Any]) -> str | None: input_field_item_type = get_args(input_field_type)[0] - if not all((are_connection_types_compatible(input_field_item_type, t) for t in output_field_types)): + if not all(are_connection_types_compatible(input_field_item_type, t) for t in output_field_types): return "Iterator outputs must connect to an input with a matching type" + return None - # Collector input type must match all iterator output types - if isinstance(input_node, CollectInvocation): - input_root_type = self._get_collector_input_root_type(input_node.id) - if input_root_type is None: - return "Iterator input collector must have at least one item or collection input edge" - if not all((are_connection_types_compatible(input_root_type, t) for t in output_field_types)): - return "Iterator collection type must match all iterator output types" + def _validate_iterator_collector_input( + self, input_node: BaseInvocation, output_field_types: list[Any] + ) -> str | None: + if not isinstance(input_node, CollectInvocation): + return None + input_root_type = self._get_collector_input_root_type(input_node.id) + if input_root_type is None: + return "Iterator input collector must have at least one item or collection input edge" + if not all(are_connection_types_compatible(input_root_type, t) for t in output_field_types): + return "Iterator collection type must match all iterator output types" return None def _resolve_collector_input_types(self, node_id: str, visited: Optional[set[str]] = None) -> set[Any]: @@ -733,6 +1498,13 @@ def _resolve_collector_input_types(self, node_id: str, visited: Optional[set[str return input_types + def _get_type_tree_root_types(self, input_types: set[Any]) -> list[Any]: + type_tree = nx.DiGraph() + type_tree.add_nodes_from(input_types) + type_tree.add_edges_from([e for e in itertools.permutations(input_types, 2) if issubclass(e[1], e[0])]) + type_degrees = type_tree.in_degree(type_tree.nodes) + return [t[0] for t in type_degrees if t[1] == 0] # type: ignore + def _get_collector_input_root_type(self, node_id: str) -> Any | None: input_types = self._resolve_collector_input_types(node_id) non_any_input_types = {t for t in input_types if t != Any} @@ -741,22 +1513,18 @@ def _get_collector_input_root_type(self, node_id: str) -> Any | None: if len(non_any_input_types) == 0: return None - type_tree = nx.DiGraph() - type_tree.add_nodes_from(non_any_input_types) - type_tree.add_edges_from([e for e in itertools.permutations(non_any_input_types, 2) if issubclass(e[1], e[0])]) - type_degrees = type_tree.in_degree(type_tree.nodes) - root_types = [t[0] for t in type_degrees if t[1] == 0] # type: ignore + root_types = self._get_type_tree_root_types(non_any_input_types) if len(root_types) != 1: return Any return root_types[0] - def _is_collector_connection_valid( + def _get_collector_connections( self, node_id: str, new_input: Optional[EdgeConnection] = None, new_input_field: Optional[str] = None, new_output: Optional[EdgeConnection] = None, - ) -> str | None: + ) -> tuple[list[EdgeConnection], list[EdgeConnection], list[EdgeConnection]]: item_inputs = [e.source for e in self._get_input_edges(node_id, ITEM_FIELD)] collection_inputs = [e.source for e in self._get_input_edges(node_id, COLLECTION_FIELD)] outputs = [e.destination for e in self._get_output_edges(node_id, COLLECTION_FIELD)] @@ -767,58 +1535,67 @@ def _is_collector_connection_valid( item_inputs.append(new_input) elif field == COLLECTION_FIELD: collection_inputs.append(new_input) + if new_output is not None: outputs.append(new_output) - if len(item_inputs) == 0 and len(collection_inputs) == 0: - return "Collector must have at least one item or collection input edge" + return item_inputs, collection_inputs, outputs - # Get input and output fields (the fields linked to the collector's input/output) + def _get_collector_port_types( + self, + item_inputs: list[EdgeConnection], + collection_inputs: list[EdgeConnection], + outputs: list[EdgeConnection], + ) -> tuple[list[Any], list[Any], list[Any]]: item_input_field_types = [get_output_field_type(self.get_node(e.node_id), e.field) for e in item_inputs] collection_input_field_types = [ get_output_field_type(self.get_node(e.node_id), e.field) for e in collection_inputs ] output_field_types = [get_input_field_type(self.get_node(e.node_id), e.field) for e in outputs] + return item_input_field_types, collection_input_field_types, output_field_types - if not all((is_list_or_contains_list(t) or is_any(t) for t in collection_input_field_types)): - return "Collector collection input must be a collection" - - # Validate that all inputs are derived from or match a single type - input_field_types = { + def _resolve_item_input_types(self, item_input_field_types: list[Any]) -> set[Any]: + return { resolved_type for input_field_type in item_input_field_types for resolved_type in ( [input_field_type] if get_origin(input_field_type) is None else get_args(input_field_type) ) if resolved_type != NoneType - } # Get unique types + } + def _resolve_collection_input_types( + self, collection_inputs: list[EdgeConnection], collection_input_field_types: list[Any] + ) -> set[Any]: + input_field_types: set[Any] = set() for input_conn, input_field_type in zip(collection_inputs, collection_input_field_types, strict=False): source_node = self.get_node(input_conn.node_id) if isinstance(source_node, CollectInvocation) and input_conn.field == COLLECTION_FIELD: input_field_types.update(self._resolve_collector_input_types(source_node.id)) continue input_field_types.update(extract_collection_item_types(input_field_type)) + return input_field_types + + def _validate_collector_collection_inputs(self, collection_input_field_types: list[Any]) -> str | None: + if not all((is_list_or_contains_list(t) or is_any(t) for t in collection_input_field_types)): + return "Collector collection input must be a collection" + return None + def _get_collector_input_root_type_from_resolved_types( + self, input_field_types: set[Any] + ) -> tuple[bool, Any | None]: non_any_input_field_types = {t for t in input_field_types if t != Any} - type_tree = nx.DiGraph() - type_tree.add_nodes_from(non_any_input_field_types) - type_tree.add_edges_from( - [e for e in itertools.permutations(non_any_input_field_types, 2) if issubclass(e[1], e[0])] - ) - type_degrees = type_tree.in_degree(type_tree.nodes) - root_types = [t[0] for t in type_degrees if t[1] == 0] # type: ignore + root_types = self._get_type_tree_root_types(non_any_input_field_types) if len(root_types) > 1: - return "Collector input collection items must be of a single type" - - # Get the input root type (if known) - input_root_type = root_types[0] if len(root_types) == 1 else None + return True, None + return False, root_types[0] if len(root_types) == 1 else None - # Verify that all outputs are lists + def _validate_collector_output_types( + self, output_field_types: list[Any], input_root_type: Any | None + ) -> str | None: if not all(is_list_or_contains_list(t) or is_any(t) for t in output_field_types): return "Collector output must connect to a collection input" - # Verify that all outputs match the input type (are a base class or the same class) if input_root_type is not None: if not all( is_any(t) @@ -830,8 +1607,11 @@ def _is_collector_connection_valid( elif any(not is_any(t) and get_args(t)[0] != Any for t in output_field_types): return "Collector outputs must connect to a collection input with a matching type" - # If this collector outputs to another collector's collection input, validate against the downstream - # collector's resolved input type (if available). + return None + + def _validate_downstream_collector_outputs( + self, outputs: list[EdgeConnection], input_root_type: Any | None + ) -> str | None: for output in outputs: output_node = self.get_node(output.node_id) if not isinstance(output_node, CollectInvocation) or output.field != COLLECTION_FIELD: @@ -845,6 +1625,46 @@ def _is_collector_connection_valid( continue if not are_connection_types_compatible(input_root_type, output_root_type): return "Collector outputs must connect to a collection input with a matching type" + return None + + def _is_collector_connection_valid( + self, + node_id: str, + new_input: Optional[EdgeConnection] = None, + new_input_field: Optional[str] = None, + new_output: Optional[EdgeConnection] = None, + ) -> str | None: + item_inputs, collection_inputs, outputs = self._get_collector_connections( + node_id, new_input=new_input, new_input_field=new_input_field, new_output=new_output + ) + + if len(item_inputs) == 0 and len(collection_inputs) == 0: + return "Collector must have at least one item or collection input edge" + + item_input_field_types, collection_input_field_types, output_field_types = self._get_collector_port_types( + item_inputs, collection_inputs, outputs + ) + + collection_input_error = self._validate_collector_collection_inputs(collection_input_field_types) + if collection_input_error is not None: + return collection_input_error + + input_field_types = self._resolve_item_input_types(item_input_field_types) + input_field_types.update(self._resolve_collection_input_types(collection_inputs, collection_input_field_types)) + + has_multiple_root_types, input_root_type = self._get_collector_input_root_type_from_resolved_types( + input_field_types + ) + if has_multiple_root_types: + return "Collector input collection items must be of a single type" + + output_type_error = self._validate_collector_output_types(output_field_types, input_root_type) + if output_type_error is not None: + return output_type_error + + downstream_output_error = self._validate_downstream_collector_outputs(outputs, input_root_type) + if downstream_output_error is not None: + return downstream_output_error return None @@ -869,7 +1689,7 @@ def nx_graph_flat(self, nx_graph: Optional[nx.DiGraph] = None) -> nx.DiGraph: class GraphExecutionState(BaseModel): - """Tracks the state of a graph execution""" + """Tracks source-graph expansion, execution progress, and runtime results.""" id: str = Field(description="The id of the execution state", default_factory=uuid_string) # TODO: Store a reference to the graph instead of the actual graph? @@ -913,62 +1733,70 @@ class GraphExecutionState(BaseModel): ready_order: list[str] = Field(default_factory=list) indegree: dict[str, int] = Field(default_factory=dict, description="Remaining unmet input count for exec nodes") _iteration_path_cache: dict[str, tuple[int, ...]] = PrivateAttr(default_factory=dict) + _if_branch_exclusive_sources: dict[str, dict[str, set[str]]] = PrivateAttr(default_factory=dict) + _resolved_if_exec_branches: dict[str, str] = PrivateAttr(default_factory=dict) + _prepared_exec_metadata: dict[str, _PreparedExecNodeMetadata] = PrivateAttr(default_factory=dict) + _prepared_exec_registry: Optional[_PreparedExecRegistry] = PrivateAttr(default=None) + _if_branch_scheduler: Optional[_IfBranchScheduler] = PrivateAttr(default=None) + _execution_materializer: Optional[_ExecutionMaterializer] = PrivateAttr(default=None) + _execution_scheduler: Optional[_ExecutionScheduler] = PrivateAttr(default=None) + _execution_runtime: Optional[_ExecutionRuntime] = PrivateAttr(default=None) def _type_key(self, node_obj: BaseInvocation) -> str: return node_obj.__class__.__name__ - def _get_iteration_path(self, exec_node_id: str) -> tuple[int, ...]: - """Best-effort outer->inner iteration indices for an execution node, stopping at collectors.""" - cached = self._iteration_path_cache.get(exec_node_id) - if cached is not None: - return cached + def _prepared_registry(self) -> _PreparedExecRegistry: + if self._prepared_exec_registry is None: + self._prepared_exec_registry = _PreparedExecRegistry( + prepared_source_mapping=self.prepared_source_mapping, + source_prepared_mapping=self.source_prepared_mapping, + metadata=self._prepared_exec_metadata, + ) + return self._prepared_exec_registry - # Only prepared execution nodes participate; otherwise treat as non-iterated. - source_node_id = self.prepared_source_mapping.get(exec_node_id) - if source_node_id is None: - self._iteration_path_cache[exec_node_id] = () - return () + def _if_scheduler(self) -> _IfBranchScheduler: + if self._if_branch_scheduler is None: + self._if_branch_scheduler = _IfBranchScheduler(self) + return self._if_branch_scheduler - # Source-graph iterator ancestry, with edges into collectors removed so iteration context doesn't leak. - it_g = self._iterator_graph(self.graph.nx_graph()) - iterator_sources = [ - n for n in nx.ancestors(it_g, source_node_id) if isinstance(self.graph.get_node(n), IterateInvocation) - ] + def _materializer(self) -> _ExecutionMaterializer: + if self._execution_materializer is None: + self._execution_materializer = _ExecutionMaterializer(self) + return self._execution_materializer - # Order iterators outer->inner via topo order of the iterator graph. - topo = list(nx.topological_sort(it_g)) - topo_index = {n: i for i, n in enumerate(topo)} - iterator_sources.sort(key=lambda n: topo_index.get(n, 0)) + def _scheduler(self) -> _ExecutionScheduler: + if self._execution_scheduler is None: + self._execution_scheduler = _ExecutionScheduler(self) + return self._execution_scheduler - # Map iterator source nodes to the prepared iterator exec nodes that are ancestors of exec_node_id. - eg = self.execution_graph.nx_graph() - path: list[int] = [] - for it_src in iterator_sources: - prepared = self.source_prepared_mapping.get(it_src) - if not prepared: - continue - it_exec = next((p for p in prepared if nx.has_path(eg, p, exec_node_id)), None) - if it_exec is None: - continue - it_node = self.execution_graph.nodes.get(it_exec) - if isinstance(it_node, IterateInvocation): - path.append(it_node.index) + def _runtime(self) -> _ExecutionRuntime: + if self._execution_runtime is None: + self._execution_runtime = _ExecutionRuntime(self) + return self._execution_runtime - # If this exec node is itself an iterator, include its own index as the innermost element. - node_obj = self.execution_graph.nodes.get(exec_node_id) - if isinstance(node_obj, IterateInvocation): - path.append(node_obj.index) + def _register_prepared_exec_node(self, exec_node_id: str, source_node_id: str) -> None: + self._prepared_registry().register(exec_node_id, source_node_id) + + def _get_prepared_exec_metadata(self, exec_node_id: str) -> _PreparedExecNodeMetadata: + return self._prepared_registry().get_metadata(exec_node_id) - result = tuple(path) - self._iteration_path_cache[exec_node_id] = result - return result + def _set_prepared_exec_state(self, exec_node_id: str, state: PreparedExecState) -> None: + self._prepared_registry().set_state(exec_node_id, state) + + def _get_iteration_path(self, exec_node_id: str) -> tuple[int, ...]: + return self._runtime().get_iteration_path(exec_node_id) def _queue_for(self, cls_name: str) -> Deque[str]: - q = self._ready_queues.get(cls_name) - if q is None: - q = deque() - self._ready_queues[cls_name] = q - return q + return self._scheduler().queue_for(cls_name) + + def _is_deferred_by_unresolved_if(self, exec_node_id: str) -> bool: + return self._if_scheduler().is_deferred_by_unresolved_if(exec_node_id) + + def _remove_from_ready_queues(self, exec_node_id: str) -> None: + self._scheduler().remove_from_ready_queues(exec_node_id) + + def _try_resolve_if_node(self, exec_node_id: str) -> None: + self._if_scheduler().try_resolve_if_node(exec_node_id) def set_ready_order(self, order: Iterable[Type[BaseInvocation] | str]) -> None: names: list[str] = [] @@ -977,24 +1805,19 @@ def set_ready_order(self, order: Iterable[Type[BaseInvocation] | str]) -> None: self.ready_order = names def _enqueue_if_ready(self, nid: str) -> None: - """Push nid to its class queue if unmet inputs == 0.""" - # Invariants: exec node exists and has an indegree entry - if nid not in self.execution_graph.nodes: - raise KeyError(f"exec node {nid} missing from execution_graph") - if nid not in self.indegree: - raise KeyError(f"indegree missing for exec node {nid}") - if self.indegree[nid] != 0 or nid in self.executed: - return - node_obj = self.execution_graph.nodes[nid] - q = self._queue_for(self._type_key(node_obj)) - nid_path = self._get_iteration_path(nid) - # Insert in lexicographic outer->inner order; preserve FIFO for equal paths. - for i, existing in enumerate(q): - if self._get_iteration_path(existing) > nid_path: - q.insert(i, nid) - break - else: - q.append(nid) + self._scheduler().enqueue_if_ready(nid) + + def _prepare_until_node_ready(self) -> Optional[BaseInvocation]: + base_graph = self.graph.nx_graph_flat() + prepared_id = self._materializer().prepare(base_graph) + next_node: Optional[BaseInvocation] = None + + while prepared_id is not None: + prepared_id = self._materializer().prepare(base_graph) + if next_node is None: + next_node = self._get_next_node() + + return next_node model_config = ConfigDict( json_schema_extra={ @@ -1027,14 +1850,7 @@ def next(self) -> Optional[BaseInvocation]: # If there are no prepared nodes, prepare some nodes next_node = self._get_next_node() if next_node is None: - base_g = self.graph.nx_graph_flat() - prepared_id = self._prepare(base_g) - - # Prepare as many nodes as we can - while prepared_id is not None: - prepared_id = self._prepare(base_g) - if next_node is None: - next_node = self._get_next_node() + next_node = self._prepare_until_node_ready() # Get values from edges if next_node is not None: @@ -1048,33 +1864,7 @@ def next(self) -> Optional[BaseInvocation]: def complete(self, node_id: str, output: BaseInvocationOutput) -> None: """Marks a node as complete""" - - if node_id not in self.execution_graph.nodes: - return # TODO: log error? - - # Mark node as executed - self.executed.add(node_id) - self.results[node_id] = output - - # Check if source node is complete (all prepared nodes are complete) - source_node = self.prepared_source_mapping[node_id] - prepared_nodes = self.source_prepared_mapping[source_node] - - if all(n in self.executed for n in prepared_nodes): - self.executed.add(source_node) - self.executed_history.append(source_node) - - # Decrement children indegree and enqueue when ready - for e in self.execution_graph._get_output_edges(node_id): - child = e.destination.node_id - if child not in self.indegree: - raise KeyError(f"indegree missing for exec node {child}") - # Only decrement if there's something to satisfy - if self.indegree[child] == 0: - raise RuntimeError(f"indegree underflow for {child} from parent {node_id}") - self.indegree[child] -= 1 - if self.indegree[child] == 0: - self._enqueue_if_ready(child) + self._scheduler().complete(node_id, output) def set_node_error(self, node_id: str, error: str): """Marks a node as errored""" @@ -1090,164 +1880,16 @@ def has_error(self) -> bool: return len(self.errors) > 0 def _create_execution_node(self, node_id: str, iteration_node_map: list[tuple[str, str]]) -> list[str]: - """Prepares an iteration node and connects all edges, returning the new node id""" - - node = self.graph.get_node(node_id) - - self_iteration_count = -1 - - # If this is an iterator node, we must create a copy for each iteration - if isinstance(node, IterateInvocation): - # Get input collection edge (should error if there are no inputs) - input_collection_edge = next(iter(self.graph._get_input_edges(node_id, COLLECTION_FIELD))) - input_collection_prepared_node_id = next( - n[1] for n in iteration_node_map if n[0] == input_collection_edge.source.node_id - ) - input_collection_prepared_node_output = self.results[input_collection_prepared_node_id] - input_collection = getattr(input_collection_prepared_node_output, input_collection_edge.source.field) - self_iteration_count = len(input_collection) - - new_nodes: list[str] = [] - if self_iteration_count == 0: - # TODO: should this raise a warning? It might just happen if an empty collection is input, and should be valid. - return new_nodes - - # Get all input edges - input_edges = self.graph._get_input_edges(node_id) - - # Create new edges for this iteration - # For collect nodes, this may contain multiple inputs to the same field - new_edges: list[Edge] = [] - for edge in input_edges: - for input_node_id in (n[1] for n in iteration_node_map if n[0] == edge.source.node_id): - new_edge = Edge( - source=EdgeConnection(node_id=input_node_id, field=edge.source.field), - destination=EdgeConnection(node_id="", field=edge.destination.field), - ) - new_edges.append(new_edge) - - # Create a new node (or one for each iteration of this iterator) - for i in range(self_iteration_count) if self_iteration_count > 0 else [-1]: - # Create a new node - new_node = node.model_copy(deep=True) - - # Create the node id (use a random uuid) - new_node.id = uuid_string() - - # Set the iteration index for iteration invocations - if isinstance(new_node, IterateInvocation): - new_node.index = i - - # Add to execution graph - self.execution_graph.add_node(new_node) - self.prepared_source_mapping[new_node.id] = node_id - if node_id not in self.source_prepared_mapping: - self.source_prepared_mapping[node_id] = set() - self.source_prepared_mapping[node_id].add(new_node.id) - - # Add new edges to execution graph - for edge in new_edges: - new_edge = Edge( - source=edge.source, - destination=EdgeConnection(node_id=new_node.id, field=edge.destination.field), - ) - self.execution_graph.add_edge(new_edge) - - # Initialize indegree as unmet inputs only and enqueue if ready - inputs = self.execution_graph._get_input_edges(new_node.id) - unmet = sum(1 for e in inputs if e.source.node_id not in self.executed) - self.indegree[new_node.id] = unmet - self._enqueue_if_ready(new_node.id) - - new_nodes.append(new_node.id) - - return new_nodes + return self._materializer().create_execution_node(node_id, iteration_node_map) def _iterator_graph(self, base: Optional[nx.DiGraph] = None) -> nx.DiGraph: - """Gets a DiGraph with edges to collectors removed so an ancestor search produces all active iterators for any node""" - g = base.copy() if base is not None else self.graph.nx_graph_flat() - collectors = (n for n in self.graph.nodes if isinstance(self.graph.get_node(n), CollectInvocation)) - for c in collectors: - g.remove_edges_from(list(g.in_edges(c))) - return g + return self._materializer().iterator_graph(base) def _get_node_iterators(self, node_id: str, it_graph: Optional[nx.DiGraph] = None) -> list[str]: - """Gets iterators for a node""" - g = it_graph or self._iterator_graph() - return [n for n in nx.ancestors(g, node_id) if isinstance(self.graph.get_node(n), IterateInvocation)] + return self._materializer().get_node_iterators(node_id, it_graph) def _prepare(self, base_g: Optional[nx.DiGraph] = None) -> Optional[str]: - # Get flattened source graph - g = base_g or self.graph.nx_graph_flat() - - # Find next node that: - # - was not already prepared - # - is not an iterate node whose inputs have not been executed - # - does not have an unexecuted iterate ancestor - sorted_nodes = nx.topological_sort(g) - - def unprepared(n: str) -> bool: - return n not in self.source_prepared_mapping - - def iter_inputs_ready(n: str) -> bool: - if not isinstance(self.graph.get_node(n), IterateInvocation): - return True - return all(u in self.executed for u, _ in g.in_edges(n)) - - def no_unexecuted_iter_ancestors(n: str) -> bool: - return not any( - isinstance(self.graph.get_node(a), IterateInvocation) and a not in self.executed - for a in nx.ancestors(g, n) - ) - - next_node_id = next( - (n for n in sorted_nodes if unprepared(n) and iter_inputs_ready(n) and no_unexecuted_iter_ancestors(n)), - None, - ) - - if next_node_id is None: - return None - - # Get all parents of the next node - next_node_parents = [u for u, _ in g.in_edges(next_node_id)] - - # Create execution nodes - next_node = self.graph.get_node(next_node_id) - new_node_ids = [] - if isinstance(next_node, CollectInvocation): - # Collapse all iterator input mappings and create a single execution node for the collect invocation - all_iteration_mappings = [] - for source_node_id in next_node_parents: - prepared_nodes = self.source_prepared_mapping[source_node_id] - all_iteration_mappings.extend([(source_node_id, p) for p in prepared_nodes]) - - create_results = self._create_execution_node(next_node_id, all_iteration_mappings) - if create_results is not None: - new_node_ids.extend(create_results) - else: # Iterators or normal nodes - # Get all iterator combinations for this node - # Will produce a list of lists of prepared iterator nodes, from which results can be iterated - it_g = self._iterator_graph(g) - iterator_nodes = self._get_node_iterators(next_node_id, it_g) - iterator_nodes_prepared = [list(self.source_prepared_mapping[n]) for n in iterator_nodes] - iterator_node_prepared_combinations = list(itertools.product(*iterator_nodes_prepared)) - - # Select the correct prepared parents for each iteration - # For every iterator, the parent must either not be a child of that iterator, or must match the prepared iteration for that iterator - eg = self.execution_graph.nx_graph_flat() - prepared_parent_mappings = [ - [(n, self._get_iteration_node(n, g, eg, it)) for n in next_node_parents] - for it in iterator_node_prepared_combinations - ] # type: ignore - prepared_parent_mappings = [m for m in prepared_parent_mappings if all(p[1] is not None for p in m)] - - # Create execution node for each iteration - for iteration_mappings in prepared_parent_mappings: - create_results = self._create_execution_node(next_node_id, iteration_mappings) # type: ignore - if create_results is not None: - new_node_ids.extend(create_results) - - return next(iter(new_node_ids), None) + return self._materializer().prepare(base_g) def _get_iteration_node( self, @@ -1256,82 +1898,13 @@ def _get_iteration_node( execution_graph: nx.DiGraph, prepared_iterator_nodes: list[str], ) -> Optional[str]: - """Gets the prepared version of the specified source node that matches every iteration specified""" - prepared_nodes = self.source_prepared_mapping[source_node_id] - if len(prepared_nodes) == 1: - return next(iter(prepared_nodes)) - - # Filter to only iterator nodes that are a parent of the specified node, in tuple format (prepared, source) - iterator_source_node_mapping = [(n, self.prepared_source_mapping[n]) for n in prepared_iterator_nodes] - parent_iterators = [itn for itn in iterator_source_node_mapping if nx.has_path(graph, itn[1], source_node_id)] - - # If the requested node is an iterator, only accept it if it is compatible with all parent iterators - prepared_iterator = next((n for n in prepared_nodes if n in prepared_iterator_nodes), None) - if prepared_iterator is not None: - if all(nx.has_path(execution_graph, pit[0], prepared_iterator) for pit in parent_iterators): - return prepared_iterator - return None - - return next( - (n for n in prepared_nodes if all(nx.has_path(execution_graph, pit[0], n) for pit in parent_iterators)), - None, - ) + return self._materializer().get_iteration_node(source_node_id, graph, execution_graph, prepared_iterator_nodes) def _get_next_node(self) -> Optional[BaseInvocation]: - """Gets the next ready node: FIFO within class, drain class before switching.""" - # 1) Continue draining the active class - if self._active_class: - q = self._ready_queues.get(self._active_class) - while q: - nid = q.popleft() - if nid not in self.executed: - return self.execution_graph.nodes[nid] - # emptied: release active class - self._active_class = None - - # 2) Pick next class by priority, then by class name - seen = set(self.ready_order) - for cls_name in self.ready_order: - q = self._ready_queues.get(cls_name) - if q: - self._active_class = cls_name - # recurse to drain newly set active class - return self._get_next_node() - for cls_name in sorted(k for k in self._ready_queues.keys() if k not in seen): - q = self._ready_queues[cls_name] - if q: - self._active_class = cls_name - return self._get_next_node() - return None + return self._scheduler().get_next_node() def _prepare_inputs(self, node: BaseInvocation): - input_edges = self.execution_graph._get_input_edges(node.id) - # Inputs must be deep-copied, else if a node mutates the object, other nodes that get the same input - # will see the mutation. - if isinstance(node, CollectInvocation): - item_edges = [e for e in input_edges if e.destination.field == ITEM_FIELD] - item_edges.sort(key=lambda e: (self._get_iteration_path(e.source.node_id), e.source.node_id)) - collection_edges = [e for e in input_edges if e.destination.field == COLLECTION_FIELD] - collection_edges.sort(key=lambda e: (self._get_iteration_path(e.source.node_id), e.source.node_id)) - - output_collection = [] - for edge in collection_edges: - source_value = copydeep(getattr(self.results[edge.source.node_id], edge.source.field)) - if isinstance(source_value, list): - output_collection.extend(source_value) - else: - output_collection.append(source_value) - output_collection.extend( - copydeep(getattr(self.results[e.source.node_id], e.source.field)) for e in item_edges - ) - node.collection = output_collection - else: - for edge in input_edges: - setattr( - node, - edge.destination.field, - copydeep(getattr(self.results[edge.source.node_id], edge.source.field)), - ) + self._runtime().prepare_inputs(node) # TODO: Add API for modifying underlying graph that checks if the change will be valid given the current execution state def _is_edge_valid(self, edge: Edge) -> bool: diff --git a/invokeai/app/services/shared/sqlite/sqlite_util.py b/invokeai/app/services/shared/sqlite/sqlite_util.py index 2478e8cdcae..fb8ca9fca38 100644 --- a/invokeai/app/services/shared/sqlite/sqlite_util.py +++ b/invokeai/app/services/shared/sqlite/sqlite_util.py @@ -31,6 +31,7 @@ from invokeai.app.services.shared.sqlite_migrator.migrations.migration_26 import build_migration_26 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_27 import build_migration_27 from invokeai.app.services.shared.sqlite_migrator.migrations.migration_28 import build_migration_28 +from invokeai.app.services.shared.sqlite_migrator.migrations.migration_29 import build_migration_29 from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_impl import SqliteMigrator @@ -79,6 +80,7 @@ def init_db(config: InvokeAIAppConfig, logger: Logger, image_files: ImageFileSto migrator.register_migration(build_migration_26(app_config=config, logger=logger)) migrator.register_migration(build_migration_27()) migrator.register_migration(build_migration_28()) + migrator.register_migration(build_migration_29()) migrator.run_migrations() return db diff --git a/invokeai/app/services/shared/sqlite_migrator/migrations/migration_29.py b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_29.py new file mode 100644 index 00000000000..c9eb7c901ba --- /dev/null +++ b/invokeai/app/services/shared/sqlite_migrator/migrations/migration_29.py @@ -0,0 +1,53 @@ +"""Migration 29: Add board_visibility column to boards table. + +This migration adds a board_visibility column to the boards table to support +three visibility levels: + - 'private': only the board owner (and admins) can view/modify + - 'shared': all users can view, but only the owner (and admins) can modify + - 'public': all users can view; only the owner (and admins) can modify the + board structure (rename/archive/delete) + +Existing boards with is_public = 1 are migrated to 'public'. +All other existing boards default to 'private'. +""" + +import sqlite3 + +from invokeai.app.services.shared.sqlite_migrator.sqlite_migrator_common import Migration + + +class Migration29Callback: + """Migration to add board_visibility column to the boards table.""" + + def __call__(self, cursor: sqlite3.Cursor) -> None: + self._update_boards_table(cursor) + + def _update_boards_table(self, cursor: sqlite3.Cursor) -> None: + """Add board_visibility column to boards table.""" + # Check if boards table exists + cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='boards';") + if cursor.fetchone() is None: + return + + cursor.execute("PRAGMA table_info(boards);") + columns = [row[1] for row in cursor.fetchall()] + + if "board_visibility" not in columns: + cursor.execute("ALTER TABLE boards ADD COLUMN board_visibility TEXT NOT NULL DEFAULT 'private';") + cursor.execute("CREATE INDEX IF NOT EXISTS idx_boards_board_visibility ON boards(board_visibility);") + # Migrate existing is_public = 1 boards to 'public' + if "is_public" in columns: + cursor.execute("UPDATE boards SET board_visibility = 'public' WHERE is_public = 1;") + + +def build_migration_29() -> Migration: + """Builds the migration object for migrating from version 28 to version 29. + + This migration adds the board_visibility column to the boards table, + supporting 'private', 'shared', and 'public' visibility levels. + """ + return Migration( + from_version=28, + to_version=29, + callback=Migration29Callback(), + ) diff --git a/invokeai/app/services/workflow_records/workflow_records_base.py b/invokeai/app/services/workflow_records/workflow_records_base.py index 8da1e97daf7..856a6c6d490 100644 --- a/invokeai/app/services/workflow_records/workflow_records_base.py +++ b/invokeai/app/services/workflow_records/workflow_records_base.py @@ -28,13 +28,13 @@ def create(self, workflow: WorkflowWithoutID, user_id: str = WORKFLOW_LIBRARY_DE pass @abstractmethod - def update(self, workflow: Workflow) -> WorkflowRecordDTO: - """Updates a workflow.""" + def update(self, workflow: Workflow, user_id: Optional[str] = None) -> WorkflowRecordDTO: + """Updates a workflow. When user_id is provided, the UPDATE is scoped to that user.""" pass @abstractmethod - def delete(self, workflow_id: str) -> None: - """Deletes a workflow.""" + def delete(self, workflow_id: str, user_id: Optional[str] = None) -> None: + """Deletes a workflow. When user_id is provided, the DELETE is scoped to that user.""" pass @abstractmethod @@ -78,8 +78,8 @@ def counts_by_tag( pass @abstractmethod - def update_opened_at(self, workflow_id: str) -> None: - """Open a workflow.""" + def update_opened_at(self, workflow_id: str, user_id: Optional[str] = None) -> None: + """Open a workflow. When user_id is provided, the UPDATE is scoped to that user.""" pass @abstractmethod @@ -93,6 +93,6 @@ def get_all_tags( pass @abstractmethod - def update_is_public(self, workflow_id: str, is_public: bool) -> WorkflowRecordDTO: - """Updates the is_public field of a workflow.""" + def update_is_public(self, workflow_id: str, is_public: bool, user_id: Optional[str] = None) -> WorkflowRecordDTO: + """Updates the is_public field of a workflow. When user_id is provided, the UPDATE is scoped to that user.""" pass diff --git a/invokeai/app/services/workflow_records/workflow_records_sqlite.py b/invokeai/app/services/workflow_records/workflow_records_sqlite.py index 0e6dfe1b700..c83d87eff68 100644 --- a/invokeai/app/services/workflow_records/workflow_records_sqlite.py +++ b/invokeai/app/services/workflow_records/workflow_records_sqlite.py @@ -67,36 +67,55 @@ def create(self, workflow: WorkflowWithoutID, user_id: str = WORKFLOW_LIBRARY_DE ) return self.get(workflow_with_id.id) - def update(self, workflow: Workflow) -> WorkflowRecordDTO: + def update(self, workflow: Workflow, user_id: Optional[str] = None) -> WorkflowRecordDTO: if workflow.meta.category is WorkflowCategory.Default: raise ValueError("Default workflows cannot be updated") with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE workflow_library - SET workflow = ? - WHERE workflow_id = ? AND category = 'user'; - """, - (workflow.model_dump_json(), workflow.id), - ) + if user_id is not None: + cursor.execute( + """--sql + UPDATE workflow_library + SET workflow = ? + WHERE workflow_id = ? AND category = 'user' AND user_id = ?; + """, + (workflow.model_dump_json(), workflow.id, user_id), + ) + else: + cursor.execute( + """--sql + UPDATE workflow_library + SET workflow = ? + WHERE workflow_id = ? AND category = 'user'; + """, + (workflow.model_dump_json(), workflow.id), + ) return self.get(workflow.id) - def delete(self, workflow_id: str) -> None: + def delete(self, workflow_id: str, user_id: Optional[str] = None) -> None: if self.get(workflow_id).workflow.meta.category is WorkflowCategory.Default: raise ValueError("Default workflows cannot be deleted") with self._db.transaction() as cursor: - cursor.execute( - """--sql - DELETE from workflow_library - WHERE workflow_id = ? AND category = 'user'; - """, - (workflow_id,), - ) + if user_id is not None: + cursor.execute( + """--sql + DELETE from workflow_library + WHERE workflow_id = ? AND category = 'user' AND user_id = ?; + """, + (workflow_id, user_id), + ) + else: + cursor.execute( + """--sql + DELETE from workflow_library + WHERE workflow_id = ? AND category = 'user'; + """, + (workflow_id,), + ) return None - def update_is_public(self, workflow_id: str, is_public: bool) -> WorkflowRecordDTO: + def update_is_public(self, workflow_id: str, is_public: bool, user_id: Optional[str] = None) -> WorkflowRecordDTO: """Updates the is_public field of a workflow and manages the 'shared' tag automatically.""" record = self.get(workflow_id) workflow = record.workflow @@ -111,14 +130,24 @@ def update_is_public(self, workflow_id: str, is_public: bool) -> WorkflowRecordD updated_workflow = workflow.model_copy(update={"tags": updated_tags}) with self._db.transaction() as cursor: - cursor.execute( - """--sql - UPDATE workflow_library - SET workflow = ?, is_public = ? - WHERE workflow_id = ? AND category = 'user'; - """, - (updated_workflow.model_dump_json(), is_public, workflow_id), - ) + if user_id is not None: + cursor.execute( + """--sql + UPDATE workflow_library + SET workflow = ?, is_public = ? + WHERE workflow_id = ? AND category = 'user' AND user_id = ?; + """, + (updated_workflow.model_dump_json(), is_public, workflow_id, user_id), + ) + else: + cursor.execute( + """--sql + UPDATE workflow_library + SET workflow = ?, is_public = ? + WHERE workflow_id = ? AND category = 'user'; + """, + (updated_workflow.model_dump_json(), is_public, workflow_id), + ) return self.get(workflow_id) def get_many( @@ -209,7 +238,8 @@ def get_many( params.extend([wildcard_query, wildcard_query, wildcard_query]) if user_id is not None: - conditions.append("user_id = ?") + # Scope to the given user but always include default workflows + conditions.append("(user_id = ? OR category = 'default')") params.append(user_id) if is_public is True: @@ -291,7 +321,8 @@ def counts_by_tag( base_conditions.append("opened_at IS NULL") if user_id is not None: - base_conditions.append("user_id = ?") + # Scope to the given user but always include default workflows + base_conditions.append("(user_id = ? OR category = 'default')") base_params.append(user_id) if is_public is True: @@ -350,7 +381,8 @@ def counts_by_category( base_conditions.append("opened_at IS NULL") if user_id is not None: - base_conditions.append("user_id = ?") + # Scope to the given user but always include default workflows + base_conditions.append("(user_id = ? OR category = 'default')") base_params.append(user_id) if is_public is True: @@ -383,16 +415,26 @@ def counts_by_category( return result - def update_opened_at(self, workflow_id: str) -> None: + def update_opened_at(self, workflow_id: str, user_id: Optional[str] = None) -> None: with self._db.transaction() as cursor: - cursor.execute( - f"""--sql - UPDATE workflow_library - SET opened_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') - WHERE workflow_id = ?; - """, - (workflow_id,), - ) + if user_id is not None: + cursor.execute( + f"""--sql + UPDATE workflow_library + SET opened_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') + WHERE workflow_id = ? AND user_id = ?; + """, + (workflow_id, user_id), + ) + else: + cursor.execute( + f"""--sql + UPDATE workflow_library + SET opened_at = STRFTIME('{SQL_TIME_FORMAT}', 'NOW') + WHERE workflow_id = ?; + """, + (workflow_id,), + ) def get_all_tags( self, @@ -414,7 +456,8 @@ def get_all_tags( params.extend([category.value for category in categories]) if user_id is not None: - conditions.append("user_id = ?") + # Scope to the given user but always include default workflows + conditions.append("(user_id = ? OR category = 'default')") params.append(user_id) if is_public is True: diff --git a/invokeai/app/util/step_callback.py b/invokeai/app/util/step_callback.py index 990fdd51d8b..08dc9a2265c 100644 --- a/invokeai/app/util/step_callback.py +++ b/invokeai/app/util/step_callback.py @@ -93,6 +93,29 @@ [-0.00955853, -0.00980067, -0.00977842], ] +# Qwen Image uses the same VAE as Wan 2.1 (16-channel). +# Factors from ComfyUI: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/latent_formats.py +QWEN_IMAGE_LATENT_RGB_FACTORS = [ + [-0.1299, -0.1692, 0.2932], + [0.0671, 0.0406, 0.0442], + [0.3568, 0.2548, 0.1747], + [0.0372, 0.2344, 0.1420], + [0.0313, 0.0189, -0.0328], + [0.0296, -0.0956, -0.0665], + [-0.3477, -0.4059, -0.2925], + [0.0166, 0.1902, 0.1975], + [-0.0412, 0.0267, -0.1364], + [-0.1293, 0.0740, 0.1636], + [0.0680, 0.3019, 0.1128], + [0.0032, 0.0581, 0.0639], + [-0.1251, 0.0927, 0.1699], + [0.0060, -0.0633, 0.0005], + [0.3477, 0.2275, 0.2950], + [0.1984, 0.0913, 0.1861], +] + +QWEN_IMAGE_LATENT_RGB_BIAS = [-0.1835, -0.0868, -0.3360] + # FLUX.2 uses 32 latent channels. # Factors from ComfyUI: https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/latent_formats.py FLUX2_LATENT_RGB_FACTORS = [ @@ -133,6 +156,29 @@ FLUX2_LATENT_RGB_BIAS = [-0.0329, -0.0718, -0.0851] +# Anima uses Wan 2.1 VAE with 16 latent channels. +# Factors from ComfyUI: https://github.com/Comfy-Org/ComfyUI/blob/main/comfy/latent_formats.py +ANIMA_LATENT_RGB_FACTORS = [ + [-0.1299, -0.1692, 0.2932], + [0.0671, 0.0406, 0.0442], + [0.3568, 0.2548, 0.1747], + [0.0372, 0.2344, 0.1420], + [0.0313, 0.0189, -0.0328], + [0.0296, -0.0956, -0.0665], + [-0.3477, -0.4059, -0.2925], + [0.0166, 0.1902, 0.1975], + [-0.0412, 0.0267, -0.1364], + [-0.1293, 0.0740, 0.1636], + [0.0680, 0.3019, 0.1128], + [0.0032, 0.0581, 0.0639], + [-0.1251, 0.0927, 0.1699], + [0.0060, -0.0633, 0.0005], + [0.3477, 0.2275, 0.2950], + [0.1984, 0.0913, 0.1861], +] + +ANIMA_LATENT_RGB_BIAS = [-0.1835, -0.0868, -0.3360] + def sample_to_lowres_estimated_image( samples: torch.Tensor, @@ -209,6 +255,9 @@ def diffusion_step_callback( latent_rgb_factors = SD3_5_LATENT_RGB_FACTORS elif base_model == BaseModelType.CogView4: latent_rgb_factors = COGVIEW4_LATENT_RGB_FACTORS + elif base_model == BaseModelType.QwenImage: + latent_rgb_factors = QWEN_IMAGE_LATENT_RGB_FACTORS + latent_rgb_bias = QWEN_IMAGE_LATENT_RGB_BIAS elif base_model == BaseModelType.Flux: latent_rgb_factors = FLUX_LATENT_RGB_FACTORS elif base_model == BaseModelType.Flux2: @@ -217,6 +266,10 @@ def diffusion_step_callback( elif base_model == BaseModelType.ZImage: # Z-Image uses FLUX-compatible VAE with 16 latent channels latent_rgb_factors = FLUX_LATENT_RGB_FACTORS + elif base_model == BaseModelType.Anima: + # Anima uses Wan 2.1 VAE with 16 latent channels + latent_rgb_factors = ANIMA_LATENT_RGB_FACTORS + latent_rgb_bias = ANIMA_LATENT_RGB_BIAS else: raise ValueError(f"Unsupported base model: {base_model}") diff --git a/invokeai/backend/anima/__init__.py b/invokeai/backend/anima/__init__.py new file mode 100644 index 00000000000..01a1a952e96 --- /dev/null +++ b/invokeai/backend/anima/__init__.py @@ -0,0 +1,6 @@ +"""Anima model backend module. + +Anima is a 2B-parameter anime-focused text-to-image model built on NVIDIA's +Cosmos Predict2 DiT architecture with a custom LLM Adapter that bridges Qwen3 +0.6B text encoder outputs to the DiT backbone. +""" diff --git a/invokeai/backend/anima/anima_transformer.py b/invokeai/backend/anima/anima_transformer.py new file mode 100644 index 00000000000..36c5764e97e --- /dev/null +++ b/invokeai/backend/anima/anima_transformer.py @@ -0,0 +1,1040 @@ +"""Anima transformer model: Cosmos Predict2 MiniTrainDIT + LLM Adapter. + +The Anima architecture combines: +1. MiniTrainDIT: A Cosmos Predict2 DiT backbone with 28 blocks, 2048-dim hidden state, + and 3D RoPE positional embeddings. +2. LLMAdapter: A 6-layer cross-attention transformer that fuses Qwen3 0.6B hidden states + with learned T5-XXL token embeddings to produce conditioning for the DiT. + +Original source code: +- MiniTrainDIT backbone and positional embeddings: https://github.com/nvidia-cosmos/cosmos-predict2 + SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + SPDX-License-Identifier: Apache-2.0 +- LLMAdapter and Anima wrapper: Clean-room implementation based on + https://github.com/hdae/diffusers-anima (Apache-2.0) +""" + +import logging +import math +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from einops.layers.torch import Rearrange +from torch import nn + +logger = logging.getLogger(__name__) + + +# ============================================================================ +# Positional Embeddings +# Original source: https://github.com/nvidia-cosmos/cosmos-predict2 +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. Apache-2.0 +# ============================================================================ + + +class VideoRopePosition3DEmb(nn.Module): + """3D Rotary Position Embedding for video/image transformers. + + Generates rotary embeddings with separate frequency components for + height, width, and temporal dimensions. + """ + + def __init__( + self, + *, + head_dim: int, + len_h: int, + len_w: int, + len_t: int, + base_fps: int = 24, + h_extrapolation_ratio: float = 1.0, + w_extrapolation_ratio: float = 1.0, + t_extrapolation_ratio: float = 1.0, + enable_fps_modulation: bool = True, + device: Optional[torch.device] = None, + **kwargs, + ): + super().__init__() + self.base_fps = base_fps + self.max_h = len_h + self.max_w = len_w + self.enable_fps_modulation = enable_fps_modulation + + dim = head_dim + dim_h = dim // 6 * 2 + dim_w = dim_h + dim_t = dim - 2 * dim_h + assert dim == dim_h + dim_w + dim_t, f"bad dim: {dim} != {dim_h} + {dim_w} + {dim_t}" + + self.register_buffer( + "dim_spatial_range", + torch.arange(0, dim_h, 2, device=device)[: (dim_h // 2)].float() / dim_h, + persistent=False, + ) + self.register_buffer( + "dim_temporal_range", + torch.arange(0, dim_t, 2, device=device)[: (dim_t // 2)].float() / dim_t, + persistent=False, + ) + + self.h_ntk_factor = h_extrapolation_ratio ** (dim_h / (dim_h - 2)) + self.w_ntk_factor = w_extrapolation_ratio ** (dim_w / (dim_w - 2)) + self.t_ntk_factor = t_extrapolation_ratio ** (dim_t / (dim_t - 2)) + + def forward( + self, + x_B_T_H_W_C: torch.Tensor, + fps: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + ) -> torch.Tensor: + return self.generate_embeddings(x_B_T_H_W_C.shape, fps=fps, device=device) + + def generate_embeddings( + self, + B_T_H_W_C: torch.Size, + fps: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + h_theta = 10000.0 * self.h_ntk_factor + w_theta = 10000.0 * self.w_ntk_factor + t_theta = 10000.0 * self.t_ntk_factor + + h_spatial_freqs = 1.0 / (h_theta ** self.dim_spatial_range.to(device=device)) + w_spatial_freqs = 1.0 / (w_theta ** self.dim_spatial_range.to(device=device)) + temporal_freqs = 1.0 / (t_theta ** self.dim_temporal_range.to(device=device)) + + B, T, H, W, _ = B_T_H_W_C + seq = torch.arange(max(H, W, T), dtype=torch.float, device=device) + + half_emb_h = torch.outer(seq[:H].to(device=device), h_spatial_freqs) + half_emb_w = torch.outer(seq[:W].to(device=device), w_spatial_freqs) + + if fps is None or self.enable_fps_modulation is False: + half_emb_t = torch.outer(seq[:T].to(device=device), temporal_freqs) + else: + half_emb_t = torch.outer(seq[:T].to(device=device) / fps * self.base_fps, temporal_freqs) + + half_emb_h = torch.stack( + [torch.cos(half_emb_h), -torch.sin(half_emb_h), torch.sin(half_emb_h), torch.cos(half_emb_h)], dim=-1 + ) + half_emb_w = torch.stack( + [torch.cos(half_emb_w), -torch.sin(half_emb_w), torch.sin(half_emb_w), torch.cos(half_emb_w)], dim=-1 + ) + half_emb_t = torch.stack( + [torch.cos(half_emb_t), -torch.sin(half_emb_t), torch.sin(half_emb_t), torch.cos(half_emb_t)], dim=-1 + ) + + em_T_H_W_D = torch.cat( + [ + repeat(half_emb_t, "t d x -> t h w d x", h=H, w=W), + repeat(half_emb_h, "h d x -> t h w d x", t=T, w=W), + repeat(half_emb_w, "w d x -> t h w d x", t=T, h=H), + ], + dim=-2, + ) + + return rearrange(em_T_H_W_D, "t h w d (i j) -> (t h w) d i j", i=2, j=2).float() + + +def _normalize(x: torch.Tensor, dim: Optional[list[int]] = None, eps: float = 0) -> torch.Tensor: + if dim is None: + dim = list(range(1, x.ndim)) + norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32) + norm = torch.add(eps, norm, alpha=math.sqrt(norm.numel() / x.numel())) + return x / norm.to(x.dtype) + + +class LearnablePosEmbAxis(nn.Module): + """Learnable per-axis positional embeddings.""" + + def __init__( + self, + *, + model_channels: int, + len_h: int, + len_w: int, + len_t: int, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + **kwargs, + ): + super().__init__() + self.pos_emb_h = nn.Parameter(torch.empty(len_h, model_channels, device=device, dtype=dtype)) + self.pos_emb_w = nn.Parameter(torch.empty(len_w, model_channels, device=device, dtype=dtype)) + self.pos_emb_t = nn.Parameter(torch.empty(len_t, model_channels, device=device, dtype=dtype)) + + def forward( + self, + x_B_T_H_W_C: torch.Tensor, + fps: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + return self.generate_embeddings(x_B_T_H_W_C.shape, device=device, dtype=dtype) + + def generate_embeddings( + self, + B_T_H_W_C: torch.Size, + fps: Optional[torch.Tensor] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> torch.Tensor: + B, T, H, W, _ = B_T_H_W_C + emb_h_H = self.pos_emb_h[:H].to(device=device, dtype=dtype) + emb_w_W = self.pos_emb_w[:W].to(device=device, dtype=dtype) + emb_t_T = self.pos_emb_t[:T].to(device=device, dtype=dtype) + emb = ( + repeat(emb_t_T, "t d -> b t h w d", b=B, h=H, w=W) + + repeat(emb_h_H, "h d -> b t h w d", b=B, t=T, w=W) + + repeat(emb_w_W, "w d -> b t h w d", b=B, t=T, h=H) + ) + return _normalize(emb, dim=-1, eps=1e-6) + + +# ============================================================================ +# Cosmos Predict2 MiniTrainDIT +# Original source: https://github.com/nvidia-cosmos/cosmos-predict2 +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. Apache-2.0 +# ============================================================================ + + +def apply_rotary_pos_emb_cosmos(t: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Apply rotary position embeddings in Cosmos format (2x2 rotation matrices).""" + t_ = t.reshape(*t.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).float() + t_out = freqs[..., 0] * t_[..., 0] + freqs[..., 1] * t_[..., 1] + t_out = t_out.movedim(-1, -2).reshape(*t.shape).type_as(t) + return t_out + + +class GPT2FeedForward(nn.Module): + def __init__(self, d_model: int, d_ff: int) -> None: + super().__init__() + self.activation = nn.GELU() + self.layer1 = nn.Linear(d_model, d_ff, bias=False) + self.layer2 = nn.Linear(d_ff, d_model, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layer2(self.activation(self.layer1(x))) + + +class CosmosAttention(nn.Module): + """Multi-head attention for the Cosmos DiT backbone. + + Supports both self-attention and cross-attention with QK normalization + and rotary position embeddings. + """ + + def __init__( + self, + query_dim: int, + context_dim: Optional[int] = None, + n_heads: int = 8, + head_dim: int = 64, + dropout: float = 0.0, + ) -> None: + super().__init__() + self.is_selfattn = context_dim is None + context_dim = query_dim if context_dim is None else context_dim + inner_dim = head_dim * n_heads + + self.n_heads = n_heads + self.head_dim = head_dim + + self.q_proj = nn.Linear(query_dim, inner_dim, bias=False) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + + self.k_proj = nn.Linear(context_dim, inner_dim, bias=False) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + + self.v_proj = nn.Linear(context_dim, inner_dim, bias=False) + self.v_norm = nn.Identity() + + self.output_proj = nn.Linear(inner_dim, query_dim, bias=False) + self.output_dropout = nn.Dropout(dropout) if dropout > 1e-4 else nn.Identity() + + def forward( + self, + x: torch.Tensor, + context: Optional[torch.Tensor] = None, + rope_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + q = self.q_proj(x) + context = x if context is None else context + k = self.k_proj(context) + v = self.v_proj(context) + q, k, v = (rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim) for t in (q, k, v)) + + q = self.q_norm(q) + k = self.k_norm(k) + v = self.v_norm(v) + + if self.is_selfattn and rope_emb is not None: + q = apply_rotary_pos_emb_cosmos(q, rope_emb) + k = apply_rotary_pos_emb_cosmos(k, rope_emb) + + # Reshape for scaled_dot_product_attention: (B, heads, seq, dim) + in_q_shape = q.shape + in_k_shape = k.shape + q = rearrange(q, "b ... h d -> b h ... d").reshape(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1]) + k = rearrange(k, "b ... h d -> b h ... d").reshape(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1]) + v = rearrange(v, "b ... h d -> b h ... d").reshape(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1]) + + result = F.scaled_dot_product_attention(q, k, v) + result = rearrange(result, "b h s d -> b s (h d)") + return self.output_dropout(self.output_proj(result)) + + +class Timesteps(nn.Module): + """Sinusoidal timestep embeddings.""" + + def __init__(self, num_channels: int): + super().__init__() + self.num_channels = num_channels + + def forward(self, timesteps_B_T: torch.Tensor) -> torch.Tensor: + assert timesteps_B_T.ndim == 2 + timesteps = timesteps_B_T.flatten().float() + half_dim = self.num_channels // 2 + exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) / half_dim + emb = timesteps[:, None].float() * torch.exp(exponent)[None, :] + emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + return rearrange(emb, "(b t) d -> b t d", b=timesteps_B_T.shape[0], t=timesteps_B_T.shape[1]) + + +class TimestepEmbedding(nn.Module): + """Projects sinusoidal timestep embeddings to model dimension.""" + + def __init__(self, in_features: int, out_features: int, use_adaln_lora: bool = False): + super().__init__() + self.use_adaln_lora = use_adaln_lora + self.linear_1 = nn.Linear(in_features, out_features, bias=not use_adaln_lora) + self.activation = nn.SiLU() + if use_adaln_lora: + self.linear_2 = nn.Linear(out_features, 3 * out_features, bias=False) + else: + self.linear_2 = nn.Linear(out_features, out_features, bias=False) + + def forward(self, sample: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + emb = self.linear_2(self.activation(self.linear_1(sample))) + if self.use_adaln_lora: + return sample, emb + return emb, None + + +class PatchEmbed(nn.Module): + """Patchify input tensor via rearrange + linear projection.""" + + def __init__( + self, + spatial_patch_size: int, + temporal_patch_size: int, + in_channels: int = 3, + out_channels: int = 768, + ): + super().__init__() + self.spatial_patch_size = spatial_patch_size + self.temporal_patch_size = temporal_patch_size + self.proj = nn.Sequential( + Rearrange( + "b c (t r) (h m) (w n) -> b t h w (c r m n)", + r=temporal_patch_size, + m=spatial_patch_size, + n=spatial_patch_size, + ), + nn.Linear( + in_channels * spatial_patch_size * spatial_patch_size * temporal_patch_size, + out_channels, + bias=False, + ), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + assert x.dim() == 5 + return self.proj(x) + + +class FinalLayer(nn.Module): + """Final AdaLN-modulated output projection.""" + + def __init__( + self, + hidden_size: int, + spatial_patch_size: int, + temporal_patch_size: int, + out_channels: int, + use_adaln_lora: bool = False, + adaln_lora_dim: int = 256, + ): + super().__init__() + self.layer_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = nn.Linear( + hidden_size, spatial_patch_size * spatial_patch_size * temporal_patch_size * out_channels, bias=False + ) + self.hidden_size = hidden_size + self.use_adaln_lora = use_adaln_lora + + if use_adaln_lora: + self.adaln_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, 2 * hidden_size, bias=False), + ) + else: + self.adaln_modulation = nn.Sequential( + nn.SiLU(), + nn.Linear(hidden_size, 2 * hidden_size, bias=False), + ) + + def forward( + self, + x_B_T_H_W_D: torch.Tensor, + emb_B_T_D: torch.Tensor, + adaln_lora_B_T_3D: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.use_adaln_lora: + assert adaln_lora_B_T_3D is not None + shift, scale = (self.adaln_modulation(emb_B_T_D) + adaln_lora_B_T_3D[:, :, : 2 * self.hidden_size]).chunk( + 2, dim=-1 + ) + else: + shift, scale = self.adaln_modulation(emb_B_T_D).chunk(2, dim=-1) + + shift = rearrange(shift, "b t d -> b t 1 1 d") + scale = rearrange(scale, "b t d -> b t 1 1 d") + + x_B_T_H_W_D = self.layer_norm(x_B_T_H_W_D) * (1 + scale) + shift + return self.linear(x_B_T_H_W_D) + + +class DiTBlock(nn.Module): + """Cosmos DiT transformer block with self-attention, cross-attention, and MLP. + + Each component uses AdaLN (Adaptive Layer Normalization) modulation from + the timestep embedding. + """ + + def __init__( + self, + x_dim: int, + context_dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + use_adaln_lora: bool = False, + adaln_lora_dim: int = 256, + ): + super().__init__() + self.x_dim = x_dim + self.use_adaln_lora = use_adaln_lora + + self.layer_norm_self_attn = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) + self.self_attn = CosmosAttention(x_dim, None, num_heads, x_dim // num_heads) + + self.layer_norm_cross_attn = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) + self.cross_attn = CosmosAttention(x_dim, context_dim, num_heads, x_dim // num_heads) + + self.layer_norm_mlp = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) + self.mlp = GPT2FeedForward(x_dim, int(x_dim * mlp_ratio)) + + # AdaLN modulation layers (shift, scale, gate for each of 3 components) + if use_adaln_lora: + self.adaln_modulation_self_attn = nn.Sequential( + nn.SiLU(), + nn.Linear(x_dim, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False), + ) + self.adaln_modulation_cross_attn = nn.Sequential( + nn.SiLU(), + nn.Linear(x_dim, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False), + ) + self.adaln_modulation_mlp = nn.Sequential( + nn.SiLU(), + nn.Linear(x_dim, adaln_lora_dim, bias=False), + nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False), + ) + else: + self.adaln_modulation_self_attn = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False)) + self.adaln_modulation_cross_attn = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False)) + self.adaln_modulation_mlp = nn.Sequential(nn.SiLU(), nn.Linear(x_dim, 3 * x_dim, bias=False)) + + def forward( + self, + x_B_T_H_W_D: torch.Tensor, + emb_B_T_D: torch.Tensor, + crossattn_emb: torch.Tensor, + rope_emb_L_1_1_D: Optional[torch.Tensor] = None, + adaln_lora_B_T_3D: Optional[torch.Tensor] = None, + extra_per_block_pos_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + residual_dtype = x_B_T_H_W_D.dtype + compute_dtype = emb_B_T_D.dtype + + if extra_per_block_pos_emb is not None: + x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb + + # Compute AdaLN modulations + if self.use_adaln_lora: + assert adaln_lora_B_T_3D is not None + shift_sa, scale_sa, gate_sa = (self.adaln_modulation_self_attn(emb_B_T_D) + adaln_lora_B_T_3D).chunk( + 3, dim=-1 + ) + shift_ca, scale_ca, gate_ca = (self.adaln_modulation_cross_attn(emb_B_T_D) + adaln_lora_B_T_3D).chunk( + 3, dim=-1 + ) + shift_mlp, scale_mlp, gate_mlp = (self.adaln_modulation_mlp(emb_B_T_D) + adaln_lora_B_T_3D).chunk(3, dim=-1) + else: + shift_sa, scale_sa, gate_sa = self.adaln_modulation_self_attn(emb_B_T_D).chunk(3, dim=-1) + shift_ca, scale_ca, gate_ca = self.adaln_modulation_cross_attn(emb_B_T_D).chunk(3, dim=-1) + shift_mlp, scale_mlp, gate_mlp = self.adaln_modulation_mlp(emb_B_T_D).chunk(3, dim=-1) + + # Reshape for broadcasting: (B, T, D) -> (B, T, 1, 1, D) + shift_sa, scale_sa, gate_sa = (rearrange(t, "b t d -> b t 1 1 d") for t in (shift_sa, scale_sa, gate_sa)) + shift_ca, scale_ca, gate_ca = (rearrange(t, "b t d -> b t 1 1 d") for t in (shift_ca, scale_ca, gate_ca)) + shift_mlp, scale_mlp, gate_mlp = (rearrange(t, "b t d -> b t 1 1 d") for t in (shift_mlp, scale_mlp, gate_mlp)) + + B, T, H, W, D = x_B_T_H_W_D.shape + + def _adaln(x: torch.Tensor, norm: nn.Module, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + return norm(x) * (1 + scale) + shift + + # Self-attention + normed = _adaln(x_B_T_H_W_D, self.layer_norm_self_attn, scale_sa, shift_sa) + result = rearrange( + self.self_attn( + rearrange(normed.to(compute_dtype), "b t h w d -> b (t h w) d"), None, rope_emb=rope_emb_L_1_1_D + ), + "b (t h w) d -> b t h w d", + t=T, + h=H, + w=W, + ) + x_B_T_H_W_D = x_B_T_H_W_D + gate_sa.to(residual_dtype) * result.to(residual_dtype) + + # Cross-attention + normed = _adaln(x_B_T_H_W_D, self.layer_norm_cross_attn, scale_ca, shift_ca) + result = rearrange( + self.cross_attn( + rearrange(normed.to(compute_dtype), "b t h w d -> b (t h w) d"), + crossattn_emb, + rope_emb=rope_emb_L_1_1_D, + ), + "b (t h w) d -> b t h w d", + t=T, + h=H, + w=W, + ) + x_B_T_H_W_D = result.to(residual_dtype) * gate_ca.to(residual_dtype) + x_B_T_H_W_D + + # MLP + normed = _adaln(x_B_T_H_W_D, self.layer_norm_mlp, scale_mlp, shift_mlp) + result = self.mlp(normed.to(compute_dtype)) + x_B_T_H_W_D = x_B_T_H_W_D + gate_mlp.to(residual_dtype) * result.to(residual_dtype) + + return x_B_T_H_W_D + + +class MiniTrainDIT(nn.Module): + """Cosmos Predict2 DiT backbone for video/image generation. + + This is the core transformer architecture that Anima extends. It processes + 3D latent tensors (B, C, T, H, W) with patch embedding, positional encoding, + and adaptive layer normalization. + + Args: + max_img_h: Maximum image height in pixels. + max_img_w: Maximum image width in pixels. + max_frames: Maximum number of video frames. + in_channels: Number of input latent channels. + out_channels: Number of output channels. + patch_spatial: Spatial patch size. + patch_temporal: Temporal patch size. + concat_padding_mask: Whether to concatenate a padding mask channel. + model_channels: Hidden dimension of the transformer. + num_blocks: Number of DiT blocks. + num_heads: Number of attention heads. + mlp_ratio: MLP expansion ratio. + crossattn_emb_channels: Cross-attention context dimension. + use_adaln_lora: Whether to use AdaLN-LoRA. + adaln_lora_dim: AdaLN-LoRA bottleneck dimension. + extra_per_block_abs_pos_emb: Whether to use extra learnable positional embeddings. + """ + + def __init__( + self, + max_img_h: int = 240, + max_img_w: int = 240, + max_frames: int = 1, + in_channels: int = 16, + out_channels: int = 16, + patch_spatial: int = 2, + patch_temporal: int = 1, + concat_padding_mask: bool = True, + model_channels: int = 2048, + num_blocks: int = 28, + num_heads: int = 16, + mlp_ratio: float = 4.0, + crossattn_emb_channels: int = 1024, + pos_emb_cls: str = "rope3d", + pos_emb_learnable: bool = False, + pos_emb_interpolation: str = "crop", + min_fps: int = 1, + max_fps: int = 30, + use_adaln_lora: bool = False, + adaln_lora_dim: int = 256, + rope_h_extrapolation_ratio: float = 1.0, + rope_w_extrapolation_ratio: float = 1.0, + rope_t_extrapolation_ratio: float = 1.0, + extra_per_block_abs_pos_emb: bool = False, + extra_h_extrapolation_ratio: float = 1.0, + extra_w_extrapolation_ratio: float = 1.0, + extra_t_extrapolation_ratio: float = 1.0, + rope_enable_fps_modulation: bool = True, + image_model: Optional[str] = None, + ) -> None: + super().__init__() + self.max_img_h = max_img_h + self.max_img_w = max_img_w + self.max_frames = max_frames + self.in_channels = in_channels + self.out_channels = out_channels + self.patch_spatial = patch_spatial + self.patch_temporal = patch_temporal + self.num_heads = num_heads + self.num_blocks = num_blocks + self.model_channels = model_channels + self.concat_padding_mask = concat_padding_mask + self.pos_emb_cls = pos_emb_cls + self.extra_per_block_abs_pos_emb = extra_per_block_abs_pos_emb + + # Positional embeddings + self.pos_embedder = VideoRopePosition3DEmb( + head_dim=model_channels // num_heads, + len_h=max_img_h // patch_spatial, + len_w=max_img_w // patch_spatial, + len_t=max_frames // patch_temporal, + max_fps=max_fps, + min_fps=min_fps, + h_extrapolation_ratio=rope_h_extrapolation_ratio, + w_extrapolation_ratio=rope_w_extrapolation_ratio, + t_extrapolation_ratio=rope_t_extrapolation_ratio, + enable_fps_modulation=rope_enable_fps_modulation, + ) + + if extra_per_block_abs_pos_emb: + self.extra_pos_embedder = LearnablePosEmbAxis( + model_channels=model_channels, + len_h=max_img_h // patch_spatial, + len_w=max_img_w // patch_spatial, + len_t=max_frames // patch_temporal, + ) + + self.use_adaln_lora = use_adaln_lora + self.adaln_lora_dim = adaln_lora_dim + + # Timestep embedding + self.t_embedder = nn.Sequential( + Timesteps(model_channels), + TimestepEmbedding(model_channels, model_channels, use_adaln_lora=use_adaln_lora), + ) + self.t_embedding_norm = nn.RMSNorm(model_channels, eps=1e-6) + + # Patch embedding + embed_in_channels = in_channels + 1 if concat_padding_mask else in_channels + self.x_embedder = PatchEmbed( + spatial_patch_size=patch_spatial, + temporal_patch_size=patch_temporal, + in_channels=embed_in_channels, + out_channels=model_channels, + ) + + # Transformer blocks + self.blocks = nn.ModuleList( + [ + DiTBlock( + x_dim=model_channels, + context_dim=crossattn_emb_channels, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + use_adaln_lora=use_adaln_lora, + adaln_lora_dim=adaln_lora_dim, + ) + for _ in range(num_blocks) + ] + ) + + # Final output layer + self.final_layer = FinalLayer( + hidden_size=model_channels, + spatial_patch_size=patch_spatial, + temporal_patch_size=patch_temporal, + out_channels=out_channels, + use_adaln_lora=use_adaln_lora, + adaln_lora_dim=adaln_lora_dim, + ) + + def _pad_to_patch_size(self, x: torch.Tensor) -> torch.Tensor: + """Pad input tensor so dimensions are divisible by patch sizes.""" + _, _, T, H, W = x.shape + pad_t = (self.patch_temporal - T % self.patch_temporal) % self.patch_temporal + pad_h = (self.patch_spatial - H % self.patch_spatial) % self.patch_spatial + pad_w = (self.patch_spatial - W % self.patch_spatial) % self.patch_spatial + if pad_t > 0 or pad_h > 0 or pad_w > 0: + x = F.pad(x, (0, pad_w, 0, pad_h, 0, pad_t)) + return x + + def prepare_embedded_sequence( + self, + x_B_C_T_H_W: torch.Tensor, + fps: Optional[torch.Tensor] = None, + padding_mask: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + if self.concat_padding_mask: + if padding_mask is None: + padding_mask = torch.zeros( + x_B_C_T_H_W.shape[0], + 1, + x_B_C_T_H_W.shape[3], + x_B_C_T_H_W.shape[4], + dtype=x_B_C_T_H_W.dtype, + device=x_B_C_T_H_W.device, + ) + x_B_C_T_H_W = torch.cat( + [x_B_C_T_H_W, padding_mask.unsqueeze(1).repeat(1, 1, x_B_C_T_H_W.shape[2], 1, 1)], dim=1 + ) + + x_B_T_H_W_D = self.x_embedder(x_B_C_T_H_W) + + extra_pos_emb = None + if self.extra_per_block_abs_pos_emb: + extra_pos_emb = self.extra_pos_embedder( + x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device, dtype=x_B_C_T_H_W.dtype + ) + + if "rope" in self.pos_emb_cls.lower(): + return x_B_T_H_W_D, self.pos_embedder(x_B_T_H_W_D, fps=fps, device=x_B_C_T_H_W.device), extra_pos_emb + + return x_B_T_H_W_D, None, extra_pos_emb + + def unpatchify(self, x_B_T_H_W_M: torch.Tensor) -> torch.Tensor: + return rearrange( + x_B_T_H_W_M, + "B T H W (p1 p2 t C) -> B C (T t) (H p1) (W p2)", + p1=self.patch_spatial, + p2=self.patch_spatial, + t=self.patch_temporal, + ) + + def forward( + self, + x: torch.Tensor, + timesteps: torch.Tensor, + context: torch.Tensor, + fps: Optional[torch.Tensor] = None, + padding_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + orig_shape = list(x.shape) + x = self._pad_to_patch_size(x) + + x_B_T_H_W_D, rope_emb_L_1_1_D, extra_pos_emb = self.prepare_embedded_sequence( + x, fps=fps, padding_mask=padding_mask + ) + + if timesteps.ndim == 1: + timesteps = timesteps.unsqueeze(1) + t_emb, adaln_lora = self.t_embedder[1](self.t_embedder[0](timesteps).to(x_B_T_H_W_D.dtype)) + t_emb = self.t_embedding_norm(t_emb) + + block_kwargs = { + "rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0) if rope_emb_L_1_1_D is not None else None, + "adaln_lora_B_T_3D": adaln_lora, + "extra_per_block_pos_emb": extra_pos_emb, + } + + # Keep residual stream in fp32 for numerical stability with fp16 compute + if x_B_T_H_W_D.dtype == torch.float16: + x_B_T_H_W_D = x_B_T_H_W_D.float() + + for block in self.blocks: + x_B_T_H_W_D = block(x_B_T_H_W_D, t_emb, context, **block_kwargs) + + x_out = self.final_layer(x_B_T_H_W_D.to(context.dtype), t_emb, adaln_lora_B_T_3D=adaln_lora) + x_out = self.unpatchify(x_out)[:, :, : orig_shape[-3], : orig_shape[-2], : orig_shape[-1]] + return x_out + + +# ============================================================================ +# LLM Adapter +# Reference implementation: https://github.com/hdae/diffusers-anima +# SPDX-License-Identifier: Apache-2.0 +# ============================================================================ + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Split the last dimension in half and negate-swap: [-x2, x1].""" + half = x.shape[-1] // 2 + first, second = x[..., :half], x[..., half:] + return torch.cat((-second, first), dim=-1) + + +def _apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + """Apply rotary position embeddings to tensor x given precomputed cos/sin.""" + return (x * cos.unsqueeze(1)) + (_rotate_half(x) * sin.unsqueeze(1)) + + +class LLMAdapterRotaryEmbedding(nn.Module): + """Rotary position embedding for the LLM Adapter's attention layers.""" + + def __init__(self, head_dim: int, theta: float = 10000.0): + super().__init__() + half_dim = head_dim // 2 + index = torch.arange(half_dim, dtype=torch.float32) + exponent = (2.0 / float(head_dim)) * index + inv_freq = torch.reciprocal(torch.pow(torch.tensor(theta, dtype=torch.float32), exponent)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + pos = position_ids.to(device=x.device, dtype=torch.float32) + inv = self.inv_freq.to(device=x.device, dtype=torch.float32) + freqs = torch.einsum("bl,d->bld", pos, inv) + emb = freqs.repeat(1, 1, 2) + return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype) + + +class LLMAdapterAttention(nn.Module): + """Attention for the LLM Adapter with QK normalization and rotary position embeddings.""" + + def __init__(self, query_dim: int, context_dim: int, n_heads: int, head_dim: int): + super().__init__() + inner_dim = head_dim * n_heads + self.n_heads = n_heads + self.head_dim = head_dim + + self.q_proj = nn.Linear(query_dim, inner_dim, bias=False) + self.q_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.k_proj = nn.Linear(context_dim, inner_dim, bias=False) + self.k_norm = nn.RMSNorm(head_dim, eps=1e-6) + self.v_proj = nn.Linear(context_dim, inner_dim, bias=False) + self.o_proj = nn.Linear(inner_dim, query_dim, bias=False) + + def forward( + self, + x: torch.Tensor, + *, + context: Optional[torch.Tensor] = None, + attn_mask: Optional[torch.Tensor] = None, + pos_q: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + pos_k: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + context = x if context is None else context + + q = self.q_proj(x).view(x.shape[0], x.shape[1], self.n_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(context).view(context.shape[0], context.shape[1], self.n_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(context).view(context.shape[0], context.shape[1], self.n_heads, self.head_dim).transpose(1, 2) + + q = self.q_norm(q) + k = self.k_norm(k) + + if pos_q is not None and pos_k is not None: + q = _apply_rope(q, *pos_q) + k = _apply_rope(k, *pos_k) + + y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + y = y.transpose(1, 2).reshape(x.shape[0], x.shape[1], -1).contiguous() + return self.o_proj(y) + + +class LLMAdapterTransformerBlock(nn.Module): + """Single transformer block in the LLM Adapter. + + Each block contains self-attention, cross-attention, and MLP with + RMSNorm pre-normalization. + """ + + def __init__( + self, + source_dim: int, + model_dim: int, + num_heads: int = 16, + ): + super().__init__() + head_dim = model_dim // num_heads + + self.norm_self_attn = nn.RMSNorm(model_dim, eps=1e-6) + self.self_attn = LLMAdapterAttention(model_dim, model_dim, num_heads, head_dim) + + self.norm_cross_attn = nn.RMSNorm(model_dim, eps=1e-6) + self.cross_attn = LLMAdapterAttention(model_dim, source_dim, num_heads, head_dim) + + self.norm_mlp = nn.RMSNorm(model_dim, eps=1e-6) + self.mlp = nn.Sequential( + nn.Linear(model_dim, model_dim * 4), + nn.GELU(), + nn.Linear(model_dim * 4, model_dim), + ) + + def forward( + self, + x: torch.Tensor, + *, + context: torch.Tensor, + target_mask: Optional[torch.Tensor] = None, + source_mask: Optional[torch.Tensor] = None, + pos_target: Tuple[torch.Tensor, torch.Tensor], + pos_source: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + x = x + self.self_attn( + self.norm_self_attn(x), + attn_mask=target_mask, + pos_q=pos_target, + pos_k=pos_target, + ) + x = x + self.cross_attn( + self.norm_cross_attn(x), + context=context, + attn_mask=source_mask, + pos_q=pos_target, + pos_k=pos_source, + ) + x = x + self.mlp(self.norm_mlp(x)) + return x + + +class LLMAdapter(nn.Module): + """LLM Adapter: bridges Qwen3 hidden states and T5-XXL token embeddings. + + Takes Qwen3 hidden states and T5-XXL token IDs, produces conditioning + embeddings for the Cosmos DiT via cross-attention through 6 transformer layers. + + Args: + vocab_size: Size of the T5 token vocabulary. + dim: Model dimension (used for embeddings, projections, and all layers). + num_layers: Number of transformer layers. + num_heads: Number of attention heads. + """ + + def __init__( + self, + vocab_size: int = 32128, + dim: int = 1024, + num_layers: int = 6, + num_heads: int = 16, + ): + super().__init__() + self.embed = nn.Embedding(vocab_size, dim) + self.blocks = nn.ModuleList( + [LLMAdapterTransformerBlock(source_dim=dim, model_dim=dim, num_heads=num_heads) for _ in range(num_layers)] + ) + self.out_proj = nn.Linear(dim, dim) + self.norm = nn.RMSNorm(dim, eps=1e-6) + self.rotary_emb = LLMAdapterRotaryEmbedding(dim // num_heads) + + def forward( + self, + source_hidden_states: torch.Tensor, + target_input_ids: torch.Tensor, + target_attention_mask: Optional[torch.Tensor] = None, + source_attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + # Expand attention masks for multi-head attention + if target_attention_mask is not None: + target_attention_mask = target_attention_mask.to(torch.bool) + if target_attention_mask.ndim == 2: + target_attention_mask = target_attention_mask[:, None, None, :] + + if source_attention_mask is not None: + source_attention_mask = source_attention_mask.to(torch.bool) + if source_attention_mask.ndim == 2: + source_attention_mask = source_attention_mask[:, None, None, :] + + context = source_hidden_states + x = self.embed(target_input_ids).to(dtype=context.dtype) + + # Build position IDs and compute rotary embeddings + target_pos_ids = torch.arange(x.shape[1], device=x.device, dtype=torch.long).unsqueeze(0) + source_pos_ids = torch.arange(context.shape[1], device=x.device, dtype=torch.long).unsqueeze(0) + pos_target = self.rotary_emb(x, target_pos_ids) + pos_source = self.rotary_emb(x, source_pos_ids) + + for block in self.blocks: + x = block( + x, + context=context, + target_mask=target_attention_mask, + source_mask=source_attention_mask, + pos_target=pos_target, + pos_source=pos_source, + ) + return self.norm(self.out_proj(x)) + + +# ============================================================================ +# Anima: MiniTrainDIT + LLMAdapter +# Reference implementation: https://github.com/hdae/diffusers-anima +# SPDX-License-Identifier: Apache-2.0 +# ============================================================================ + + +class AnimaTransformer(MiniTrainDIT): + """Anima transformer: Cosmos Predict2 DiT with integrated LLM Adapter. + + Extends MiniTrainDIT by adding the LLMAdapter component that preprocesses + text embeddings before they are fed to the DiT cross-attention layers. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.llm_adapter = LLMAdapter() + + def preprocess_text_embeds( + self, + text_embeds: torch.Tensor, + text_ids: Optional[torch.Tensor], + t5xxl_weights: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Run the LLM Adapter to produce conditioning for the DiT. + + Args: + text_embeds: Qwen3 hidden states. Shape: (batch, seq_len, 1024). + text_ids: T5-XXL token IDs. Shape: (batch, seq_len). If None, returns text_embeds directly. + t5xxl_weights: Optional per-token weights. Shape: (batch, seq_len, 1). + + Returns: + Conditioning tensor. Shape: (batch, 512, 1024), zero-padded if needed. + """ + if text_ids is None: + return text_embeds + out = self.llm_adapter(text_embeds, text_ids) + if t5xxl_weights is not None: + out = out * t5xxl_weights + if out.shape[1] < 512: + out = F.pad(out, (0, 0, 0, 512 - out.shape[1])) + return out + + def forward( + self, + x: torch.Tensor, + timesteps: torch.Tensor, + context: torch.Tensor, + t5xxl_ids: Optional[torch.Tensor] = None, + t5xxl_weights: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """Forward pass with LLM Adapter preprocessing. + + Args: + x: Input latent tensor. Shape: (B, C, T, H, W). + timesteps: Timestep values. Shape: (B,) or (B, T). + context: Qwen3 hidden states. Shape: (B, seq_len, 1024). + t5xxl_ids: T5-XXL token IDs. Shape: (B, seq_len). + t5xxl_weights: Per-token weights. Shape: (B, seq_len, 1). + + Returns: + Denoised output. Shape: (B, C, T, H, W). + """ + if t5xxl_ids is not None: + context = self.preprocess_text_embeds(context, t5xxl_ids, t5xxl_weights=t5xxl_weights) + return super().forward(x, timesteps, context, **kwargs) diff --git a/invokeai/backend/anima/anima_transformer_patch.py b/invokeai/backend/anima/anima_transformer_patch.py new file mode 100644 index 00000000000..4eff79830e9 --- /dev/null +++ b/invokeai/backend/anima/anima_transformer_patch.py @@ -0,0 +1,106 @@ +"""Utilities for patching the AnimaTransformer to support regional cross-attention masks.""" + +from contextlib import contextmanager +from typing import Optional + +import torch +import torch.nn.functional as F +from einops import rearrange + +from invokeai.backend.anima.regional_prompting import AnimaRegionalPromptingExtension + + +def _patched_cross_attn_forward( + original_forward, + attn_mask: torch.Tensor, +): + """Create a patched forward for CosmosAttention that injects a cross-attention mask. + + Args: + original_forward: The original CosmosAttention.forward method (bound to self). + attn_mask: Cross-attention mask of shape (img_seq_len, context_seq_len). + """ + + def forward(x, context=None, rope_emb=None): + # If the context sequence length doesn't match the mask (e.g. negative conditioning + # has a different number of tokens than positive regional conditioning), skip masking + # and use the original unmasked forward. + actual_context = x if context is None else context + if actual_context.shape[-2] != attn_mask.shape[1]: + return original_forward(x, context, rope_emb=rope_emb) + + self = original_forward.__self__ + + q = self.q_proj(x) + context = x if context is None else context + k = self.k_proj(context) + v = self.v_proj(context) + q, k, v = (rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim) for t in (q, k, v)) + + q = self.q_norm(q) + k = self.k_norm(k) + v = self.v_norm(v) + + if self.is_selfattn and rope_emb is not None: + from invokeai.backend.anima.anima_transformer import apply_rotary_pos_emb_cosmos + + q = apply_rotary_pos_emb_cosmos(q, rope_emb) + k = apply_rotary_pos_emb_cosmos(k, rope_emb) + + in_q_shape = q.shape + in_k_shape = k.shape + q = rearrange(q, "b ... h d -> b h ... d").reshape(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1]) + k = rearrange(k, "b ... h d -> b h ... d").reshape(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1]) + v = rearrange(v, "b ... h d -> b h ... d").reshape(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1]) + + # Convert boolean mask to float additive mask for SDPA + # True (attend) -> 0.0, False (block) -> -inf + # Shape: (img_seq_len, context_seq_len) -> (1, 1, img_seq_len, context_seq_len) + float_mask = torch.zeros_like(attn_mask, dtype=q.dtype) + float_mask[~attn_mask] = float("-inf") + expanded_mask = float_mask.unsqueeze(0).unsqueeze(0) + + result = F.scaled_dot_product_attention(q, k, v, attn_mask=expanded_mask) + result = rearrange(result, "b h s d -> b s (h d)") + return self.output_dropout(self.output_proj(result)) + + return forward + + +@contextmanager +def patch_anima_for_regional_prompting( + transformer, + regional_extension: Optional[AnimaRegionalPromptingExtension], +): + """Context manager to temporarily patch the Anima transformer for regional prompting. + + Patches the cross-attention in each DiT block to use a regional attention mask. + Uses alternating pattern: masked on even blocks, unmasked on odd blocks for + global coherence. + + Args: + transformer: The AnimaTransformer instance. + regional_extension: The regional prompting extension. If None or no mask, no patching. + + Yields: + The (possibly patched) transformer. + """ + if regional_extension is None or regional_extension.cross_attn_mask is None: + yield transformer + return + + # Store original forwards + original_forwards = [] + for block_idx, block in enumerate(transformer.blocks): + original_forwards.append(block.cross_attn.forward) + + mask = regional_extension.get_cross_attn_mask(block_idx) + if mask is not None: + block.cross_attn.forward = _patched_cross_attn_forward(block.cross_attn.forward, mask) + + try: + yield transformer + finally: + # Restore original forwards + for block_idx, block in enumerate(transformer.blocks): + block.cross_attn.forward = original_forwards[block_idx] diff --git a/invokeai/backend/anima/conditioning_data.py b/invokeai/backend/anima/conditioning_data.py new file mode 100644 index 00000000000..b96c807835d --- /dev/null +++ b/invokeai/backend/anima/conditioning_data.py @@ -0,0 +1,64 @@ +"""Anima text conditioning data structures. + +Anima uses a dual-conditioning scheme: +- Qwen3 0.6B hidden states (continuous embeddings) +- T5-XXL token IDs (discrete IDs, embedded by the LLM Adapter inside the transformer) + +Both are produced by the text encoder invocation and stored together. + +For regional prompting, multiple conditionings (each with an optional spatial mask) +are concatenated and processed together. The LLM Adapter runs on each region's +conditioning separately, producing per-region context vectors that are concatenated +for the DiT's cross-attention layers. An attention mask restricts which image tokens +attend to which regional context tokens. +""" + +from dataclasses import dataclass + +import torch + +from invokeai.backend.stable_diffusion.diffusion.conditioning_data import Range + + +@dataclass +class AnimaTextConditioning: + """Anima text conditioning with Qwen3 hidden states, T5-XXL token IDs, and optional mask. + + Attributes: + qwen3_embeds: Text embeddings from Qwen3 0.6B encoder. + Shape: (seq_len, hidden_size) where hidden_size=1024. + t5xxl_ids: T5-XXL token IDs for the same prompt. + Shape: (seq_len,). + t5xxl_weights: Per-token weights for prompt weighting. + Shape: (seq_len,). Defaults to all ones if not provided. + mask: Optional binary mask for regional prompting. If None, the prompt is global. + Shape: (1, 1, img_seq_len) where img_seq_len = (H // patch_size) * (W // patch_size). + """ + + qwen3_embeds: torch.Tensor + t5xxl_ids: torch.Tensor + t5xxl_weights: torch.Tensor | None = None + mask: torch.Tensor | None = None + + +@dataclass +class AnimaRegionalTextConditioning: + """Container for multiple regional text conditionings processed by the LLM Adapter. + + After the LLM Adapter processes each region's conditioning, the outputs are concatenated. + The DiT cross-attention then uses an attention mask to restrict which image tokens + attend to which region's context tokens. + + Attributes: + context_embeds: Concatenated LLM Adapter outputs from all regional prompts. + Shape: (total_context_len, 1024). + image_masks: List of binary masks for each regional prompt. + If None, the prompt is global (applies to entire image). + Shape: (1, 1, img_seq_len). + context_ranges: List of ranges indicating which portion of context_embeds + corresponds to each regional prompt. + """ + + context_embeds: torch.Tensor + image_masks: list[torch.Tensor | None] + context_ranges: list[Range] diff --git a/invokeai/backend/anima/regional_prompting.py b/invokeai/backend/anima/regional_prompting.py new file mode 100644 index 00000000000..c0af366332f --- /dev/null +++ b/invokeai/backend/anima/regional_prompting.py @@ -0,0 +1,173 @@ +"""Regional prompting extension for Anima. + +Anima's architecture uses separate cross-attention in each DiT block: image tokens +(in 5D spatial layout) cross-attend to context tokens (LLM Adapter output). This is +different from Z-Image's unified [img, txt] sequence with self-attention. + +For regional prompting, we: +1. Run the LLM Adapter separately for each regional prompt +2. Concatenate the resulting context vectors +3. Build a cross-attention mask that restricts each image region to attend only to + its corresponding context tokens +4. Patch the DiT's cross-attention to use this mask + +The mask alternation strategy (masked on even blocks, full on odd blocks) helps +maintain global coherence across regions. +""" + +from typing import Optional + +import torch +import torchvision + +from invokeai.backend.anima.conditioning_data import AnimaRegionalTextConditioning +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.mask import to_standard_float_mask + + +class AnimaRegionalPromptingExtension: + """Manages regional prompting for Anima's cross-attention. + + Unlike Z-Image which uses a unified [img, txt] sequence, Anima has separate + cross-attention where image tokens (query) attend to context tokens (key/value). + The cross-attention mask shape is (img_seq_len, context_seq_len). + """ + + def __init__( + self, + regional_text_conditioning: AnimaRegionalTextConditioning, + cross_attn_mask: torch.Tensor | None = None, + ): + self.regional_text_conditioning = regional_text_conditioning + self.cross_attn_mask = cross_attn_mask + + def get_cross_attn_mask(self, block_index: int) -> torch.Tensor | None: + """Get the cross-attention mask for a given block index. + + Uses alternating pattern: apply mask on even blocks, no mask on odd blocks. + This helps balance regional control with global coherence. + """ + if block_index % 2 == 0: + return self.cross_attn_mask + return None + + @classmethod + def from_regional_conditioning( + cls, + regional_text_conditioning: AnimaRegionalTextConditioning, + img_seq_len: int, + ) -> "AnimaRegionalPromptingExtension": + """Create extension from pre-processed regional conditioning. + + Args: + regional_text_conditioning: Regional conditioning with concatenated context and masks. + img_seq_len: Number of image tokens (H_patches * W_patches). + """ + cross_attn_mask = cls._prepare_cross_attn_mask(regional_text_conditioning, img_seq_len) + return cls( + regional_text_conditioning=regional_text_conditioning, + cross_attn_mask=cross_attn_mask, + ) + + @classmethod + def _prepare_cross_attn_mask( + cls, + regional_text_conditioning: AnimaRegionalTextConditioning, + img_seq_len: int, + ) -> torch.Tensor | None: + """Prepare a cross-attention mask for regional prompting. + + The mask shape is (img_seq_len, context_seq_len) where: + - Each image token can attend to context tokens from its assigned region + - Global prompts (mask=None) attend to background regions + + Args: + regional_text_conditioning: The regional text conditioning data. + img_seq_len: Number of image tokens. + + Returns: + Cross-attention mask of shape (img_seq_len, context_seq_len), or None + if no regional masks are present. + """ + has_regional_masks = any(mask is not None for mask in regional_text_conditioning.image_masks) + if not has_regional_masks: + return None + + # Identify background region (area not covered by any mask) + background_region_mask: torch.Tensor | None = None + for image_mask in regional_text_conditioning.image_masks: + if image_mask is not None: + mask_flat = image_mask.view(-1) + if background_region_mask is None: + background_region_mask = torch.ones_like(mask_flat) + background_region_mask = background_region_mask * (1 - mask_flat) + + device = TorchDevice.choose_torch_device() + context_seq_len = regional_text_conditioning.context_embeds.shape[0] + + # Cross-attention mask: (img_seq_len, context_seq_len) + # img tokens are queries, context tokens are keys/values + cross_attn_mask = torch.zeros((img_seq_len, context_seq_len), device=device, dtype=torch.float16) + + for image_mask, context_range in zip( + regional_text_conditioning.image_masks, + regional_text_conditioning.context_ranges, + strict=True, + ): + ctx_start = context_range.start + ctx_end = context_range.end + + if image_mask is not None: + # Regional prompt: only masked image tokens attend to this region's context + mask_flat = image_mask.view(img_seq_len) + cross_attn_mask[:, ctx_start:ctx_end] = mask_flat.view(img_seq_len, 1) + else: + # Global prompt: background image tokens attend to this context + if background_region_mask is not None: + cross_attn_mask[:, ctx_start:ctx_end] = background_region_mask.view(img_seq_len, 1) + else: + cross_attn_mask[:, ctx_start:ctx_end] = 1.0 + + # Convert to boolean + cross_attn_mask = cross_attn_mask > 0.5 + return cross_attn_mask + + @staticmethod + def preprocess_regional_prompt_mask( + mask: Optional[torch.Tensor], + target_height: int, + target_width: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Preprocess a regional prompt mask to match the target image token grid. + + Args: + mask: Input mask tensor. If None, returns a mask of all ones. + target_height: Height of the image token grid (H // patch_size). + target_width: Width of the image token grid (W // patch_size). + dtype: Target dtype for the mask. + device: Target device for the mask. + + Returns: + Processed mask of shape (1, 1, target_height * target_width). + """ + img_seq_len = target_height * target_width + + if mask is None: + return torch.ones((1, 1, img_seq_len), dtype=dtype, device=device) + + mask = to_standard_float_mask(mask, out_dtype=dtype) + + tf = torchvision.transforms.Resize( + (target_height, target_width), + interpolation=torchvision.transforms.InterpolationMode.NEAREST, + ) + + if mask.ndim == 2: + mask = mask.unsqueeze(0) + if mask.ndim == 3: + mask = mask.unsqueeze(0) + + resized_mask = tf(mask) + return resized_mask.flatten(start_dim=2).to(device=device) diff --git a/invokeai/backend/flux/modules/conditioner.py b/invokeai/backend/flux/modules/conditioner.py index 9deb4429292..d48d78cd4a1 100644 --- a/invokeai/backend/flux/modules/conditioner.py +++ b/invokeai/backend/flux/modules/conditioner.py @@ -3,6 +3,8 @@ from torch import Tensor, nn from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast +from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device + class HFEncoder(nn.Module): def __init__( @@ -32,7 +34,7 @@ def forward(self, text: list[str]) -> Tensor: ) # Move inputs to the same device as the model to support cpu_only models - model_device = next(self.hf_module.parameters()).device + model_device = get_effective_device(self.hf_module) outputs = self.hf_module( input_ids=batch_encoding["input_ids"].to(model_device), diff --git a/invokeai/backend/flux/schedulers.py b/invokeai/backend/flux/schedulers.py index e5a8a7137c2..05e6bb085f0 100644 --- a/invokeai/backend/flux/schedulers.py +++ b/invokeai/backend/flux/schedulers.py @@ -60,3 +60,23 @@ if _HAS_LCM: ZIMAGE_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler + + +# Anima scheduler types (same Flow Matching schedulers as Flux/Z-Image) +# Anima uses rectified flow with shift=3.0 and multiplier=1000. +# Recommended: 30 steps with Euler, CFG 4-5. +ANIMA_SCHEDULER_NAME_VALUES = Literal["euler", "heun", "lcm"] + +ANIMA_SCHEDULER_LABELS: dict[str, str] = { + "euler": "Euler", + "heun": "Heun (2nd order)", + "lcm": "LCM", +} + +ANIMA_SCHEDULER_MAP: dict[str, Type[SchedulerMixin]] = { + "euler": FlowMatchEulerDiscreteScheduler, + "heun": FlowMatchHeunDiscreteScheduler, +} + +if _HAS_LCM: + ANIMA_SCHEDULER_MAP["lcm"] = FlowMatchLCMScheduler diff --git a/invokeai/backend/model_manager/configs/factory.py b/invokeai/backend/model_manager/configs/factory.py index 7702d4a5d9b..5a376cfad95 100644 --- a/invokeai/backend/model_manager/configs/factory.py +++ b/invokeai/backend/model_manager/configs/factory.py @@ -46,8 +46,10 @@ LoRA_Diffusers_SD2_Config, LoRA_Diffusers_SDXL_Config, LoRA_Diffusers_ZImage_Config, + LoRA_LyCORIS_Anima_Config, LoRA_LyCORIS_Flux2_Config, LoRA_LyCORIS_FLUX_Config, + LoRA_LyCORIS_QwenImage_Config, LoRA_LyCORIS_SD1_Config, LoRA_LyCORIS_SD2_Config, LoRA_LyCORIS_SDXL_Config, @@ -58,6 +60,7 @@ ) from invokeai.backend.model_manager.configs.main import ( Main_BnBNF4_FLUX_Config, + Main_Checkpoint_Anima_Config, Main_Checkpoint_Flux2_Config, Main_Checkpoint_FLUX_Config, Main_Checkpoint_SD1_Config, @@ -68,6 +71,7 @@ Main_Diffusers_CogView4_Config, Main_Diffusers_Flux2_Config, Main_Diffusers_FLUX_Config, + Main_Diffusers_QwenImage_Config, Main_Diffusers_SD1_Config, Main_Diffusers_SD2_Config, Main_Diffusers_SD3_Config, @@ -76,6 +80,7 @@ Main_Diffusers_ZImage_Config, Main_GGUF_Flux2_Config, Main_GGUF_FLUX_Config, + Main_GGUF_QwenImage_Config, Main_GGUF_ZImage_Config, MainModelDefaultSettings, ) @@ -101,6 +106,7 @@ ) from invokeai.backend.model_manager.configs.unknown import Unknown_Config from invokeai.backend.model_manager.configs.vae import ( + VAE_Checkpoint_Anima_Config, VAE_Checkpoint_Flux2_Config, VAE_Checkpoint_FLUX_Config, VAE_Checkpoint_SD1_Config, @@ -159,6 +165,7 @@ Annotated[Main_Diffusers_FLUX_Config, Main_Diffusers_FLUX_Config.get_tag()], Annotated[Main_Diffusers_Flux2_Config, Main_Diffusers_Flux2_Config.get_tag()], Annotated[Main_Diffusers_CogView4_Config, Main_Diffusers_CogView4_Config.get_tag()], + Annotated[Main_Diffusers_QwenImage_Config, Main_Diffusers_QwenImage_Config.get_tag()], Annotated[Main_Diffusers_ZImage_Config, Main_Diffusers_ZImage_Config.get_tag()], # Main (Pipeline) - checkpoint format # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation @@ -170,12 +177,14 @@ Annotated[Main_Checkpoint_Flux2_Config, Main_Checkpoint_Flux2_Config.get_tag()], Annotated[Main_Checkpoint_FLUX_Config, Main_Checkpoint_FLUX_Config.get_tag()], Annotated[Main_Checkpoint_ZImage_Config, Main_Checkpoint_ZImage_Config.get_tag()], + Annotated[Main_Checkpoint_Anima_Config, Main_Checkpoint_Anima_Config.get_tag()], # Main (Pipeline) - quantized formats # IMPORTANT: FLUX.2 must be checked BEFORE FLUX.1 because FLUX.2 has specific validation # that will reject FLUX.1 models, but FLUX.1 validation may incorrectly match FLUX.2 models Annotated[Main_BnBNF4_FLUX_Config, Main_BnBNF4_FLUX_Config.get_tag()], Annotated[Main_GGUF_Flux2_Config, Main_GGUF_Flux2_Config.get_tag()], Annotated[Main_GGUF_FLUX_Config, Main_GGUF_FLUX_Config.get_tag()], + Annotated[Main_GGUF_QwenImage_Config, Main_GGUF_QwenImage_Config.get_tag()], Annotated[Main_GGUF_ZImage_Config, Main_GGUF_ZImage_Config.get_tag()], # VAE - checkpoint format Annotated[VAE_Checkpoint_SD1_Config, VAE_Checkpoint_SD1_Config.get_tag()], @@ -183,6 +192,7 @@ Annotated[VAE_Checkpoint_SDXL_Config, VAE_Checkpoint_SDXL_Config.get_tag()], Annotated[VAE_Checkpoint_FLUX_Config, VAE_Checkpoint_FLUX_Config.get_tag()], Annotated[VAE_Checkpoint_Flux2_Config, VAE_Checkpoint_Flux2_Config.get_tag()], + Annotated[VAE_Checkpoint_Anima_Config, VAE_Checkpoint_Anima_Config.get_tag()], # VAE - diffusers format Annotated[VAE_Diffusers_SD1_Config, VAE_Diffusers_SD1_Config.get_tag()], Annotated[VAE_Diffusers_SDXL_Config, VAE_Diffusers_SDXL_Config.get_tag()], @@ -207,6 +217,8 @@ Annotated[LoRA_LyCORIS_Flux2_Config, LoRA_LyCORIS_Flux2_Config.get_tag()], Annotated[LoRA_LyCORIS_FLUX_Config, LoRA_LyCORIS_FLUX_Config.get_tag()], Annotated[LoRA_LyCORIS_ZImage_Config, LoRA_LyCORIS_ZImage_Config.get_tag()], + Annotated[LoRA_LyCORIS_QwenImage_Config, LoRA_LyCORIS_QwenImage_Config.get_tag()], + Annotated[LoRA_LyCORIS_Anima_Config, LoRA_LyCORIS_Anima_Config.get_tag()], # LoRA - OMI format Annotated[LoRA_OMI_SDXL_Config, LoRA_OMI_SDXL_Config.get_tag()], Annotated[LoRA_OMI_FLUX_Config, LoRA_OMI_FLUX_Config.get_tag()], diff --git a/invokeai/backend/model_manager/configs/lora.py b/invokeai/backend/model_manager/configs/lora.py index 1619c9d6f06..88f917d0d3b 100644 --- a/invokeai/backend/model_manager/configs/lora.py +++ b/invokeai/backend/model_manager/configs/lora.py @@ -31,6 +31,10 @@ ZImageVariantType, ) from invokeai.backend.model_manager.util.model_util import lora_token_vector_length +from invokeai.backend.patches.lora_conversions.anima_lora_constants import ( + has_cosmos_dit_kohya_keys, + has_cosmos_dit_peft_keys, +) from invokeai.backend.patches.lora_conversions.flux_control_lora_utils import is_state_dict_likely_flux_control @@ -637,6 +641,13 @@ def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: return BaseModelType.Flux state_dict = mod.load_state_dict() + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + + # Rule out Anima LoRAs — their lora_te_ keys have shapes that + # lora_token_vector_length() misidentifies as SD2/SDXL. + if has_cosmos_dit_kohya_keys(str_keys) or has_cosmos_dit_peft_keys(str_keys): + raise NotAMatchError("model looks like an Anima LoRA, not a Stable Diffusion LoRA") + # If we've gotten here, we assume that the model is a Stable Diffusion model token_vector_length = lora_token_vector_length(state_dict) if token_vector_length == 768: @@ -711,6 +722,8 @@ def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None: state_dict, { "diffusion_model.layers.", # Z-Image S3-DiT layer pattern + "transformer.layers.", # OneTrainer/diffusers prefix variant + "base_model.model.transformer.layers.", # PEFT-wrapped variant }, ) @@ -747,6 +760,8 @@ def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: state_dict, { "diffusion_model.layers.", # Z-Image S3-DiT layer pattern + "transformer.layers.", # OneTrainer/diffusers prefix variant + "base_model.model.transformer.layers.", # PEFT-wrapped variant }, ) @@ -757,6 +772,142 @@ def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: raise NotAMatchError("model does not look like a Z-Image LoRA") +class LoRA_LyCORIS_QwenImage_Config(LoRA_LyCORIS_Config_Base, Config_Base): + """Model config for Qwen Image Edit LoRA models in LyCORIS format.""" + + base: Literal[BaseModelType.QwenImage] = Field(default=BaseModelType.QwenImage) + + @classmethod + def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None: + """Qwen Image Edit LoRAs have keys like transformer_blocks.X.attn.to_k.lora_down.weight.""" + state_dict = mod.load_state_dict() + + has_qwen_ie_keys = state_dict_has_any_keys_starting_with( + state_dict, + { + "transformer_blocks.", + "transformer.transformer_blocks.", + "lora_unet_transformer_blocks_", # Kohya format + }, + ) + has_lora_suffix = state_dict_has_any_keys_ending_with( + state_dict, + { + "lora_A.weight", + "lora_B.weight", + "lora_down.weight", + "lora_up.weight", + "dora_scale", + "lokr_w1", + "lokr_w2", # LoKR format + }, + ) + # Must NOT have diffusion_model.layers (Z-Image) or Flux-style keys. + # Flux LoRAs can have transformer.single_transformer_blocks or transformer.transformer_blocks + # (with the "transformer." prefix and "single_" variant) which would falsely match our check. + # Flux Kohya LoRAs use lora_unet_double_blocks or lora_unet_single_blocks. + has_z_image_keys = state_dict_has_any_keys_starting_with(state_dict, {"diffusion_model.layers."}) + has_flux_keys = state_dict_has_any_keys_starting_with( + state_dict, + { + "double_blocks.", + "single_blocks.", + "single_transformer_blocks.", + "transformer.single_transformer_blocks.", + "lora_unet_double_blocks_", + "lora_unet_single_blocks_", + "lora_unet_single_transformer_blocks_", + }, + ) + + if has_qwen_ie_keys and has_lora_suffix and not has_z_image_keys and not has_flux_keys: + return + + raise NotAMatchError("model does not match Qwen Image LoRA heuristics") + + @classmethod + def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: + state_dict = mod.load_state_dict() + has_qwen_ie_keys = state_dict_has_any_keys_starting_with( + state_dict, + {"transformer_blocks.", "transformer.transformer_blocks.", "lora_unet_transformer_blocks_"}, + ) + has_z_image_keys = state_dict_has_any_keys_starting_with(state_dict, {"diffusion_model.layers."}) + has_flux_keys = state_dict_has_any_keys_starting_with( + state_dict, + { + "double_blocks.", + "single_blocks.", + "single_transformer_blocks.", + "transformer.single_transformer_blocks.", + "lora_unet_double_blocks_", + "lora_unet_single_blocks_", + "lora_unet_single_transformer_blocks_", + }, + ) + + if has_qwen_ie_keys and not has_z_image_keys and not has_flux_keys: + return BaseModelType.QwenImage + raise NotAMatchError("model does not look like a Qwen Image Edit LoRA") + + +class LoRA_LyCORIS_Anima_Config(LoRA_LyCORIS_Config_Base, Config_Base): + """Model config for Anima LoRA models in LyCORIS format.""" + + base: Literal[BaseModelType.Anima] = Field(default=BaseModelType.Anima) + + @classmethod + def _validate_looks_like_lora(cls, mod: ModelOnDisk) -> None: + """Anima LoRAs use Kohya-style keys targeting Cosmos DiT blocks. + + Anima LoRAs have keys like: + - lora_unet_blocks_0_cross_attn_k_proj.lora_down.weight (Kohya format) + - diffusion_model.blocks.0.cross_attn.k_proj.lora_A.weight (diffusers PEFT format) + - transformer.blocks.0.cross_attn.k_proj.lora_A.weight (diffusers PEFT format) + + Detection requires Cosmos DiT-specific subcomponent names (cross_attn, + self_attn, mlp, adaln_modulation) to avoid false-positives on other + architectures that also use ``blocks`` in their paths. + """ + state_dict = mod.load_state_dict() + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + + has_cosmos_keys = has_cosmos_dit_kohya_keys(str_keys) or has_cosmos_dit_peft_keys(str_keys) + + # Also check for LoRA/LoKR weight suffixes + has_lora_suffix = state_dict_has_any_keys_ending_with( + state_dict, + { + "lora_A.weight", + "lora_B.weight", + "lora_down.weight", + "lora_up.weight", + "dora_scale", + ".lokr_w1", + ".lokr_w2", + }, + ) + + if has_cosmos_keys and has_lora_suffix: + return + + raise NotAMatchError("model does not match Anima LoRA heuristics") + + @classmethod + def _get_base_or_raise(cls, mod: ModelOnDisk) -> BaseModelType: + """Anima LoRAs target Cosmos DiT blocks (blocks.X.cross_attn, blocks.X.self_attn, etc.). + + Uses Cosmos DiT-specific subcomponent names to avoid false-positives. + """ + state_dict = mod.load_state_dict() + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + + if has_cosmos_dit_kohya_keys(str_keys) or has_cosmos_dit_peft_keys(str_keys): + return BaseModelType.Anima + + raise NotAMatchError("model does not look like an Anima LoRA") + + class ControlAdapter_Config_Base(ABC, BaseModel): default_settings: ControlAdapterDefaultSettings | None = Field(None) diff --git a/invokeai/backend/model_manager/configs/main.py b/invokeai/backend/model_manager/configs/main.py index 6f737ceb92d..1be349f3942 100644 --- a/invokeai/backend/model_manager/configs/main.py +++ b/invokeai/backend/model_manager/configs/main.py @@ -28,6 +28,7 @@ ModelFormat, ModelType, ModelVariantType, + QwenImageVariantType, SchedulerPredictionType, SubModelType, ZImageVariantType, @@ -76,6 +77,8 @@ def from_base( else: # Turbo (distilled) uses fewer steps, no CFG return cls(steps=9, cfg_scale=1.0, width=1024, height=1024) + case BaseModelType.Anima: + return cls(steps=35, cfg_scale=4.5, width=1024, height=1024) case BaseModelType.Flux2: # Different defaults based on variant if variant == Flux2VariantType.Klein9BBase: @@ -84,6 +87,8 @@ def from_base( else: # Distilled models (Klein 4B, Klein 9B) use fewer steps return cls(steps=4, cfg_scale=1.0, width=1024, height=1024) + case BaseModelType.QwenImage: + return cls(steps=40, cfg_scale=4.0, width=1024, height=1024) case _: # TODO(psyche): Do we want defaults for other base types? return None @@ -194,9 +199,11 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_base(mod) - prediction_type = override_fields.get("prediction_type") or cls._get_scheduler_prediction_type_or_raise(mod) + prediction_type = override_fields.pop("prediction_type", None) or cls._get_scheduler_prediction_type_or_raise( + mod + ) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) return cls(**override_fields, prediction_type=prediction_type, variant=variant) @@ -323,6 +330,16 @@ def _is_flux2_model(state_dict: dict[str | int, Any]) -> bool: return False +def _filename_suggests_base(name: str) -> bool: + """Check if a model name/filename suggests it is a Base (undistilled) variant. + + Klein 9B Base and Klein 9B have identical architectures and cannot be distinguished + from the state dict. We use the filename as a heuristic: filenames containing "base" + (e.g. "flux-2-klein-base-9b", "FLUX.2-klein-base-9B") indicate the undistilled model. + """ + return "base" in name.lower() + + def _get_flux2_variant(state_dict: dict[str | int, Any]) -> Flux2VariantType | None: """Determine FLUX.2 variant from state dict. @@ -330,9 +347,9 @@ def _get_flux2_variant(state_dict: dict[str | int, Any]) -> Flux2VariantType | N - Klein 4B: context_in_dim = 7680 (3 × Qwen3-4B hidden_size 2560) - Klein 9B: context_in_dim = 12288 (3 × Qwen3-8B hidden_size 4096) - Note: Klein 9B Base (undistilled) also has context_in_dim = 12288 but is rare. - We default to Klein9B (distilled) for all 9B models since GGUF models may not - include guidance embedding keys needed to distinguish them. + Note: Klein 9B (distilled) and Klein 9B Base (undistilled) have identical architectures + and cannot be distinguished from the state dict alone. This function defaults to Klein9B + for all 9B models. Callers should use filename heuristics to detect Klein9BBase. Supports both BFL format (checkpoint) and diffusers format keys: - BFL format: txt_in.weight (context embedder) @@ -366,7 +383,7 @@ def _get_flux2_variant(state_dict: dict[str | int, Any]) -> Flux2VariantType | N context_in_dim = shape[1] # Determine variant based on context dimension if context_in_dim == KLEIN_9B_CONTEXT_DIM: - # Default to Klein9B (distilled) - the official/common 9B model + # Default to Klein9B - callers use filename heuristics to detect Klein9BBase return Flux2VariantType.Klein9B elif context_in_dim == KLEIN_4B_CONTEXT_DIM: return Flux2VariantType.Klein4B @@ -459,7 +476,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_does_not_look_like_gguf_quantized(mod) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) return cls(**override_fields, variant=variant) @@ -534,7 +551,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_does_not_look_like_gguf_quantized(mod) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) return cls(**override_fields, variant=variant) @@ -553,6 +570,11 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: if variant is None: raise NotAMatchError("unable to determine FLUX.2 model variant from state dict") + # Klein 9B Base and Klein 9B have identical architectures. + # Use filename heuristic to detect the Base (undistilled) variant. + if variant == Flux2VariantType.Klein9B and _filename_suggests_base(mod.name): + return Flux2VariantType.Klein9BBase + return variant @classmethod @@ -592,7 +614,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_model_looks_like_bnb_quantized(mod) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) return cls(**override_fields, variant=variant) @@ -643,7 +665,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_is_not_flux2(mod) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) return cls(**override_fields, variant=variant) @@ -701,7 +723,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_is_flux2(mod) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) return cls(**override_fields, variant=variant) @@ -720,6 +742,11 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: if variant is None: raise NotAMatchError("unable to determine FLUX.2 model variant from state dict") + # Klein 9B Base and Klein 9B have identical architectures. + # Use filename heuristic to detect the Base (undistilled) variant. + if variant == Flux2VariantType.Klein9B and _filename_suggests_base(mod.name): + return Flux2VariantType.Klein9BBase + return variant @classmethod @@ -757,9 +784,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - }, ) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) - repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod) + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) return cls( **override_fields, @@ -811,9 +838,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - }, ) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) - repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod) + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) return cls( **override_fields, @@ -829,12 +856,8 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: - Klein 4B: joint_attention_dim = 7680 (3×Qwen3-4B hidden size) - Klein 9B/9B Base: joint_attention_dim = 12288 (3×Qwen3-8B hidden size) - To distinguish Klein 9B (distilled) from Klein 9B Base (undistilled), - we check guidance_embeds: - - Klein 9B (distilled): guidance_embeds = False (guidance is "baked in" during distillation) - - Klein 9B Base (undistilled): guidance_embeds = True (needs guidance at inference) - - Note: The official BFL Klein 9B model is the distilled version with guidance_embeds=False. + Klein 9B (distilled) and Klein 9B Base (undistilled) have identical architectures + and both have guidance_embeds=False. We use a filename heuristic to detect Base models. """ KLEIN_4B_CONTEXT_DIM = 7680 # 3 × 2560 KLEIN_9B_CONTEXT_DIM = 12288 # 3 × 4096 @@ -842,17 +865,12 @@ def _get_variant_or_raise(cls, mod: ModelOnDisk) -> Flux2VariantType: transformer_config = get_config_dict_or_raise(mod.path / "transformer" / "config.json") joint_attention_dim = transformer_config.get("joint_attention_dim", 4096) - guidance_embeds = transformer_config.get("guidance_embeds", False) # Determine variant based on joint_attention_dim if joint_attention_dim == KLEIN_9B_CONTEXT_DIM: - # Check guidance_embeds to distinguish distilled from undistilled - # Klein 9B (distilled): guidance_embeds = False (guidance is baked in) - # Klein 9B Base (undistilled): guidance_embeds = True (needs guidance) - if guidance_embeds: + if _filename_suggests_base(mod.name): return Flux2VariantType.Klein9BBase - else: - return Flux2VariantType.Klein9B + return Flux2VariantType.Klein9B elif joint_attention_dim == KLEIN_4B_CONTEXT_DIM: return Flux2VariantType.Klein4B elif joint_attention_dim > 4096: @@ -891,11 +909,13 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_base(mod) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) - prediction_type = override_fields.get("prediction_type") or cls._get_scheduler_prediction_type_or_raise(mod) + prediction_type = override_fields.pop("prediction_type", None) or cls._get_scheduler_prediction_type_or_raise( + mod + ) - repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod) + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) return cls( **override_fields, @@ -1001,9 +1021,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - }, ) - submodels = override_fields.get("submodels") or cls._get_submodels_or_raise(mod) + submodels = override_fields.pop("submodels", None) or cls._get_submodels_or_raise(mod) - repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod) + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) return cls( **override_fields, @@ -1076,7 +1096,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - }, ) - repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod) + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) return cls( **override_fields, @@ -1084,6 +1104,44 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - ) +def _has_anima_keys(state_dict: dict[str | int, Any]) -> bool: + """Check if state dict contains Anima model keys. + + Anima models are identified by the presence of `llm_adapter` keys + (unique to Anima - the LLM Adapter that bridges Qwen3 text encoder to the Cosmos DiT) + alongside Cosmos Predict2 DiT keys (blocks, t_embedder, x_embedder, final_layer). + + The checkpoint keys may have a `net.` prefix (e.g. `net.llm_adapter.`, `net.blocks.`). + """ + has_llm_adapter = False + has_cosmos_dit = False + + # Cosmos DiT key prefixes — support both with and without `net.` prefix + cosmos_prefixes = ( + "blocks.", + "t_embedder.", + "x_embedder.", + "final_layer.", + "net.blocks.", + "net.t_embedder.", + "net.x_embedder.", + "net.final_layer.", + ) + + for key in state_dict.keys(): + if isinstance(key, int): + continue + if key.startswith("llm_adapter.") or key.startswith("net.llm_adapter."): + has_llm_adapter = True + for prefix in cosmos_prefixes: + if key.startswith(prefix): + has_cosmos_dit = True + if has_llm_adapter and has_cosmos_dit: + return True + + return False + + class Main_Diffusers_ZImage_Config(Diffusers_Config_Base, Main_Config_Base, Config_Base): """Model config for Z-Image diffusers models (Z-Image-Turbo, Z-Image-Base).""" @@ -1104,9 +1162,9 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - }, ) - variant = override_fields.get("variant") or cls._get_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_variant_or_raise(mod) - repo_variant = override_fields.get("repo_variant") or cls._get_repo_variant_or_raise(mod) + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) return cls( **override_fields, @@ -1150,7 +1208,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_does_not_look_like_gguf_quantized(mod) - variant = override_fields.get("variant", ZImageVariantType.Turbo) + variant = override_fields.pop("variant", None) or ZImageVariantType.Turbo return cls(**override_fields, variant=variant) @@ -1184,7 +1242,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - cls._validate_looks_like_gguf_quantized(mod) - variant = override_fields.get("variant", ZImageVariantType.Turbo) + variant = override_fields.pop("variant", None) or ZImageVariantType.Turbo return cls(**override_fields, variant=variant) @@ -1199,3 +1257,130 @@ def _validate_looks_like_gguf_quantized(cls, mod: ModelOnDisk) -> None: has_ggml_tensors = _has_ggml_tensors(mod.load_state_dict()) if not has_ggml_tensors: raise NotAMatchError("state dict does not look like GGUF quantized") + + +class Main_Diffusers_QwenImage_Config(Diffusers_Config_Base, Main_Config_Base, Config_Base): + """Model config for Qwen Image diffusers models (both txt2img and edit).""" + + base: Literal[BaseModelType.QwenImage] = Field(BaseModelType.QwenImage) + variant: QwenImageVariantType | None = Field(default=None) + + @classmethod + def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: + raise_if_not_dir(mod) + + raise_for_override_fields(cls, override_fields) + + # This check implies the base type - no further validation needed. + raise_for_class_name( + common_config_paths(mod.path), + { + "QwenImagePlusPipeline", + "QwenImageEditPlusPipeline", + "QwenImagePipeline", + }, + ) + + repo_variant = override_fields.pop("repo_variant", None) or cls._get_repo_variant_or_raise(mod) + variant = override_fields.pop("variant", None) or cls._get_qwen_image_variant(mod) + + return cls( + **override_fields, + repo_variant=repo_variant, + variant=variant, + ) + + @classmethod + def _get_qwen_image_variant(cls, mod: ModelOnDisk) -> QwenImageVariantType: + """Detect whether this is an edit or txt2img model from the pipeline class name.""" + import json + + model_index = mod.path / "model_index.json" + if model_index.exists(): + with open(model_index) as f: + config = json.load(f) + class_name = config.get("_class_name", "") + if "Edit" in class_name: + return QwenImageVariantType.Edit + return QwenImageVariantType.Generate + + +def _has_qwen_image_keys(state_dict: dict[str | int, Any]) -> bool: + """Check if state dict contains Qwen Image Edit transformer keys. + + Qwen Image Edit uses 'txt_in' and 'txt_norm' instead of 'context_embedder' (FLUX). + This distinguishes it from FLUX and other architectures. + """ + has_txt_in = any(isinstance(k, str) and k.startswith("txt_in.") for k in state_dict.keys()) + has_txt_norm = any(isinstance(k, str) and k.startswith("txt_norm.") for k in state_dict.keys()) + has_img_in = any(isinstance(k, str) and k.startswith("img_in.") for k in state_dict.keys()) + # Must NOT have context_embedder (which would indicate FLUX) + has_context_embedder = any(isinstance(k, str) and "context_embedder" in k for k in state_dict.keys()) + return has_txt_in and has_txt_norm and has_img_in and not has_context_embedder + + +class Main_GGUF_QwenImage_Config(Checkpoint_Config_Base, Main_Config_Base, Config_Base): + """Model config for GGUF-quantized Qwen Image transformer models.""" + + base: Literal[BaseModelType.QwenImage] = Field(default=BaseModelType.QwenImage) + format: Literal[ModelFormat.GGUFQuantized] = Field(default=ModelFormat.GGUFQuantized) + variant: QwenImageVariantType | None = Field(default=None) + + @classmethod + def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: + raise_if_not_file(mod) + + raise_for_override_fields(cls, override_fields) + + sd = mod.load_state_dict() + + if not _has_qwen_image_keys(sd): + raise NotAMatchError("state dict does not look like a Qwen Image Edit model") + + if not _has_ggml_tensors(sd): + raise NotAMatchError("state dict does not look like GGUF quantized") + + # Infer variant from the state dict if not explicitly provided. + # The Edit variant includes an extra tensor `__index_timestep_zero__` (used by the + # `zero_cond_t` dual-modulation path in diffusers' QwenImageTransformer2DModel). + # If the marker tensor is missing, fall back to the filename heuristic since older + # or alternate GGUF converters may not emit it. + explicit_variant = override_fields.pop("variant", None) + if explicit_variant is None: + if "__index_timestep_zero__" in sd: + explicit_variant = QwenImageVariantType.Edit + else: + filename = mod.path.stem.lower() + if "edit" in filename: + explicit_variant = QwenImageVariantType.Edit + else: + explicit_variant = QwenImageVariantType.Generate + + return cls(**override_fields, variant=explicit_variant) + + +class Main_Checkpoint_Anima_Config(Checkpoint_Config_Base, Main_Config_Base, Config_Base): + """Model config for Anima single-file checkpoint models (safetensors). + + Anima is built on NVIDIA Cosmos Predict2 DiT with a custom LLM Adapter + that bridges Qwen3 0.6B text encoder outputs to the DiT. + """ + + base: Literal[BaseModelType.Anima] = Field(default=BaseModelType.Anima) + format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) + + @classmethod + def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: + raise_if_not_file(mod) + + raise_for_override_fields(cls, override_fields) + + cls._validate_looks_like_anima_model(mod) + + return cls(**override_fields) + + @classmethod + def _validate_looks_like_anima_model(cls, mod: ModelOnDisk) -> None: + has_anima_keys = _has_anima_keys(mod.load_state_dict()) + if not has_anima_keys: + raise NotAMatchError("state dict does not look like an Anima model") diff --git a/invokeai/backend/model_manager/configs/qwen3_encoder.py b/invokeai/backend/model_manager/configs/qwen3_encoder.py index 2e24fee9185..82cf3b62c8f 100644 --- a/invokeai/backend/model_manager/configs/qwen3_encoder.py +++ b/invokeai/backend/model_manager/configs/qwen3_encoder.py @@ -47,9 +47,10 @@ def _has_ggml_tensors(state_dict: dict[str | int, Any]) -> bool: def _get_qwen3_variant_from_state_dict(state_dict: dict[str | int, Any]) -> Optional[Qwen3VariantType]: - """Determine Qwen3 variant (4B vs 8B) from state dict based on hidden_size. + """Determine Qwen3 variant (0.6B, 4B, or 8B) from state dict based on hidden_size. The hidden_size can be determined from the embed_tokens.weight tensor shape: + - Qwen3 0.6B: hidden_size = 1024 - Qwen3 4B: hidden_size = 2560 - Qwen3 8B: hidden_size = 4096 @@ -57,6 +58,7 @@ def _get_qwen3_variant_from_state_dict(state_dict: dict[str | int, Any]) -> Opti For PyTorch format, the key is 'model.embed_tokens.weight'. """ # Hidden size thresholds + QWEN3_06B_HIDDEN_SIZE = 1024 QWEN3_4B_HIDDEN_SIZE = 2560 QWEN3_8B_HIDDEN_SIZE = 4096 @@ -91,7 +93,9 @@ def _get_qwen3_variant_from_state_dict(state_dict: dict[str | int, Any]) -> Opti return None # Determine variant based on hidden_size - if hidden_size == QWEN3_4B_HIDDEN_SIZE: + if hidden_size == QWEN3_06B_HIDDEN_SIZE: + return Qwen3VariantType.Qwen3_06B + elif hidden_size == QWEN3_4B_HIDDEN_SIZE: return Qwen3VariantType.Qwen3_4B elif hidden_size == QWEN3_8B_HIDDEN_SIZE: return Qwen3VariantType.Qwen3_8B @@ -206,6 +210,7 @@ def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) - @classmethod def _get_variant_from_config(cls, config_path) -> Qwen3VariantType: """Get variant from config.json based on hidden_size.""" + QWEN3_06B_HIDDEN_SIZE = 1024 QWEN3_4B_HIDDEN_SIZE = 2560 QWEN3_8B_HIDDEN_SIZE = 4096 @@ -217,6 +222,8 @@ def _get_variant_from_config(cls, config_path) -> Qwen3VariantType: return Qwen3VariantType.Qwen3_8B elif hidden_size == QWEN3_4B_HIDDEN_SIZE: return Qwen3VariantType.Qwen3_4B + elif hidden_size == QWEN3_06B_HIDDEN_SIZE: + return Qwen3VariantType.Qwen3_06B else: # Default to 4B for unknown sizes return Qwen3VariantType.Qwen3_4B diff --git a/invokeai/backend/model_manager/configs/vae.py b/invokeai/backend/model_manager/configs/vae.py index cc079cb9aae..ce26a94a6e9 100644 --- a/invokeai/backend/model_manager/configs/vae.py +++ b/invokeai/backend/model_manager/configs/vae.py @@ -175,6 +175,43 @@ def _validate_is_flux2_vae(cls, mod: ModelOnDisk) -> None: raise NotAMatchError("state dict does not look like a FLUX.2 VAE") +def _has_anima_vae_keys(state_dict: dict[str | int, Any]) -> bool: + """Check if state dict looks like an Anima QwenImage VAE (AutoencoderKLQwenImage). + + The Anima VAE has a distinctive structure with: + - encoder.downsamples.* (instead of encoder.down_blocks) + - decoder.upsamples.* (instead of decoder.up_blocks) + - decoder.head.* / decoder.middle.* + - Top-level conv1/conv2 weights + """ + required_prefixes = { + "encoder.downsamples.", + "decoder.upsamples.", + "decoder.middle.", + } + return all(any(str(k).startswith(prefix) for k in state_dict) for prefix in required_prefixes) + + +class VAE_Checkpoint_Anima_Config(Checkpoint_Config_Base, Config_Base): + """Model config for Anima QwenImage VAE checkpoint models (AutoencoderKLQwenImage).""" + + type: Literal[ModelType.VAE] = Field(default=ModelType.VAE) + format: Literal[ModelFormat.Checkpoint] = Field(default=ModelFormat.Checkpoint) + base: Literal[BaseModelType.Anima] = Field(default=BaseModelType.Anima) + + @classmethod + def from_model_on_disk(cls, mod: ModelOnDisk, override_fields: dict[str, Any]) -> Self: + raise_if_not_file(mod) + + raise_for_override_fields(cls, override_fields) + + state_dict = mod.load_state_dict() + if not _has_anima_vae_keys(state_dict): + raise NotAMatchError("state dict does not look like an Anima QwenImage VAE") + + return cls(**override_fields) + + class VAE_Diffusers_Config_Base(Diffusers_Config_Base): """Model config for standalone VAE models (diffusers version).""" diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index b972969a687..4609a2e92ab 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -58,7 +58,12 @@ def __init__(self, cache_record: CacheRecord, cache: ModelCache): def __enter__(self) -> AnyModel: self._cache.lock(self._cache_record, None) - return self.model + try: + self.repair_required_tensors_on_device() + return self.model + except Exception: + self._cache.unlock(self._cache_record) + raise def __exit__(self, *args: Any, **kwargs: Any) -> None: self._cache.unlock(self._cache_record) @@ -74,6 +79,7 @@ def model_on_device( """ self._cache.lock(self._cache_record, working_mem_bytes) try: + self.repair_required_tensors_on_device() yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model) finally: self._cache.unlock(self._cache_record) diff --git a/invokeai/backend/model_manager/load/model_loaders/anima.py b/invokeai/backend/model_manager/load/model_loaders/anima.py new file mode 100644 index 00000000000..6549c220a86 --- /dev/null +++ b/invokeai/backend/model_manager/load/model_loaders/anima.py @@ -0,0 +1,140 @@ +# Copyright (c) 2024, Lincoln D. Stein and the InvokeAI Development Team +"""Class for Anima model loading in InvokeAI.""" + +from pathlib import Path +from typing import Optional + +import accelerate + +from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base +from invokeai.backend.model_manager.configs.factory import AnyModelConfig +from invokeai.backend.model_manager.configs.main import Main_Checkpoint_Anima_Config +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + BaseModelType, + ModelFormat, + ModelType, + SubModelType, +) +from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.logging import InvokeAILogger + +logger = InvokeAILogger.get_logger(__name__) + + +@ModelLoaderRegistry.register(base=BaseModelType.Anima, type=ModelType.Main, format=ModelFormat.Checkpoint) +class AnimaCheckpointModel(ModelLoader): + """Class to load Anima transformer models from single-file checkpoints. + + The Anima checkpoint contains both the MiniTrainDIT backbone and the LLM Adapter + under a shared `net.` prefix. The loader strips this prefix and instantiates + the AnimaTransformer model with the correct architecture parameters. + """ + + def _load_model( + self, + config: AnyModelConfig, + submodel_type: Optional[SubModelType] = None, + ) -> AnyModel: + if not isinstance(config, Checkpoint_Config_Base): + raise ValueError("Only CheckpointConfigBase models are currently supported here.") + + match submodel_type: + case SubModelType.Transformer: + return self._load_from_singlefile(config) + + raise ValueError( + f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}" + ) + + def _load_from_singlefile( + self, + config: AnyModelConfig, + ) -> AnyModel: + from safetensors.torch import load_file + + from invokeai.backend.anima.anima_transformer import AnimaTransformer + + if not isinstance(config, Main_Checkpoint_Anima_Config): + raise TypeError( + f"Expected Main_Checkpoint_Anima_Config, got {type(config).__name__}. " + "Model configuration type mismatch." + ) + model_path = Path(config.path) + + # Load the state dict from safetensors + sd = load_file(model_path) + + # Strip the `net.` prefix that all Anima checkpoint keys have + # e.g., "net.blocks.0.self_attn.q_proj.weight" -> "blocks.0.self_attn.q_proj.weight" + prefix_to_strip = None + for prefix in ["net."]: + if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)): + prefix_to_strip = prefix + break + + if prefix_to_strip: + stripped_sd = {} + for key, value in sd.items(): + if isinstance(key, str) and key.startswith(prefix_to_strip): + stripped_sd[key[len(prefix_to_strip) :]] = value + else: + stripped_sd[key] = value + sd = stripped_sd + + # Create an empty AnimaTransformer with Anima's default architecture parameters + with accelerate.init_empty_weights(): + model = AnimaTransformer( + max_img_h=240, + max_img_w=240, + max_frames=1, + in_channels=16, + out_channels=16, + patch_spatial=2, + patch_temporal=1, + concat_padding_mask=True, + model_channels=2048, + num_blocks=28, + num_heads=16, + mlp_ratio=4.0, + crossattn_emb_channels=1024, + pos_emb_cls="rope3d", + use_adaln_lora=True, + adaln_lora_dim=256, + extra_per_block_abs_pos_emb=False, + image_model="anima", + ) + + # Determine safe dtype + target_device = TorchDevice.choose_torch_device() + model_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + + # Handle memory management + new_sd_size = sum(ten.nelement() * model_dtype.itemsize for ten in sd.values()) + self._ram_cache.make_room(new_sd_size) + + # Convert to target dtype (skip non-float tensors like embedding indices) + for k in sd.keys(): + if sd[k].is_floating_point(): + sd[k] = sd[k].to(model_dtype) + + # Filter out rotary embedding inv_freq buffers that are regenerated at runtime + keys_to_remove = [k for k in sd.keys() if k.endswith(".inv_freq")] + for k in keys_to_remove: + del sd[k] + + load_result = model.load_state_dict(sd, assign=True, strict=False) + if load_result.unexpected_keys: + raise RuntimeError( + f"Checkpoint contains {len(load_result.unexpected_keys)} unexpected keys. " + f"This may indicate a corrupted or incompatible checkpoint. " + f"First 5 unexpected keys: {load_result.unexpected_keys[:5]}" + ) + if load_result.missing_keys: + logger.warning( + f"Checkpoint is missing {len(load_result.missing_keys)} keys " + f"(expected for inv_freq buffers). First 5: {load_result.missing_keys[:5]}" + ) + return model diff --git a/invokeai/backend/model_manager/load/model_loaders/lora.py b/invokeai/backend/model_manager/load/model_loaders/lora.py index d39982456af..6cf06d48074 100644 --- a/invokeai/backend/model_manager/load/model_loaders/lora.py +++ b/invokeai/backend/model_manager/load/model_loaders/lora.py @@ -21,6 +21,7 @@ ModelType, SubModelType, ) +from invokeai.backend.patches.lora_conversions.anima_lora_conversion_utils import lora_model_from_anima_state_dict from invokeai.backend.patches.lora_conversions.flux_aitoolkit_lora_conversion_utils import ( is_state_dict_likely_in_flux_aitoolkit_format, lora_model_from_flux_aitoolkit_state_dict, @@ -44,6 +45,10 @@ is_state_dict_likely_in_flux_kohya_format, lora_model_from_flux_kohya_state_dict, ) +from invokeai.backend.patches.lora_conversions.flux_onetrainer_bfl_lora_conversion_utils import ( + is_state_dict_likely_in_flux_onetrainer_bfl_format, + lora_model_from_flux_onetrainer_bfl_state_dict, +) from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_utils import ( is_state_dict_likely_in_flux_onetrainer_format, lora_model_from_flux_onetrainer_state_dict, @@ -52,6 +57,9 @@ is_state_dict_likely_in_flux_xlabs_format, lora_model_from_flux_xlabs_state_dict, ) +from invokeai.backend.patches.lora_conversions.qwen_image_lora_conversion_utils import ( + lora_model_from_qwen_image_state_dict, +) from invokeai.backend.patches.lora_conversions.sd_lora_conversion_utils import lora_model_from_sd_state_dict from invokeai.backend.patches.lora_conversions.sdxl_lora_conversion_utils import convert_sdxl_keys_to_diffusers_format from invokeai.backend.patches.lora_conversions.z_image_lora_conversion_utils import lora_model_from_z_image_state_dict @@ -128,6 +136,8 @@ def _load_model( model = lora_model_from_flux_diffusers_state_dict(state_dict=state_dict, alpha=None) elif is_state_dict_likely_in_flux_kohya_format(state_dict=state_dict): model = lora_model_from_flux_kohya_state_dict(state_dict=state_dict) + elif is_state_dict_likely_in_flux_onetrainer_bfl_format(state_dict=state_dict): + model = lora_model_from_flux_onetrainer_bfl_state_dict(state_dict=state_dict) elif is_state_dict_likely_in_flux_onetrainer_format(state_dict=state_dict): model = lora_model_from_flux_onetrainer_state_dict(state_dict=state_dict) elif is_state_dict_likely_flux_control(state_dict=state_dict): @@ -155,6 +165,11 @@ def _load_model( # Z-Image LoRAs use diffusers PEFT format with transformer and/or Qwen3 encoder layers. # We set alpha=None to use rank as alpha (common default). model = lora_model_from_z_image_state_dict(state_dict=state_dict, alpha=None) + elif self._model_base == BaseModelType.QwenImage: + model = lora_model_from_qwen_image_state_dict(state_dict=state_dict, alpha=None) + elif self._model_base == BaseModelType.Anima: + # Anima LoRAs use Kohya-style or diffusers PEFT format targeting Cosmos DiT blocks. + model = lora_model_from_anima_state_dict(state_dict=state_dict, alpha=None) else: raise ValueError(f"Unsupported LoRA base model: {self._model_base}") diff --git a/invokeai/backend/model_manager/load/model_loaders/qwen_image.py b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py new file mode 100644 index 00000000000..a025e727945 --- /dev/null +++ b/invokeai/backend/model_manager/load/model_loaders/qwen_image.py @@ -0,0 +1,177 @@ +from pathlib import Path +from typing import Optional + +import accelerate +import torch + +from invokeai.backend.model_manager.configs.base import Checkpoint_Config_Base, Diffusers_Config_Base +from invokeai.backend.model_manager.configs.factory import AnyModelConfig +from invokeai.backend.model_manager.configs.main import Main_GGUF_QwenImage_Config +from invokeai.backend.model_manager.load.load_default import ModelLoader +from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader +from invokeai.backend.model_manager.taxonomy import ( + AnyModel, + BaseModelType, + ModelFormat, + ModelType, + QwenImageVariantType, + SubModelType, +) +from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor +from invokeai.backend.quantization.gguf.loaders import gguf_sd_loader +from invokeai.backend.util.devices import TorchDevice + + +@ModelLoaderRegistry.register(base=BaseModelType.QwenImage, type=ModelType.Main, format=ModelFormat.Diffusers) +class QwenImageDiffusersModel(GenericDiffusersLoader): + """Class to load Qwen Image Edit main models.""" + + def _load_model( + self, + config: AnyModelConfig, + submodel_type: Optional[SubModelType] = None, + ) -> AnyModel: + if isinstance(config, Checkpoint_Config_Base): + raise NotImplementedError("CheckpointConfigBase is not implemented for Qwen Image Edit models.") + + if submodel_type is None: + raise Exception("A submodel type must be provided when loading main pipelines.") + + model_path = Path(config.path) + load_class = self.get_hf_load_class(model_path, submodel_type) + repo_variant = config.repo_variant if isinstance(config, Diffusers_Config_Base) else None + variant = repo_variant.value if repo_variant else None + model_path = model_path / submodel_type.value + + # We force bfloat16 for Qwen Image Edit models. + # Use `dtype` (newer) with fallback to `torch_dtype` (older diffusers). + dtype_kwarg = {"dtype": torch.bfloat16} + try: + result: AnyModel = load_class.from_pretrained( + model_path, + **dtype_kwarg, + variant=variant, + local_files_only=True, + ) + except TypeError: + # Older diffusers uses torch_dtype instead of dtype + dtype_kwarg = {"torch_dtype": torch.bfloat16} + result = load_class.from_pretrained( + model_path, + **dtype_kwarg, + variant=variant, + local_files_only=True, + ) + except OSError as e: + if variant and "no file named" in str(e): + result = load_class.from_pretrained(model_path, **dtype_kwarg, local_files_only=True) + else: + raise e + + return result + + +@ModelLoaderRegistry.register(base=BaseModelType.QwenImage, type=ModelType.Main, format=ModelFormat.GGUFQuantized) +class QwenImageGGUFCheckpointModel(ModelLoader): + """Class to load GGUF-quantized Qwen Image Edit transformer models.""" + + def _load_model( + self, + config: AnyModelConfig, + submodel_type: Optional[SubModelType] = None, + ) -> AnyModel: + if not isinstance(config, Checkpoint_Config_Base): + raise ValueError("Only CheckpointConfigBase models are currently supported here.") + + match submodel_type: + case SubModelType.Transformer: + return self._load_from_singlefile(config) + + raise ValueError( + f"Only Transformer submodels are currently supported. Received: {submodel_type.value if submodel_type else 'None'}" + ) + + def _load_from_singlefile(self, config: AnyModelConfig) -> AnyModel: + from diffusers import QwenImageTransformer2DModel + + if not isinstance(config, Main_GGUF_QwenImage_Config): + raise TypeError(f"Expected Main_GGUF_QwenImage_Config, got {type(config).__name__}.") + model_path = Path(config.path) + + target_device = TorchDevice.choose_torch_device() + compute_dtype = TorchDevice.choose_bfloat16_safe_dtype(target_device) + + sd = gguf_sd_loader(model_path, compute_dtype=compute_dtype) + + # Strip ComfyUI-style prefixes if present + prefix_to_strip = None + for prefix in ["model.diffusion_model.", "diffusion_model."]: + if any(k.startswith(prefix) for k in sd.keys() if isinstance(k, str)): + prefix_to_strip = prefix + break + + if prefix_to_strip: + stripped_sd = {} + for key, value in sd.items(): + if isinstance(key, str) and key.startswith(prefix_to_strip): + stripped_sd[key[len(prefix_to_strip) :]] = value + else: + stripped_sd[key] = value + sd = stripped_sd + + # Auto-detect architecture from state dict + num_layers = 0 + for key in sd.keys(): + if isinstance(key, str) and key.startswith("transformer_blocks."): + parts = key.split(".") + if len(parts) >= 2: + try: + layer_idx = int(parts[1]) + num_layers = max(num_layers, layer_idx + 1) + except ValueError: + pass + + # Detect dimensions from weights + num_attention_heads = 24 # default + attention_head_dim = 128 # default + + if "img_in.weight" in sd: + w = sd["img_in.weight"] + shape = w.tensor_shape if isinstance(w, GGMLTensor) else w.shape + hidden_dim = shape[0] + in_channels = shape[1] + num_attention_heads = hidden_dim // attention_head_dim + + joint_attention_dim = 3584 # default + if "txt_in.weight" in sd: + w = sd["txt_in.weight"] + shape = w.tensor_shape if isinstance(w, GGMLTensor) else w.shape + joint_attention_dim = shape[1] + + model_config: dict = { + "patch_size": 2, + "in_channels": in_channels if "img_in.weight" in sd else 64, + "out_channels": 16, + "num_layers": num_layers if num_layers > 0 else 60, + "attention_head_dim": attention_head_dim, + "num_attention_heads": num_attention_heads, + "joint_attention_dim": joint_attention_dim, + "guidance_embeds": False, + "axes_dims_rope": (16, 56, 56), + } + + # zero_cond_t is only used by edit-variant models. It enables dual modulation + # for noisy vs reference patches. Setting it on txt2img models produces garbage. + # Also requires diffusers 0.37+ (the parameter doesn't exist in older versions). + import inspect + + is_edit = getattr(config, "variant", None) == QwenImageVariantType.Edit + if is_edit and "zero_cond_t" in inspect.signature(QwenImageTransformer2DModel.__init__).parameters: + model_config["zero_cond_t"] = True + + with accelerate.init_empty_weights(): + model = QwenImageTransformer2DModel(**model_config) + + model.load_state_dict(sd, strict=False, assign=True) + return model diff --git a/invokeai/backend/model_manager/load/model_loaders/vae.py b/invokeai/backend/model_manager/load/model_loaders/vae.py index e91903ccdad..db26e8c6547 100644 --- a/invokeai/backend/model_manager/load/model_loaders/vae.py +++ b/invokeai/backend/model_manager/load/model_loaders/vae.py @@ -6,7 +6,7 @@ from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from invokeai.backend.model_manager.configs.factory import AnyModelConfig -from invokeai.backend.model_manager.configs.vae import VAE_Checkpoint_Config_Base +from invokeai.backend.model_manager.configs.vae import VAE_Checkpoint_Anima_Config, VAE_Checkpoint_Config_Base from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader from invokeai.backend.model_manager.taxonomy import ( @@ -28,7 +28,14 @@ def _load_model( config: AnyModelConfig, submodel_type: Optional[SubModelType] = None, ) -> AnyModel: - if isinstance(config, VAE_Checkpoint_Config_Base): + if isinstance(config, VAE_Checkpoint_Anima_Config): + from diffusers.models.autoencoders import AutoencoderKLWan + + return AutoencoderKLWan.from_single_file( + config.path, + torch_dtype=self._torch_dtype, + ) + elif isinstance(config, VAE_Checkpoint_Config_Base): return AutoencoderKL.from_single_file( config.path, torch_dtype=self._torch_dtype, diff --git a/invokeai/backend/model_manager/metadata/fetch/huggingface.py b/invokeai/backend/model_manager/metadata/fetch/huggingface.py index 1b2b6c36742..30fe418fe14 100644 --- a/invokeai/backend/model_manager/metadata/fetch/huggingface.py +++ b/invokeai/backend/model_manager/metadata/fetch/huggingface.py @@ -19,8 +19,7 @@ from typing import Optional import requests -from huggingface_hub import HfApi, configure_http_backend, hf_hub_url -from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundError +from huggingface_hub import hf_hub_url from pydantic.networks import AnyHttpUrl from requests.sessions import Session @@ -47,7 +46,6 @@ def __init__(self, session: Optional[Session] = None): this module without an internet connection. """ self._requests = session or requests.Session() - configure_http_backend(backend_factory=lambda: self._requests) @classmethod def from_json(cls, json: str) -> HuggingFaceMetadata: @@ -55,6 +53,22 @@ def from_json(cls, json: str) -> HuggingFaceMetadata: metadata = HuggingFaceMetadata.model_validate_json(json) return metadata + def _fetch_model_info(self, repo_id: str, variant: Optional[ModelRepoVariant] = None) -> dict: + """Fetch model info from HuggingFace API using self._requests session. + + This allows the session to be mocked in tests via requests_testadapter. + """ + url = f"https://huggingface.co/api/models/{repo_id}" + params: dict[str, str] = {"blobs": "True"} + if variant is not None: + params["revision"] = str(variant) + + response = self._requests.get(url, params=params) + if response.status_code == 404: + raise UnknownMetadataException(f"'{repo_id}' not found.") + response.raise_for_status() + return response.json() + def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyModelRepoMetadata: """Return a HuggingFaceMetadata object given the model's repo_id.""" # Little loop which tries fetching a revision corresponding to the selected variant. @@ -67,10 +81,10 @@ def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyMod repo_id = id.split("::")[0] or id while not model_info: try: - model_info = HfApi().model_info(repo_id=repo_id, files_metadata=True, revision=variant) - except RepositoryNotFoundError as excp: - raise UnknownMetadataException(f"'{repo_id}' not found. See trace for details.") from excp - except RevisionNotFoundError: + model_info = self._fetch_model_info(repo_id, variant) + except UnknownMetadataException: + raise + except requests.HTTPError: if variant is None: raise else: @@ -80,15 +94,18 @@ def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyMod _, name = repo_id.split("/") - for s in model_info.siblings or []: - assert s.rfilename is not None - assert s.size is not None + for s in model_info.get("siblings") or []: + rfilename = s.get("rfilename") + size = s.get("size") + assert rfilename is not None + assert size is not None + lfs = s.get("lfs") files.append( RemoteModelFile( - url=hf_hub_url(repo_id, s.rfilename, revision=variant or "main"), - path=Path(name, s.rfilename), - size=s.size, - sha256=s.lfs.get("sha256") if s.lfs else None, + url=hf_hub_url(repo_id, rfilename, revision=variant or "main"), + path=Path(name, rfilename), + size=size, + sha256=lfs.get("sha256") if lfs else None, ) ) @@ -115,10 +132,10 @@ def from_id(self, id: str, variant: Optional[ModelRepoVariant] = None) -> AnyMod ) return HuggingFaceMetadata( - id=model_info.id, + id=model_info["id"], name=name, files=files, - api_response=json.dumps(model_info.__dict__, default=str), + api_response=json.dumps(model_info, default=str), is_diffusers=is_diffusers, ckpt_urls=ckpt_urls, ) diff --git a/invokeai/backend/model_manager/metadata/metadata_base.py b/invokeai/backend/model_manager/metadata/metadata_base.py index e16ad4cbc47..b048144e547 100644 --- a/invokeai/backend/model_manager/metadata/metadata_base.py +++ b/invokeai/backend/model_manager/metadata/metadata_base.py @@ -17,7 +17,7 @@ from pathlib import Path from typing import List, Literal, Optional, Union -from huggingface_hub import configure_http_backend, hf_hub_url +from huggingface_hub import hf_hub_url from pydantic import BaseModel, Field, TypeAdapter from pydantic.networks import AnyHttpUrl from requests.sessions import Session @@ -111,7 +111,6 @@ def download_urls( full-precision model is returned. """ session = session or Session() - configure_http_backend(backend_factory=lambda: session) # used in testing paths = filter_files([x.path for x in self.files], variant, subfolder, subfolders) # all files in the model diff --git a/invokeai/backend/model_manager/starter_models.py b/invokeai/backend/model_manager/starter_models.py index 9f86f83dc59..66637a59707 100644 --- a/invokeai/backend/model_manager/starter_models.py +++ b/invokeai/backend/model_manager/starter_models.py @@ -2,7 +2,13 @@ from pydantic import BaseModel -from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType +from invokeai.backend.model_manager.taxonomy import ( + AnyVariant, + BaseModelType, + ModelFormat, + ModelType, + QwenImageVariantType, +) class StarterModelWithoutDependencies(BaseModel): @@ -12,6 +18,7 @@ class StarterModelWithoutDependencies(BaseModel): base: BaseModelType type: ModelType format: Optional[ModelFormat] = None + variant: Optional[AnyVariant] = None is_installed: bool = False # allows us to track what models a user has installed across name changes within starter models # if you update a starter model name, please add the old one to this list for that starter model @@ -71,7 +78,7 @@ class StarterModelBundle(BaseModel): name="t5_base_encoder", base=BaseModelType.Any, source="InvokeAI/t5-v1_1-xxl::bfloat16", - description="T5-XXL text encoder (used in FLUX pipelines). ~8GB", + description="T5-XXL text encoder (used in FLUX pipelines). ~9.5GB", type=ModelType.T5Encoder, ) @@ -156,7 +163,7 @@ class StarterModelBundle(BaseModel): name="FLUX.1 Kontext dev (quantized)", base=BaseModelType.Flux, source="https://huggingface.co/unsloth/FLUX.1-Kontext-dev-GGUF/resolve/main/flux1-kontext-dev-Q4_K_M.gguf", - description="FLUX.1 Kontext dev quantized (q4_k_m). Total size with dependencies: ~14GB", + description="FLUX.1 Kontext dev quantized (q4_k_m). Total size with dependencies: ~12GB", type=ModelType.Main, dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder], ) @@ -164,7 +171,7 @@ class StarterModelBundle(BaseModel): name="FLUX.1 Krea dev", base=BaseModelType.Flux, source="https://huggingface.co/InvokeAI/FLUX.1-Krea-dev/resolve/main/flux1-krea-dev.safetensors", - description="FLUX.1 Krea dev. Total size with dependencies: ~33GB", + description="FLUX.1 Krea dev. Total size with dependencies: ~29GB", type=ModelType.Main, dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder], ) @@ -172,7 +179,7 @@ class StarterModelBundle(BaseModel): name="FLUX.1 Krea dev (quantized)", base=BaseModelType.Flux, source="https://huggingface.co/InvokeAI/FLUX.1-Krea-dev-GGUF/resolve/main/flux1-krea-dev-Q4_K_M.gguf", - description="FLUX.1 Krea dev quantized (q4_k_m). Total size with dependencies: ~14GB", + description="FLUX.1 Krea dev quantized (q4_k_m). Total size with dependencies: ~12GB", type=ModelType.Main, dependencies=[t5_8b_quantized_encoder, flux_vae, clip_l_encoder], ) @@ -180,7 +187,7 @@ class StarterModelBundle(BaseModel): name="SD3.5 Medium", base=BaseModelType.StableDiffusion3, source="stabilityai/stable-diffusion-3.5-medium", - description="Medium SD3.5 Model: ~15GB", + description="Medium SD3.5 Model: ~16GB", type=ModelType.Main, dependencies=[], ) @@ -188,7 +195,7 @@ class StarterModelBundle(BaseModel): name="SD3.5 Large", base=BaseModelType.StableDiffusion3, source="stabilityai/stable-diffusion-3.5-large", - description="Large SD3.5 Model: ~19G", + description="Large SD3.5 Model: ~28GB", type=ModelType.Main, dependencies=[], ) @@ -644,8 +651,140 @@ class StarterModelBundle(BaseModel): name="CogView4", base=BaseModelType.CogView4, source="THUDM/CogView4-6B", - description="The base CogView4 model (~29GB).", + description="The base CogView4 model (~31GB).", + type=ModelType.Main, +) +# endregion + +# region Qwen Image Edit +qwen_image_edit = StarterModel( + name="Qwen Image Edit 2511", + base=BaseModelType.QwenImage, + source="Qwen/Qwen-Image-Edit-2511", + description="Qwen Image Edit 2511 full diffusers model. Supports text-guided image editing with multiple reference images. (~40GB)", + type=ModelType.Main, + variant=QwenImageVariantType.Edit, +) + +qwen_image_edit_gguf_q4_k_m = StarterModel( + name="Qwen Image Edit 2511 (Q4_K_M)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-Edit-2511-GGUF/resolve/main/qwen-image-edit-2511-Q4_K_M.gguf", + description="Qwen Image Edit 2511 - Q4_K_M quantized transformer. Good quality/size balance. (~13GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, + variant=QwenImageVariantType.Edit, +) + +qwen_image_edit_gguf_q2_k = StarterModel( + name="Qwen Image Edit 2511 (Q2_K)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-Edit-2511-GGUF/resolve/main/qwen-image-edit-2511-Q2_K.gguf", + description="Qwen Image Edit 2511 - Q2_K heavily quantized transformer. Smallest size, lower quality. (~7.5GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, + variant=QwenImageVariantType.Edit, +) + +qwen_image_edit_gguf_q6_k = StarterModel( + name="Qwen Image Edit 2511 (Q6_K)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-Edit-2511-GGUF/resolve/main/qwen-image-edit-2511-Q6_K.gguf", + description="Qwen Image Edit 2511 - Q6_K quantized transformer. Near-lossless quality. (~17GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, + variant=QwenImageVariantType.Edit, +) + +qwen_image_edit_gguf_q8_0 = StarterModel( + name="Qwen Image Edit 2511 (Q8_0)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-Edit-2511-GGUF/resolve/main/qwen-image-edit-2511-Q8_0.gguf", + description="Qwen Image Edit 2511 - Q8_0 quantized transformer. Highest quality quantization. (~22GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, + variant=QwenImageVariantType.Edit, +) + +qwen_image_edit_lightning_4step = StarterModel( + name="Qwen Image Edit Lightning (4-step, bf16)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/lightx2v/Qwen-Image-Edit-2511-Lightning/resolve/main/Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors", + description="Lightning distillation LoRA for Qwen Image Edit — enables generation in just 4 steps. " + "Settings: Steps=4, CFG=1, Shift Override=3.", + type=ModelType.LoRA, +) + +qwen_image_edit_lightning_8step = StarterModel( + name="Qwen Image Edit Lightning (8-step, bf16)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/lightx2v/Qwen-Image-Edit-2511-Lightning/resolve/main/Qwen-Image-Edit-2511-Lightning-8steps-V1.0-bf16.safetensors", + description="Lightning distillation LoRA for Qwen Image Edit — enables generation in 8 steps with better quality. " + "Settings: Steps=8, CFG=1, Shift Override=3.", + type=ModelType.LoRA, +) + +# Qwen Image (txt2img) +qwen_image = StarterModel( + name="Qwen Image 2512", + base=BaseModelType.QwenImage, + source="Qwen/Qwen-Image-2512", + description="Qwen Image 2512 full diffusers model. High-quality text-to-image generation. (~40GB)", + type=ModelType.Main, +) + +qwen_image_gguf_q4_k_m = StarterModel( + name="Qwen Image 2512 (Q4_K_M)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-2512-GGUF/resolve/main/qwen-image-2512-Q4_K_M.gguf", + description="Qwen Image 2512 - Q4_K_M quantized transformer. Good quality/size balance. (~13GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, +) + +qwen_image_gguf_q2_k = StarterModel( + name="Qwen Image 2512 (Q2_K)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-2512-GGUF/resolve/main/qwen-image-2512-Q2_K.gguf", + description="Qwen Image 2512 - Q2_K heavily quantized transformer. Smallest size, lower quality. (~7.5GB)", type=ModelType.Main, + format=ModelFormat.GGUFQuantized, +) + +qwen_image_gguf_q6_k = StarterModel( + name="Qwen Image 2512 (Q6_K)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-2512-GGUF/resolve/main/qwen-image-2512-Q6_K.gguf", + description="Qwen Image 2512 - Q6_K quantized transformer. Near-lossless quality. (~17GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, +) + +qwen_image_gguf_q8_0 = StarterModel( + name="Qwen Image 2512 (Q8_0)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/unsloth/Qwen-Image-2512-GGUF/resolve/main/qwen-image-2512-Q8_0.gguf", + description="Qwen Image 2512 - Q8_0 quantized transformer. Highest quality quantization. (~22GB)", + type=ModelType.Main, + format=ModelFormat.GGUFQuantized, +) + +qwen_image_lightning_4step = StarterModel( + name="Qwen Image Lightning (4-step, V2.0, bf16)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors", + description="Lightning distillation LoRA for Qwen Image — enables generation in just 4 steps. " + "Settings: Steps=4, CFG=1, Shift Override=3.", + type=ModelType.LoRA, +) + +qwen_image_lightning_8step = StarterModel( + name="Qwen Image Lightning (8-step, V2.0, bf16)", + base=BaseModelType.QwenImage, + source="https://huggingface.co/lightx2v/Qwen-Image-Lightning/resolve/main/Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors", + description="Lightning distillation LoRA for Qwen Image — enables generation in 8 steps with better quality. " + "Settings: Steps=8, CFG=1, Shift Override=3.", + type=ModelType.LoRA, ) # endregion @@ -695,7 +834,7 @@ class StarterModelBundle(BaseModel): name="FLUX.2 VAE", base=BaseModelType.Flux2, source="black-forest-labs/FLUX.2-klein-4B::vae", - description="FLUX.2 VAE (16-channel, same architecture as FLUX.1 VAE). ~335MB", + description="FLUX.2 VAE (16-channel, same architecture as FLUX.1 VAE). ~168MB", type=ModelType.VAE, ) @@ -719,7 +858,7 @@ class StarterModelBundle(BaseModel): name="FLUX.2 Klein 4B (Diffusers)", base=BaseModelType.Flux2, source="black-forest-labs/FLUX.2-klein-4B", - description="FLUX.2 Klein 4B in Diffusers format - includes transformer, VAE and Qwen3 encoder. ~10GB", + description="FLUX.2 Klein 4B in Diffusers format - includes transformer, VAE and Qwen3 encoder. ~16GB", type=ModelType.Main, ) @@ -745,7 +884,7 @@ class StarterModelBundle(BaseModel): name="FLUX.2 Klein 9B (Diffusers)", base=BaseModelType.Flux2, source="black-forest-labs/FLUX.2-klein-9B", - description="FLUX.2 Klein 9B in Diffusers format - includes transformer, VAE and Qwen3 encoder. ~20GB", + description="FLUX.2 Klein 9B in Diffusers format - includes transformer, VAE and Qwen3 encoder. ~35GB", type=ModelType.Main, ) @@ -821,7 +960,7 @@ class StarterModelBundle(BaseModel): name="Z-Image Turbo", base=BaseModelType.ZImage, source="Tongyi-MAI/Z-Image-Turbo", - description="Z-Image Turbo - fast 6B parameter text-to-image model with 8 inference steps. Supports bilingual prompts (English & Chinese). ~30.6GB", + description="Z-Image Turbo - fast 6B parameter text-to-image model with 8 inference steps. Supports bilingual prompts (English & Chinese). ~33GB", type=ModelType.Main, ) @@ -862,6 +1001,36 @@ class StarterModelBundle(BaseModel): ) # endregion +# region Anima +anima_qwen3_encoder = StarterModel( + name="Anima Qwen3 0.6B Text Encoder", + base=BaseModelType.Any, + source="https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/text_encoders/qwen_3_06b_base.safetensors", + description="Qwen3 0.6B text encoder for Anima. ~1.2GB", + type=ModelType.Qwen3Encoder, + format=ModelFormat.Checkpoint, +) + +anima_vae = StarterModel( + name="Anima QwenImage VAE", + base=BaseModelType.Anima, + source="https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/vae/qwen_image_vae.safetensors", + description="QwenImage VAE for Anima (fine-tuned Wan 2.1 VAE, 16 latent channels). ~200MB", + type=ModelType.VAE, + format=ModelFormat.Checkpoint, +) + +anima_preview3 = StarterModel( + name="Anima Preview 3", + base=BaseModelType.Anima, + source="https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/diffusion_models/anima-preview3-base.safetensors", + description="Anima Preview 3 - 2B parameter anime-focused text-to-image model built on Cosmos Predict2 DiT. ~4.5GB", + type=ModelType.Main, + format=ModelFormat.Checkpoint, + dependencies=[anima_qwen3_encoder, anima_vae, t5_base_encoder], +) +# endregion + # List of starter models, displayed on the frontend. # The order/sort of this list is not changed by the frontend - set it how you want it here. STARTER_MODELS: list[StarterModel] = [ @@ -948,6 +1117,20 @@ class StarterModelBundle(BaseModel): flux2_klein_qwen3_4b_encoder, flux2_klein_qwen3_8b_encoder, cogview4, + qwen_image_edit, + qwen_image_edit_gguf_q2_k, + qwen_image_edit_gguf_q4_k_m, + qwen_image_edit_gguf_q6_k, + qwen_image_edit_gguf_q8_0, + qwen_image_edit_lightning_4step, + qwen_image_edit_lightning_8step, + qwen_image, + qwen_image_gguf_q2_k, + qwen_image_gguf_q4_k_m, + qwen_image_gguf_q6_k, + qwen_image_gguf_q8_0, + qwen_image_lightning_4step, + qwen_image_lightning_8step, flux_krea, flux_krea_quantized, z_image_turbo, @@ -957,6 +1140,9 @@ class StarterModelBundle(BaseModel): z_image_qwen3_encoder_quantized, z_image_controlnet_union, z_image_controlnet_tile, + anima_preview3, + anima_qwen3_encoder, + anima_vae, ] sd1_bundle: list[StarterModel] = [ @@ -1025,12 +1211,34 @@ class StarterModelBundle(BaseModel): flux2_klein_qwen3_4b_encoder, ] +qwen_image_bundle: list[StarterModel] = [ + qwen_image_edit, + qwen_image_edit_gguf_q4_k_m, + qwen_image_edit_gguf_q8_0, + qwen_image_edit_lightning_4step, + qwen_image_edit_lightning_8step, + qwen_image, + qwen_image_gguf_q4_k_m, + qwen_image_gguf_q8_0, + qwen_image_lightning_4step, + qwen_image_lightning_8step, +] + +anima_bundle: list[StarterModel] = [ + anima_preview3, + anima_qwen3_encoder, + anima_vae, + t5_base_encoder, +] + STARTER_BUNDLES: dict[str, StarterModelBundle] = { BaseModelType.StableDiffusion1: StarterModelBundle(name="Stable Diffusion 1.5", models=sd1_bundle), BaseModelType.StableDiffusionXL: StarterModelBundle(name="SDXL", models=sdxl_bundle), BaseModelType.Flux: StarterModelBundle(name="FLUX.1 dev", models=flux_bundle), BaseModelType.Flux2: StarterModelBundle(name="FLUX.2 Klein", models=flux2_klein_bundle), BaseModelType.ZImage: StarterModelBundle(name="Z-Image Turbo", models=zimage_bundle), + BaseModelType.QwenImage: StarterModelBundle(name="Qwen Image", models=qwen_image_bundle), + BaseModelType.Anima: StarterModelBundle(name="Anima", models=anima_bundle), } assert len(STARTER_MODELS) == len({m.source for m in STARTER_MODELS}), "Duplicate starter models" diff --git a/invokeai/backend/model_manager/taxonomy.py b/invokeai/backend/model_manager/taxonomy.py index c002418a6bd..861984f304d 100644 --- a/invokeai/backend/model_manager/taxonomy.py +++ b/invokeai/backend/model_manager/taxonomy.py @@ -52,6 +52,10 @@ class BaseModelType(str, Enum): """Indicates the model is associated with CogView 4 model architecture.""" ZImage = "z-image" """Indicates the model is associated with Z-Image model architecture, including Z-Image-Turbo.""" + QwenImage = "qwen-image" + """Indicates the model is associated with Qwen Image Edit 2511 model architecture.""" + Anima = "anima" + """Indicates the model is associated with Anima model architecture (Cosmos Predict2 DiT + LLM Adapter).""" Unknown = "unknown" """Indicates the model's base architecture is unknown.""" @@ -143,6 +147,16 @@ class ZImageVariantType(str, Enum): """Z-Image Base - undistilled foundation model with full CFG and negative prompt support.""" +class QwenImageVariantType(str, Enum): + """Qwen Image model variants.""" + + Generate = "generate" + """Qwen Image - text-to-image generation model.""" + + Edit = "edit" + """Qwen Image Edit - image editing model with reference image support.""" + + class Qwen3VariantType(str, Enum): """Qwen3 text encoder variants based on model size.""" @@ -152,6 +166,9 @@ class Qwen3VariantType(str, Enum): Qwen3_8B = "qwen3_8b" """Qwen3 8B text encoder (hidden_size=4096). Used by FLUX.2 Klein 9B.""" + Qwen3_06B = "qwen3_06b" + """Qwen3 0.6B text encoder (hidden_size=1024). Used by Anima.""" + class ModelFormat(str, Enum): """Storage format of model.""" @@ -210,11 +227,32 @@ class FluxLoRAFormat(str, Enum): AIToolkit = "flux.aitoolkit" XLabs = "flux.xlabs" BflPeft = "flux.bfl_peft" + OneTrainerBfl = "flux.onetrainer_bfl" AnyVariant: TypeAlias = Union[ - ModelVariantType, ClipVariantType, FluxVariantType, Flux2VariantType, ZImageVariantType, Qwen3VariantType + ModelVariantType, + ClipVariantType, + FluxVariantType, + Flux2VariantType, + ZImageVariantType, + QwenImageVariantType, + Qwen3VariantType, ] variant_type_adapter = TypeAdapter[ - ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType -](ModelVariantType | ClipVariantType | FluxVariantType | Flux2VariantType | ZImageVariantType | Qwen3VariantType) + ModelVariantType + | ClipVariantType + | FluxVariantType + | Flux2VariantType + | ZImageVariantType + | QwenImageVariantType + | Qwen3VariantType +]( + ModelVariantType + | ClipVariantType + | FluxVariantType + | Flux2VariantType + | ZImageVariantType + | QwenImageVariantType + | Qwen3VariantType +) diff --git a/invokeai/backend/patches/lora_conversions/anima_lora_constants.py b/invokeai/backend/patches/lora_conversions/anima_lora_constants.py new file mode 100644 index 00000000000..380e31998a7 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/anima_lora_constants.py @@ -0,0 +1,45 @@ +# Anima LoRA prefix constants +# These prefixes are used for key mapping when applying LoRA patches to Anima models + +import re + +# Prefix for Anima transformer (Cosmos DiT architecture) LoRA layers +ANIMA_LORA_TRANSFORMER_PREFIX = "lora_transformer-" + +# Prefix for Qwen3 text encoder LoRA layers +ANIMA_LORA_QWEN3_PREFIX = "lora_qwen3-" + +# --------------------------------------------------------------------------- +# Cosmos DiT detection helpers +# +# Shared between ``anima_lora_conversion_utils.is_state_dict_likely_anima_lora`` +# and the config probing code in ``configs/lora.py``. Kept here (rather than +# in ``anima_lora_conversion_utils``) to avoid circular imports. +# --------------------------------------------------------------------------- + +# Cosmos DiT subcomponent names unique to the Anima / Cosmos Predict2 architecture. +_COSMOS_DIT_SUBCOMPONENTS_RE = r"(cross_attn|self_attn|mlp|adaln_modulation)" + +# Kohya format: lora_unet_[llm_adapter_]blocks_N_ +_KOHYA_ANIMA_RE = re.compile(r"lora_unet_(llm_adapter_)?blocks_\d+_" + _COSMOS_DIT_SUBCOMPONENTS_RE) + +# PEFT format: .blocks.N. +_PEFT_ANIMA_RE = re.compile( + r"(diffusion_model|transformer|base_model\.model\.transformer)\.blocks\.\d+\." + _COSMOS_DIT_SUBCOMPONENTS_RE +) + + +def has_cosmos_dit_kohya_keys(str_keys: list[str]) -> bool: + """Check for Kohya-style keys targeting Cosmos DiT blocks with specific subcomponents. + + Requires both the ``lora_unet_[llm_adapter_]blocks_N_`` prefix **and** a + Cosmos DiT subcomponent name (cross_attn, self_attn, mlp, adaln_modulation) + to avoid false-positives on other architectures that might also use bare + ``blocks`` in their key paths. + """ + return any(_KOHYA_ANIMA_RE.search(k) is not None for k in str_keys) + + +def has_cosmos_dit_peft_keys(str_keys: list[str]) -> bool: + """Check for diffusers PEFT keys targeting Cosmos DiT blocks with specific subcomponents.""" + return any(_PEFT_ANIMA_RE.search(k) is not None for k in str_keys) diff --git a/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py new file mode 100644 index 00000000000..b55a96dca75 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/anima_lora_conversion_utils.py @@ -0,0 +1,300 @@ +"""Anima LoRA conversion utilities. + +Anima uses a Cosmos Predict2 DiT transformer architecture. +LoRAs for Anima typically follow the Kohya-style format with underscore-separated keys +(e.g., lora_unet_blocks_0_cross_attn_k_proj) that map to model parameter paths +(e.g., blocks.0.cross_attn.k_proj). + +Some Anima LoRAs also target the Qwen3 text encoder with lora_te_ prefix keys +(e.g., lora_te_layers_0_self_attn_q_proj -> layers.0.self_attn.q_proj). +""" + +import re +from typing import Dict + +import torch + +from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch +from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict +from invokeai.backend.patches.lora_conversions.anima_lora_constants import ( + ANIMA_LORA_QWEN3_PREFIX, + ANIMA_LORA_TRANSFORMER_PREFIX, + has_cosmos_dit_kohya_keys, + has_cosmos_dit_peft_keys, +) +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw +from invokeai.backend.util.logging import InvokeAILogger + +logger = InvokeAILogger.get_logger(__name__) + + +def is_state_dict_likely_anima_lora(state_dict: dict[str | int, torch.Tensor]) -> bool: + """Checks if the provided state dict is likely an Anima LoRA. + + Anima LoRAs use Kohya-style naming with lora_unet_ prefix and underscore-separated + model key paths targeting Cosmos DiT blocks. Detection requires Cosmos DiT-specific + subcomponent names (cross_attn, self_attn, mlp, adaln_modulation) to avoid + false-positives on other architectures that also use ``blocks`` in their paths. + """ + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + + if has_cosmos_dit_kohya_keys(str_keys): + return True + + return has_cosmos_dit_peft_keys(str_keys) + + +# Mapping from Kohya underscore-style substrings to model parameter names. +# Order matters: longer/more specific patterns should come first to avoid partial matches. +_KOHYA_UNET_KEY_REPLACEMENTS = [ + ("adaln_modulation_cross_attn_", "adaln_modulation_cross_attn."), + ("adaln_modulation_self_attn_", "adaln_modulation_self_attn."), + ("adaln_modulation_mlp_", "adaln_modulation_mlp."), + ("cross_attn_k_proj", "cross_attn.k_proj"), + ("cross_attn_q_proj", "cross_attn.q_proj"), + ("cross_attn_v_proj", "cross_attn.v_proj"), + ("cross_attn_output_proj", "cross_attn.output_proj"), + ("cross_attn_o_proj", "cross_attn.o_proj"), + ("self_attn_k_proj", "self_attn.k_proj"), + ("self_attn_q_proj", "self_attn.q_proj"), + ("self_attn_v_proj", "self_attn.v_proj"), + ("self_attn_output_proj", "self_attn.output_proj"), + ("self_attn_o_proj", "self_attn.o_proj"), + ("mlp_layer1", "mlp.layer1"), + ("mlp_layer2", "mlp.layer2"), +] + +# Mapping for Qwen3 text encoder Kohya keys. +_KOHYA_TE_KEY_REPLACEMENTS = [ + ("self_attn_k_proj", "self_attn.k_proj"), + ("self_attn_q_proj", "self_attn.q_proj"), + ("self_attn_v_proj", "self_attn.v_proj"), + ("self_attn_o_proj", "self_attn.o_proj"), + ("mlp_down_proj", "mlp.down_proj"), + ("mlp_gate_proj", "mlp.gate_proj"), + ("mlp_up_proj", "mlp.up_proj"), +] + + +def _convert_kohya_unet_key(kohya_layer_name: str) -> str: + """Convert a Kohya-style LoRA layer name to a model parameter path. + + Example: lora_unet_blocks_0_cross_attn_k_proj -> blocks.0.cross_attn.k_proj + Example: lora_unet_llm_adapter_blocks_0_cross_attn_k_proj -> llm_adapter.blocks.0.cross_attn.k_proj + """ + key = kohya_layer_name + if key.startswith("lora_unet_"): + key = key[len("lora_unet_") :] + + # Handle llm_adapter prefix: strip it, run the standard block conversion, then re-add with dot + llm_adapter_prefix = "" + if key.startswith("llm_adapter_"): + key = key[len("llm_adapter_") :] + llm_adapter_prefix = "llm_adapter." + + # Convert blocks_N_ to blocks.N. + key = re.sub(r"^blocks_(\d+)_", r"blocks.\1.", key) + + # Apply known replacements for subcomponent names + for old, new in _KOHYA_UNET_KEY_REPLACEMENTS: + if old in key: + key = key.replace(old, new, 1) + break + + return llm_adapter_prefix + key + + +def _convert_kohya_te_key(kohya_layer_name: str) -> str: + """Convert a Kohya-style text encoder LoRA layer name to a model parameter path. + + The Qwen3 text encoder is loaded as Qwen3ForCausalLM which wraps the base model + under a `model.` prefix, so the final path must include it. + + Example: lora_te_layers_0_self_attn_q_proj -> model.layers.0.self_attn.q_proj + """ + key = kohya_layer_name + if key.startswith("lora_te_"): + key = key[len("lora_te_") :] + + # Convert layers_N_ to layers.N. + key = re.sub(r"^layers_(\d+)_", r"layers.\1.", key) + + # Apply known replacements + for old, new in _KOHYA_TE_KEY_REPLACEMENTS: + if old in key: + key = key.replace(old, new, 1) + break + + # Qwen3ForCausalLM wraps the base Qwen3Model under `model.` + key = f"model.{key}" + + return key + + +def _make_layer_patch(layer_dict: dict[str, torch.Tensor]) -> BaseLayerPatch: + """Create a layer patch from a layer dict, handling DoRA+LoKR edge case. + + Some Anima LoRAs combine DoRA (dora_scale) with LoKR (lokr_w1/lokr_w2) weights. + The shared any_lora_layer_from_state_dict checks dora_scale first and expects + lora_up/lora_down keys, which don't exist in LoKR layers. We strip dora_scale + from LoKR layers so they fall through to the LoKR handler instead. + """ + has_lokr = "lokr_w1" in layer_dict or "lokr_w1_a" in layer_dict + has_dora = "dora_scale" in layer_dict + if has_lokr and has_dora: + layer_dict = {k: v for k, v in layer_dict.items() if k != "dora_scale"} + logger.warning("Stripped dora_scale from LoKR layer (DoRA+LoKR combination not supported, using LoKR only)") + return any_lora_layer_from_state_dict(layer_dict) + + +# Known suffixes for Kohya format +_KOHYA_KNOWN_SUFFIXES = [ + ".lora_A.weight", + ".lora_B.weight", + ".lora_down.weight", + ".lora_up.weight", + ".dora_scale", + ".alpha", +] + +# Additional suffixes for PEFT/LoKR format +_PEFT_EXTRA_SUFFIXES = [ + ".lokr_w1", + ".lokr_w2", + ".lokr_w1_a", + ".lokr_w1_b", + ".lokr_w2_a", + ".lokr_w2_b", +] + + +def _group_keys_by_layer( + state_dict: Dict[str, torch.Tensor], + extra_suffixes: list[str] | None = None, +) -> dict[str, dict[str, torch.Tensor]]: + """Group state dict keys by layer name based on known suffixes. + + Args: + state_dict: The LoRA state dict to group. + extra_suffixes: Additional suffixes to recognize beyond the base Kohya set. + + Returns: + Dict mapping layer names to their component tensors. + """ + layer_dict: dict[str, dict[str, torch.Tensor]] = {} + + known_suffixes = list(_KOHYA_KNOWN_SUFFIXES) + if extra_suffixes: + known_suffixes.extend(extra_suffixes) + + for key in state_dict: + if not isinstance(key, str): + continue + + layer_name = None + key_name = None + for suffix in known_suffixes: + if key.endswith(suffix): + layer_name = key[: -len(suffix)] + key_name = suffix[1:] # Remove leading dot + break + + if layer_name is None: + parts = key.rsplit(".", maxsplit=2) + layer_name = parts[0] + key_name = ".".join(parts[1:]) + + if layer_name not in layer_dict: + layer_dict[layer_name] = {} + layer_dict[layer_name][key_name] = state_dict[key] + + return layer_dict + + +def _get_lora_layer_values(layer_dict: dict[str, torch.Tensor], alpha: float | None) -> dict[str, torch.Tensor]: + """Convert layer dict keys from PEFT format to internal format.""" + if "lora_A.weight" in layer_dict: + values = { + "lora_down.weight": layer_dict["lora_A.weight"], + "lora_up.weight": layer_dict["lora_B.weight"], + } + if alpha is not None: + values["alpha"] = torch.tensor(alpha) + return values + elif "lora_down.weight" in layer_dict: + return layer_dict + else: + return layer_dict + + +def lora_model_from_anima_state_dict(state_dict: Dict[str, torch.Tensor], alpha: float | None = None) -> ModelPatchRaw: + """Convert an Anima LoRA state dict to a ModelPatchRaw. + + Supports both Kohya-style keys (lora_unet_blocks_0_...) and diffusers PEFT format. + Also supports text encoder LoRA keys (lora_te_layers_0_...) targeting the Qwen3 encoder. + + Args: + state_dict: The LoRA state dict + alpha: The alpha value for LoRA scaling. If None, uses rank as alpha. + + Returns: + A ModelPatchRaw containing the LoRA layers + """ + layers: dict[str, BaseLayerPatch] = {} + + # Detect format + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + is_kohya = any(k.startswith(("lora_unet_", "lora_te_")) for k in str_keys) + + if is_kohya: + # Kohya format: group by layer name (everything before .lora_down/.lora_up/.alpha) + grouped = _group_keys_by_layer(state_dict) + for kohya_layer_name, layer_dict in grouped.items(): + if kohya_layer_name.startswith("lora_te_"): + model_key = _convert_kohya_te_key(kohya_layer_name) + final_key = f"{ANIMA_LORA_QWEN3_PREFIX}{model_key}" + else: + model_key = _convert_kohya_unet_key(kohya_layer_name) + final_key = f"{ANIMA_LORA_TRANSFORMER_PREFIX}{model_key}" + layer = _make_layer_patch(layer_dict) + layers[final_key] = layer + else: + # Diffusers PEFT format + grouped = _group_keys_by_layer(state_dict, extra_suffixes=_PEFT_EXTRA_SUFFIXES) + for layer_key, layer_dict in grouped.items(): + values = _get_lora_layer_values(layer_dict, alpha) + clean_key = layer_key + + # Check for text encoder prefixes + text_encoder_prefixes = [ + "base_model.model.text_encoder.", + "text_encoder.", + ] + + is_text_encoder = False + for prefix in text_encoder_prefixes: + if layer_key.startswith(prefix): + clean_key = layer_key[len(prefix) :] + is_text_encoder = True + break + + # If not text encoder, check transformer prefixes + if not is_text_encoder: + for prefix in [ + "base_model.model.transformer.", + "transformer.", + "diffusion_model.", + ]: + if layer_key.startswith(prefix): + clean_key = layer_key[len(prefix) :] + break + + if is_text_encoder: + final_key = f"{ANIMA_LORA_QWEN3_PREFIX}{clean_key}" + else: + final_key = f"{ANIMA_LORA_TRANSFORMER_PREFIX}{clean_key}" + + layer = _make_layer_patch(values) + layers[final_key] = layer + + return ModelPatchRaw(layers=layers) diff --git a/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py new file mode 100644 index 00000000000..b2109222a31 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/flux_onetrainer_bfl_lora_conversion_utils.py @@ -0,0 +1,168 @@ +"""Utilities for detecting and converting FLUX LoRAs in OneTrainer BFL format. + +This format is produced by newer versions of OneTrainer and uses BFL internal key names +(double_blocks, single_blocks, img_attn, etc.) with a 'transformer.' prefix and +InvokeAI-native LoRA suffixes (lora_down.weight, lora_up.weight, alpha). + +Unlike the standard BFL PEFT format (which uses 'diffusion_model.' prefix and lora_A/lora_B), +this format also has split QKV projections: + - double_blocks.{i}.img_attn.qkv.{0,1,2} (Q, K, V separate) + - double_blocks.{i}.txt_attn.qkv.{0,1,2} (Q, K, V separate) + - single_blocks.{i}.linear1.{0,1,2,3} (Q, K, V, MLP separate) + +Example keys: + transformer.double_blocks.0.img_attn.qkv.0.lora_down.weight + transformer.double_blocks.0.img_attn.qkv.0.lora_up.weight + transformer.double_blocks.0.img_attn.qkv.0.alpha + transformer.single_blocks.0.linear1.3.lora_down.weight + transformer.double_blocks.0.img_mlp.0.lora_down.weight +""" + +import re +from typing import Any, Dict + +import torch + +from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch +from invokeai.backend.patches.layers.merged_layer_patch import MergedLayerPatch, Range +from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict +from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_TRANSFORMER_PREFIX +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw + +_TRANSFORMER_PREFIX = "transformer." + +# Valid LoRA weight suffixes in this format. +_LORA_SUFFIXES = ("lora_down.weight", "lora_up.weight", "alpha") + +# Regex to detect split QKV keys in double blocks: e.g. "double_blocks.0.img_attn.qkv.1" +_SPLIT_QKV_RE = re.compile(r"^(double_blocks\.\d+\.(img_attn|txt_attn)\.qkv)\.\d+$") + +# Regex to detect split linear1 keys in single blocks: e.g. "single_blocks.0.linear1.2" +_SPLIT_LINEAR1_RE = re.compile(r"^(single_blocks\.\d+\.linear1)\.\d+$") + + +def is_state_dict_likely_in_flux_onetrainer_bfl_format( + state_dict: dict[str | int, Any], + metadata: dict[str, Any] | None = None, +) -> bool: + """Checks if the provided state dict is likely in the OneTrainer BFL FLUX LoRA format. + + This format uses BFL internal key names with 'transformer.' prefix and split QKV projections. + """ + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + if not str_keys: + return False + + # All keys must start with 'transformer.' + if not all(k.startswith(_TRANSFORMER_PREFIX) for k in str_keys): + return False + + # All keys must end with recognized LoRA suffixes. + if not all(k.endswith(_LORA_SUFFIXES) for k in str_keys): + return False + + # Must have BFL block structure (double_blocks or single_blocks) under transformer prefix. + has_bfl_blocks = any( + k.startswith("transformer.double_blocks.") or k.startswith("transformer.single_blocks.") for k in str_keys + ) + if not has_bfl_blocks: + return False + + # Must have split QKV pattern (qkv.0, qkv.1, qkv.2) to distinguish from other formats + # that might use transformer. prefix in the future. + has_split_qkv = any(".qkv.0." in k or ".qkv.1." in k or ".qkv.2." in k or ".linear1.0." in k for k in str_keys) + if not has_split_qkv: + return False + + return True + + +def _split_key(key: str) -> tuple[str, str]: + """Split a key into (layer_name, weight_suffix). + + Handles: + - 2-component suffixes ending with '.weight': e.g., 'lora_down.weight' → split at 2nd-to-last dot + - 1-component suffixes: e.g., 'alpha' → split at last dot + """ + if key.endswith(".weight"): + parts = key.rsplit(".", maxsplit=2) + return parts[0], f"{parts[1]}.{parts[2]}" + else: + parts = key.rsplit(".", maxsplit=1) + return parts[0], parts[1] + + +def lora_model_from_flux_onetrainer_bfl_state_dict(state_dict: Dict[str, torch.Tensor]) -> ModelPatchRaw: + """Convert a OneTrainer BFL format FLUX LoRA state dict to a ModelPatchRaw. + + Strips the 'transformer.' prefix, groups by layer, and merges split QKV/linear1 + layers into MergedLayerPatch instances. + """ + # Step 1: Strip prefix and group by layer name. + grouped_state_dict: dict[str, dict[str, torch.Tensor]] = {} + for key, value in state_dict.items(): + if not isinstance(key, str): + continue + + # Strip 'transformer.' prefix. + key = key[len(_TRANSFORMER_PREFIX) :] + + layer_name, suffix = _split_key(key) + + if layer_name not in grouped_state_dict: + grouped_state_dict[layer_name] = {} + grouped_state_dict[layer_name][suffix] = value + + # Step 2: Build LoRA layers, merging split QKV and linear1. + layers: dict[str, BaseLayerPatch] = {} + + # Identify which layers need merging. + merge_groups: dict[str, list[str]] = {} + standalone_keys: list[str] = [] + + for layer_key in grouped_state_dict: + qkv_match = _SPLIT_QKV_RE.match(layer_key) + linear1_match = _SPLIT_LINEAR1_RE.match(layer_key) + + if qkv_match: + parent = qkv_match.group(1) + if parent not in merge_groups: + merge_groups[parent] = [] + merge_groups[parent].append(layer_key) + elif linear1_match: + parent = linear1_match.group(1) + if parent not in merge_groups: + merge_groups[parent] = [] + merge_groups[parent].append(layer_key) + else: + standalone_keys.append(layer_key) + + # Process standalone layers. + for layer_key in standalone_keys: + layer_sd = grouped_state_dict[layer_key] + layers[f"{FLUX_LORA_TRANSFORMER_PREFIX}{layer_key}"] = any_lora_layer_from_state_dict(layer_sd) + + # Process merged layers. + for parent_key, sub_keys in merge_groups.items(): + # Sort by the numeric index at the end (e.g., qkv.0, qkv.1, qkv.2). + sub_keys.sort(key=lambda k: int(k.rsplit(".", maxsplit=1)[1])) + + sub_layers: list[BaseLayerPatch] = [] + sub_ranges: list[Range] = [] + dim_0_offset = 0 + + for sub_key in sub_keys: + layer_sd = grouped_state_dict[sub_key] + sub_layer = any_lora_layer_from_state_dict(layer_sd) + + # Determine the output dimension from the up weight shape. + up_weight = layer_sd["lora_up.weight"] + out_dim = up_weight.shape[0] + + sub_layers.append(sub_layer) + sub_ranges.append(Range(dim_0_offset, dim_0_offset + out_dim)) + dim_0_offset += out_dim + + layers[f"{FLUX_LORA_TRANSFORMER_PREFIX}{parent_key}"] = MergedLayerPatch(sub_layers, sub_ranges) + + return ModelPatchRaw(layers=layers) diff --git a/invokeai/backend/patches/lora_conversions/formats.py b/invokeai/backend/patches/lora_conversions/formats.py index 0b316602fcd..b3e00c288bd 100644 --- a/invokeai/backend/patches/lora_conversions/formats.py +++ b/invokeai/backend/patches/lora_conversions/formats.py @@ -14,6 +14,9 @@ from invokeai.backend.patches.lora_conversions.flux_kohya_lora_conversion_utils import ( is_state_dict_likely_in_flux_kohya_format, ) +from invokeai.backend.patches.lora_conversions.flux_onetrainer_bfl_lora_conversion_utils import ( + is_state_dict_likely_in_flux_onetrainer_bfl_format, +) from invokeai.backend.patches.lora_conversions.flux_onetrainer_lora_conversion_utils import ( is_state_dict_likely_in_flux_onetrainer_format, ) @@ -28,6 +31,8 @@ def flux_format_from_state_dict( ) -> FluxLoRAFormat | None: if is_state_dict_likely_in_flux_kohya_format(state_dict): return FluxLoRAFormat.Kohya + elif is_state_dict_likely_in_flux_onetrainer_bfl_format(state_dict, metadata): + return FluxLoRAFormat.OneTrainerBfl elif is_state_dict_likely_in_flux_onetrainer_format(state_dict): return FluxLoRAFormat.OneTrainer elif is_state_dict_likely_in_flux_diffusers_format(state_dict): diff --git a/invokeai/backend/patches/lora_conversions/qwen_image_lora_constants.py b/invokeai/backend/patches/lora_conversions/qwen_image_lora_constants.py new file mode 100644 index 00000000000..727ee5a4281 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/qwen_image_lora_constants.py @@ -0,0 +1,5 @@ +# Qwen Image Edit LoRA prefix constants +# These prefixes are used for key mapping when applying LoRA patches to Qwen Image Edit models + +# Prefix for Qwen Image Edit transformer LoRA layers +QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX = "lora_transformer-" diff --git a/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py b/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py new file mode 100644 index 00000000000..7fc01f72315 --- /dev/null +++ b/invokeai/backend/patches/lora_conversions/qwen_image_lora_conversion_utils.py @@ -0,0 +1,197 @@ +"""Qwen Image LoRA conversion utilities. + +Qwen Image uses QwenImageTransformer2DModel architecture. +Supports multiple LoRA formats: +- Diffusers/PEFT: transformer_blocks.0.attn.to_k.lora_down.weight +- With prefix: transformer.transformer_blocks.0.attn.to_k.lora_down.weight +- Kohya: lora_unet_transformer_blocks_0_attn_to_k.lora_down.weight (underscores instead of dots) +""" + +import re +from typing import Dict + +import torch + +from invokeai.backend.patches.layers.base_layer_patch import BaseLayerPatch +from invokeai.backend.patches.layers.utils import any_lora_layer_from_state_dict +from invokeai.backend.patches.lora_conversions.qwen_image_lora_constants import ( + QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX, +) +from invokeai.backend.patches.model_patch_raw import ModelPatchRaw + +# Regex for Kohya-format Qwen Image LoRA keys. +# Example: lora_unet_transformer_blocks_0_attn_to_k +# Groups: (block_idx, sub_module_with_underscores) +_KOHYA_KEY_REGEX = re.compile(r"lora_unet_transformer_blocks_(\d+)_(.*)") + +# Mapping from Kohya underscore-separated sub-module names to dot-separated model paths. +# The Kohya format uses underscores everywhere, but some underscores are part of the +# module name (e.g., add_k_proj, to_out). We match the longest prefix first. +_KOHYA_MODULE_MAP: list[tuple[str, str]] = [ + # Attention projections + ("attn_add_k_proj", "attn.add_k_proj"), + ("attn_add_q_proj", "attn.add_q_proj"), + ("attn_add_v_proj", "attn.add_v_proj"), + ("attn_to_add_out", "attn.to_add_out"), + ("attn_to_out_0", "attn.to_out.0"), + ("attn_to_k", "attn.to_k"), + ("attn_to_q", "attn.to_q"), + ("attn_to_v", "attn.to_v"), + # Image stream MLP and modulation + ("img_mlp_net_0_proj", "img_mlp.net.0.proj"), + ("img_mlp_net_2", "img_mlp.net.2"), + ("img_mod_1", "img_mod.1"), + # Text stream MLP and modulation + ("txt_mlp_net_0_proj", "txt_mlp.net.0.proj"), + ("txt_mlp_net_2", "txt_mlp.net.2"), + ("txt_mod_1", "txt_mod.1"), +] + + +def is_state_dict_likely_kohya_qwen_image(state_dict: dict[str | int, torch.Tensor]) -> bool: + """Check if the state dict uses Kohya-format Qwen Image LoRA keys.""" + str_keys = [k for k in state_dict.keys() if isinstance(k, str)] + if not str_keys: + return False + # Check if any key matches the Kohya pattern + return any(k.startswith("lora_unet_transformer_blocks_") for k in str_keys) + + +def _convert_kohya_key(kohya_layer: str) -> str | None: + """Convert a Kohya-format layer name to a dot-separated model module path. + + Example: lora_unet_transformer_blocks_0_attn_to_k -> transformer_blocks.0.attn.to_k + """ + m = _KOHYA_KEY_REGEX.match(kohya_layer) + if not m: + return None + + block_idx = m.group(1) + sub_module = m.group(2) + + for kohya_name, model_path in _KOHYA_MODULE_MAP: + if sub_module == kohya_name: + return f"transformer_blocks.{block_idx}.{model_path}" + + # Fallback: unknown sub-module, return None so caller can warn/skip + return None + + +def lora_model_from_qwen_image_state_dict( + state_dict: Dict[str, torch.Tensor], alpha: float | None = None +) -> ModelPatchRaw: + """Convert a Qwen Image LoRA state dict to a ModelPatchRaw. + + Handles three key formats: + - Diffusers/PEFT: transformer_blocks.0.attn.to_k.lora_down.weight + - With prefix: transformer.transformer_blocks.0.attn.to_k.lora_down.weight + - Kohya: lora_unet_transformer_blocks_0_attn_to_k.lora_down.weight + """ + is_kohya = is_state_dict_likely_kohya_qwen_image(state_dict) + + if is_kohya: + return _convert_kohya_format(state_dict, alpha) + else: + return _convert_diffusers_format(state_dict, alpha) + + +def _convert_kohya_format(state_dict: Dict[str, torch.Tensor], alpha: float | None) -> ModelPatchRaw: + """Convert Kohya-format state dict. Keys are like lora_unet_transformer_blocks_0_attn_to_k.lokr_w1""" + layers: dict[str, BaseLayerPatch] = {} + + # Group by layer (split at first dot: layer_name.param_name) + grouped: dict[str, dict[str, torch.Tensor]] = {} + for key, value in state_dict.items(): + if not isinstance(key, str): + continue + layer_name, param_name = key.split(".", 1) + if layer_name not in grouped: + grouped[layer_name] = {} + grouped[layer_name][param_name] = value + + for kohya_layer, layer_dict in grouped.items(): + model_path = _convert_kohya_key(kohya_layer) + if model_path is None: + continue # Skip unrecognized layers + + layer = any_lora_layer_from_state_dict(layer_dict) + final_key = f"{QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX}{model_path}" + layers[final_key] = layer + + return ModelPatchRaw(layers=layers) + + +def _convert_diffusers_format(state_dict: Dict[str, torch.Tensor], alpha: float | None) -> ModelPatchRaw: + """Convert Diffusers/PEFT format state dict.""" + layers: dict[str, BaseLayerPatch] = {} + + # Some LoRAs use a "transformer." prefix on keys + strip_prefixes = ["transformer."] + + grouped = _group_by_layer(state_dict) + + for layer_key, layer_dict in grouped.items(): + values = _normalize_lora_keys(layer_dict, alpha) + layer = any_lora_layer_from_state_dict(values) + clean_key = layer_key + for prefix in strip_prefixes: + if clean_key.startswith(prefix): + clean_key = clean_key[len(prefix) :] + break + final_key = f"{QWEN_IMAGE_EDIT_LORA_TRANSFORMER_PREFIX}{clean_key}" + layers[final_key] = layer + + return ModelPatchRaw(layers=layers) + + +def _normalize_lora_keys(layer_dict: dict[str, torch.Tensor], alpha: float | None) -> dict[str, torch.Tensor]: + """Normalize LoRA key names to internal format.""" + if "lora_A.weight" in layer_dict: + values: dict[str, torch.Tensor] = { + "lora_down.weight": layer_dict["lora_A.weight"], + "lora_up.weight": layer_dict["lora_B.weight"], + } + if alpha is not None: + values["alpha"] = torch.tensor(alpha) + return values + elif "lora_down.weight" in layer_dict: + return layer_dict + else: + return layer_dict + + +def _group_by_layer(state_dict: Dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]: + """Group state dict keys by layer path.""" + layer_dict: dict[str, dict[str, torch.Tensor]] = {} + + known_suffixes = [ + ".lora_A.weight", + ".lora_B.weight", + ".lora_down.weight", + ".lora_up.weight", + ".dora_scale", + ".alpha", + ] + + for key in state_dict: + if not isinstance(key, str): + continue + + layer_name = None + key_name = None + for suffix in known_suffixes: + if key.endswith(suffix): + layer_name = key[: -len(suffix)] + key_name = suffix[1:] + break + + if layer_name is None: + parts = key.rsplit(".", maxsplit=2) + layer_name = parts[0] + key_name = ".".join(parts[1:]) + + if layer_name not in layer_dict: + layer_dict[layer_name] = {} + layer_dict[layer_name][key_name] = state_dict[key] + + return layer_dict diff --git a/invokeai/backend/stable_diffusion/diffusers_pipeline.py b/invokeai/backend/stable_diffusion/diffusers_pipeline.py index de5253f0733..054e04dcb28 100644 --- a/invokeai/backend/stable_diffusion/diffusers_pipeline.py +++ b/invokeai/backend/stable_diffusion/diffusers_pipeline.py @@ -17,7 +17,7 @@ from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin from diffusers.utils.import_utils import is_xformers_available from pydantic import Field -from transformers import CLIPFeatureExtractor, CLIPTextModel, CLIPTokenizer +from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer from invokeai.app.services.config.config_default import get_config from invokeai.backend.stable_diffusion.diffusion.conditioning_data import IPAdapterData, TextConditioningData @@ -139,7 +139,7 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline): safety_checker ([`StableDiffusionSafetyChecker`]): Classification module that estimates whether generated images could be considered offensive or harmful. Please, refer to the [model card](https://huggingface.co/CompVis/stable-diffusion-v1-4) for details. - feature_extractor ([`CLIPFeatureExtractor`]): + feature_extractor ([`CLIPImageProcessor`]): Model that extracts features from generated images to be used as inputs for the `safety_checker`. """ @@ -151,7 +151,7 @@ def __init__( unet: UNet2DConditionModel, scheduler: KarrasDiffusionSchedulers, safety_checker: Optional[StableDiffusionSafetyChecker], - feature_extractor: Optional[CLIPFeatureExtractor], + feature_extractor: Optional[CLIPImageProcessor], requires_safety_checker: bool = False, ): super().__init__( diff --git a/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py b/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py index 9d1bd676174..6a9959f1e87 100644 --- a/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py +++ b/invokeai/backend/stable_diffusion/diffusion/conditioning_data.py @@ -88,6 +88,48 @@ def to(self, device: torch.device | None = None, dtype: torch.dtype | None = Non return self +@dataclass +class QwenImageConditioningInfo: + """Qwen Image Edit conditioning information from Qwen2.5-VL encoder.""" + + prompt_embeds: torch.Tensor + """Text/image embeddings from Qwen2.5-VL encoder. Shape: (batch_size, seq_len, hidden_size).""" + + prompt_embeds_mask: torch.Tensor | None = None + """Attention mask for prompt_embeds. Shape: (batch_size, seq_len). 1 for valid, 0 for padding.""" + + def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None): + self.prompt_embeds = self.prompt_embeds.to(device=device, dtype=dtype) + if self.prompt_embeds_mask is not None: + self.prompt_embeds_mask = self.prompt_embeds_mask.to(device=device) + return self + + +@dataclass +class AnimaConditioningInfo: + """Anima text conditioning information from Qwen3 0.6B encoder + T5-XXL tokenizer. + + Anima uses a dual-conditioning scheme where Qwen3 hidden states are combined + with T5-XXL token IDs inside the LLM Adapter (part of the transformer). + """ + + qwen3_embeds: torch.Tensor + """Qwen3 0.6B hidden states. Shape: (seq_len, hidden_size) where hidden_size=1024.""" + + t5xxl_ids: torch.Tensor + """T5-XXL token IDs. Shape: (seq_len,).""" + + t5xxl_weights: Optional[torch.Tensor] = None + """Per-token weights for prompt weighting. Shape: (seq_len,). None means uniform weight.""" + + def to(self, device: torch.device | None = None, dtype: torch.dtype | None = None): + self.qwen3_embeds = self.qwen3_embeds.to(device=device, dtype=dtype) + self.t5xxl_ids = self.t5xxl_ids.to(device=device) + if self.t5xxl_weights is not None: + self.t5xxl_weights = self.t5xxl_weights.to(device=device, dtype=dtype) + return self + + @dataclass class ConditioningFieldData: # If you change this class, adding more types, you _must_ update the instantiation of ObjectSerializerDisk in @@ -100,6 +142,8 @@ class ConditioningFieldData: | List[SD3ConditioningInfo] | List[CogView4ConditioningInfo] | List[ZImageConditioningInfo] + | List[QwenImageConditioningInfo] + | List[AnimaConditioningInfo] ) diff --git a/invokeai/frontend/web/package.json b/invokeai/frontend/web/package.json index da4e31142f2..e9a896f1b4e 100644 --- a/invokeai/frontend/web/package.json +++ b/invokeai/frontend/web/package.json @@ -21,7 +21,7 @@ "scripts": { "dev": "vite dev", "dev:host": "vite dev --host", - "build": "pnpm run lint && vite build", + "build": "pnpm run lint && vitest run && vite build", "typegen": "node scripts/typegen.js", "preview": "vite preview", "lint:knip": "knip --tags=-knipignore", @@ -35,6 +35,7 @@ "storybook": "storybook dev -p 6006", "build-storybook": "storybook build", "test": "vitest", + "test:run": "vitest run", "test:ui": "vitest --coverage --ui", "test:no-watch": "vitest --no-watch" }, diff --git a/invokeai/frontend/web/public/locales/en.json b/invokeai/frontend/web/public/locales/en.json index 00eab834ca4..201ea8badba 100644 --- a/invokeai/frontend/web/public/locales/en.json +++ b/invokeai/frontend/web/public/locales/en.json @@ -25,7 +25,8 @@ "rememberMe": "Remember me for 7 days", "signIn": "Sign In", "signingIn": "Signing in...", - "loginFailed": "Login failed. Please check your credentials." + "loginFailed": "Login failed. Please check your credentials.", + "sessionExpired": "Your credentials have expired. Please log in again to resume." }, "setup": { "title": "Welcome to InvokeAI", @@ -160,7 +161,17 @@ "imagesWithCount_other": "{{count}} images", "assetsWithCount_one": "{{count}} asset", "assetsWithCount_other": "{{count}} assets", - "updateBoardError": "Error updating board" + "updateBoardError": "Error updating board", + "setBoardVisibility": "Set Board Visibility", + "setVisibilityPrivate": "Set Private", + "setVisibilityShared": "Set Shared", + "setVisibilityPublic": "Set Public", + "visibilityPrivate": "Private", + "visibilityShared": "Shared", + "visibilityPublic": "Public", + "visibilityBadgeShared": "Shared board", + "visibilityBadgePublic": "Public board", + "updateBoardVisibilityError": "Error updating board visibility" }, "accordions": { "generation": { @@ -223,6 +234,7 @@ "discordLabel": "Discord", "dontAskMeAgain": "Don't ask me again", "dontShowMeThese": "Don't show me these", + "editName": "Edit name", "editor": "Editor", "error": "Error", "error_withCount_one": "{{count}} error", @@ -230,50 +242,67 @@ "model_withCount_one": "{{count}} model", "model_withCount_other": "{{count}} models", "file": "File", + "fitView": "Fit View", "folder": "Folder", "format": "format", "githubLabel": "Github", "goTo": "Go to", "hotkeysLabel": "Hotkeys", - "loadingImage": "Loading Image", - "loadingModel": "Loading Model", + "hex": "Hex", "imageFailedToLoad": "Unable to Load Image", "img2img": "Image To Image", "inpaint": "inpaint", "input": "Input", "installed": "Installed", + "json": "JSON", "languagePickerLabel": "Language", "linear": "Linear", "load": "Load", "loading": "Loading", + "loadingImage": "Loading Image", + "loadingModel": "Loading Model", "localSystem": "Local System", + "minimize": "Minimize", + "next": "Next", + "noMatchingItems": "No matching items", + "notifications": "Notifications", "learnMore": "Learn More", "modelManager": "Model Manager", "noMatches": "No matches", "noOptions": "No options", "nodes": "Workflows", "notInstalled": "Not $t(common.installed)", + "openSlider": "Open slider", "openInNewTab": "Open in New Tab", "openInViewer": "Open in Viewer", "orderBy": "Order By", "outpaint": "outpaint", "outputs": "Outputs", "postprocessing": "Post Processing", + "previous": "Previous", "random": "Random", + "removeFromCollection": "Remove from Collection", "reportBugLabel": "Report Bug", + "resetView": "Reset View", "safetensors": "Safetensors", "save": "Save", "saveAs": "Save As", "saveChanges": "Save Changes", + "saveToAssets": "Save to Assets", + "settings": "Settings", "settingsLabel": "Settings", "simple": "Simple", "somethingWentWrong": "Something went wrong", "statusDisconnected": "Disconnected", "template": "Template", + "toggleRgbHex": "Toggle RGB/HEX", "toResolve": "To resolve", "txt2img": "Text To Image", "unknown": "Unknown", + "unpin": "Unpin", "upload": "Upload", + "zoomIn": "Zoom In", + "zoomOut": "Zoom Out", "updated": "Updated", "created": "Created", "prevPage": "Previous Page", @@ -343,13 +372,19 @@ "discard": "Discard", "noPromptHistory": "No prompt history recorded.", "noMatchingPrompts": "No matching prompts in history.", - "toSwitchBetweenPrompts": "to switch between prompts." + "toSwitchBetweenPrompts": "to switch between prompts.", + "promptHistory": "Prompt History", + "clearHistory": "Clear History", + "usePrompt": "Use prompt", + "searchPrompts": "Search..." }, "queue": { "queue": "Queue", "queueFront": "Add to Front of Queue", "queueBack": "Add to Queue", + "queueActionsMenu": "Queue Actions Menu", "queueEmpty": "Queue Empty", + "queueItem": "Queue Item", "enqueueing": "Queueing Batch", "resume": "Resume", "resumeTooltip": "Resume Processor", @@ -492,7 +527,10 @@ "imagesSettings": "Gallery Images Settings", "jump": "Jump", "loading": "Loading", + "loadingGallery": "Loading gallery...", + "loadingMetadata": "Loading metadata...", "newestFirst": "Newest First", + "noImagesFound": "No images found", "oldestFirst": "Oldest First", "sortDirection": "Sort Direction", "showStarredImagesFirst": "Show Starred Images First", @@ -504,6 +542,8 @@ "unableToLoad": "Unable to load Gallery", "deleteSelection": "Delete Selection", "downloadSelection": "Download Selection", + "bulkDownloadReady": "Download ready", + "clickToDownload": "Click here to download", "bulkDownloadRequested": "Preparing Download", "bulkDownloadRequestedDesc": "Your download request is being prepared. This may take a few moments.", "bulkDownloadRequestFailed": "Problem Preparing Download", @@ -935,7 +975,8 @@ } }, "lora": { - "weight": "Weight" + "weight": "Weight", + "removeLoRA": "Remove LoRA" }, "metadata": { "allPrompts": "All Prompts", @@ -970,6 +1011,7 @@ "seedVarianceEnabled": "Seed Variance Enabled", "seedVarianceStrength": "Seed Variance Strength", "seedVarianceRandomizePercent": "Seed Variance Randomize %", + "zImageShift": "Z-Image Shift", "seed": "Seed", "steps": "Steps", "strength": "Image to image strength", @@ -982,6 +1024,14 @@ "modelManager": { "active": "active", "actions": "Bulk Actions", + "deleteModelsConfirm": "Are you sure you want to delete {{count}} model(s)? This action cannot be undone.", + "deleteWarning": "Models in your Invoke models directory will be permanently deleted from disk.", + "modelsDeleted": "Successfully deleted {{count}} model(s)", + "modelsDeleteFailed": "Failed to delete models", + "someModelsFailedToDelete": "{{count}} model(s) could not be deleted", + "modelsDeletedPartial": "Partially completed", + "someModelsDeleted": "{{deleted}} deleted, {{failed}} failed", + "modelsDeleteError": "Error deleting models", "pause": "Pause", "pauseAll": "Pause All", "pauseAllTooltip": "Pause all active downloads", @@ -1203,11 +1253,18 @@ "triggerPhrases": "Trigger Phrases", "loraTriggerPhrases": "LoRA Trigger Phrases", "mainModelTriggerPhrases": "Main Model Trigger Phrases", + "queueEmpty": "The install queue is empty.", "selectAll": "Select All", "selectModelToView": "Select a model to view its details", "typePhraseHere": "Type phrase here", "t5Encoder": "T5 Encoder", "qwen3Encoder": "Qwen3 Encoder", + "animaVae": "VAE", + "animaVaePlaceholder": "Select Anima-compatible VAE", + "animaQwen3Encoder": "Qwen3 0.6B Encoder", + "animaQwen3EncoderPlaceholder": "Select Qwen3 0.6B encoder", + "animaT5Encoder": "T5-XXL Encoder", + "animaT5EncoderPlaceholder": "Select T5-XXL encoder", "zImageVae": "VAE (optional)", "zImageVaePlaceholder": "From VAE source model", "zImageQwen3Encoder": "Qwen3 Encoder (optional)", @@ -1218,6 +1275,12 @@ "flux2KleinVaePlaceholder": "From main model", "flux2KleinQwen3Encoder": "Qwen3 Encoder (optional)", "flux2KleinQwen3EncoderPlaceholder": "From main model", + "qwenImageComponentSource": "VAE/Encoder Source (Diffusers)", + "qwenImageComponentSourcePlaceholder": "Required for GGUF models", + "qwenImageQuantization": "Encoder Quantization", + "qwenImageQuantizationNone": "None (bf16)", + "qwenImageQuantizationInt8": "8-bit (int8)", + "qwenImageQuantizationNf4": "4-bit (nf4)", "upcastAttention": "Upcast Attention", "uploadImage": "Upload Image", "urlOrLocalPath": "URL or Local Path", @@ -1326,6 +1389,7 @@ "noWorkflows": "No Workflows", "noMatchingWorkflows": "No Matching Workflows", "noWorkflow": "No Workflow", + "noWorkflowToSave": "No workflow to save", "unableToUpdateNode": "Node update failed: node {{node}} of type {{type}} (may require deleting and recreating)", "mismatchedVersion": "Invalid node: node {{node}} of type {{type}} has mismatched version (try updating?)", "missingTemplate": "Invalid node: node {{node}} of type {{type}} missing template (not installed?)", @@ -1349,6 +1413,7 @@ "nodeOpacity": "Node Opacity", "nodeVersion": "Node Version", "noOutputRecorded": "No outputs recorded", + "nodeData": "Node Data", "notes": "Notes", "description": "Description", "notesDescription": "Add notes about your workflow", @@ -1482,6 +1547,7 @@ "info": "Info", "invoke": { "addingImagesTo": "Adding images to", + "boardNotWritable": "You do not have write access to board \"{{boardName}}\". Select a board you own or switch to Uncategorized.", "modelDisabledForTrial": "Generating with {{modelName}} is not available on trial accounts. Visit your account settings to upgrade.", "invoke": "Invoke", "missingFieldTemplate": "Missing field template", @@ -1508,8 +1574,12 @@ "noFLUXVAEModelSelected": "No VAE model selected for FLUX generation", "noCLIPEmbedModelSelected": "No CLIP Embed model selected for FLUX generation", "noQwen3EncoderModelSelected": "No Qwen3 Encoder model selected for FLUX2 Klein generation", + "noQwenImageComponentSourceSelected": "GGUF Qwen Image models require a Diffusers Component Source for VAE/encoder", "noZImageVaeSourceSelected": "No VAE source: Select VAE (FLUX) or Qwen3 Source model", "noZImageQwen3EncoderSourceSelected": "No Qwen3 Encoder source: Select Qwen3 Encoder or Qwen3 Source model", + "noAnimaVaeModelSelected": "No Anima VAE model selected", + "noAnimaQwen3EncoderModelSelected": "No Anima Qwen3 Encoder model selected", + "noAnimaT5EncoderModelSelected": "No Anima T5 Encoder model selected", "fluxModelIncompatibleBboxWidth": "$t(parameters.invoke.fluxRequiresDimensionsToBeMultipleOf16), bbox width is {{width}}", "fluxModelIncompatibleBboxHeight": "$t(parameters.invoke.fluxRequiresDimensionsToBeMultipleOf16), bbox height is {{height}}", "fluxModelIncompatibleScaledBboxWidth": "$t(parameters.invoke.fluxRequiresDimensionsToBeMultipleOf16), scaled bbox width is {{width}}", @@ -1558,6 +1628,7 @@ "sendToCanvas": "Send To Canvas", "sendToUpscale": "Send To Upscale", "showOptionsPanel": "Show Side Panel (O or T)", + "shift": "Shift", "shuffle": "Shuffle Seed", "steps": "Steps", "strength": "Strength", @@ -1595,6 +1666,7 @@ "perPromptDesc": "Use a different seed for each image" }, "loading": "Generating Dynamic Prompts...", + "problemGeneratingPrompts": "Problem generating prompts", "promptsToGenerate": "Prompts to Generate" }, "sdxl": { @@ -2380,6 +2452,27 @@ "pullBboxIntoReferenceImageError": "Problem Pulling BBox Into ReferenceImage", "addAdjustments": "Add Adjustments", "removeAdjustments": "Remove Adjustments", + "workflowIntegration": { + "title": "Run Workflow on Canvas", + "description": "Select a workflow with a Canvas Output node and an image parameter to run on the current canvas layer. You can adjust parameters before executing. The result will be added back to the canvas.", + "execute": "Execute Workflow", + "executing": "Executing...", + "runWorkflow": "Run Workflow", + "filteringWorkflows": "Filtering workflows...", + "loadingWorkflows": "Loading workflows...", + "noWorkflowsFound": "No workflows found.", + "noWorkflowsWithImageField": "No compatible workflows found. A workflow needs a Form Builder with an image input field and a Canvas Output node.", + "selectWorkflow": "Select Workflow", + "selectPlaceholder": "Choose a workflow...", + "unnamedWorkflow": "Unnamed Workflow", + "loadingParameters": "Loading workflow parameters...", + "noFormBuilderError": "This workflow has no form builder and cannot be used. Please select a different workflow.", + "imageFieldSelected": "This field will receive the canvas image", + "imageFieldNotSelected": "Click to use this field for canvas image", + "executionStarted": "Workflow execution started", + "executionStartedDescription": "The result will appear in the staging area when complete.", + "executionFailed": "Failed to execute workflow" + }, "compositeOperation": { "label": "Blend Mode", "add": "Add Blend Mode", @@ -2452,6 +2545,8 @@ "disableAutoNegative": "Disable Auto Negative", "deletePrompt": "Delete Prompt", "deleteReferenceImage": "Delete Reference Image", + "disableReferenceImage": "Disable Reference Image", + "enableReferenceImage": "Enable Reference Image", "showHUD": "Show HUD", "rectangle": "Rectangle", "maskFill": "Mask Fill", @@ -2472,6 +2567,7 @@ "controlLayer": "Control Layer", "inpaintMask": "Inpaint Mask", "invertMask": "Invert Mask", + "invertRegion": "Invert Region", "regionalGuidance": "Regional Guidance", "referenceImageRegional": "Reference Image (Regional)", "referenceImageGlobal": "Reference Image (Global)", @@ -2479,7 +2575,10 @@ "asRasterLayerResize": "As $t(controlLayers.rasterLayer) (Resize)", "asControlLayer": "As $t(controlLayers.controlLayer)", "asControlLayerResize": "As $t(controlLayers.controlLayer) (Resize)", + "invalidReferenceImage": "Invalid Reference Image:", "referenceImage": "Reference Image", + "removeImageFromCollection": "Remove Image from Collection", + "selectRefImage": "Select Ref Image", "maxRefImages": "Max Ref Images", "useAsReferenceImage": "Use as Reference Image", "regionalReferenceImage": "Regional Reference Image", @@ -2500,7 +2599,11 @@ "alignLeft": "Align Left", "alignCenter": "Align Center", "alignRight": "Align Right", - "px": "px" + "px": "px", + "lineHeight": "Spacing", + "lineHeightDense": "Dense", + "lineHeightNormal": "Normal", + "lineHeightSpacious": "Spacious" }, "newCanvasFromImage": "New Canvas from Image", "newImg2ImgCanvasFromImage": "New Img2Img from Image", @@ -2647,7 +2750,8 @@ "crosshatch": "Crosshatch", "vertical": "Vertical", "horizontal": "Horizontal", - "diagonal": "Diagonal" + "diagonal": "Diagonal", + "switchColors": "Switch FG/BG (X)" }, "gradient": { "linear": "Linear", @@ -2911,8 +3015,11 @@ }, "autoSwitch": { "off": "Off", + "doNotAutoSwitch": "Do not auto-switch", "switchOnStart": "On Start", - "switchOnFinish": "On Finish" + "switchOnStartDesc": "Switch on start", + "switchOnFinish": "On Finish", + "switchOnFinishDesc": "Switch on finish" } }, "upscaling": { @@ -3135,9 +3242,21 @@ "watchRecentReleaseVideos": "Watch Recent Release Videos", "watchUiUpdatesOverview": "Watch UI Updates Overview" }, + "cropper": { + "cropImage": "Crop Image", + "aspectRatio": "Aspect Ratio", + "free": "Free", + "mouseWheelZoom": "Mouse wheel: Zoom", + "spaceDragPan": "Space + Drag: Pan", + "dragCropBoxToAdjust": "Drag crop box or handles to adjust" + }, "supportVideos": { "supportVideos": "Support Videos", "gettingStarted": "Getting Started", + "gettingStartedPlaylist": "Getting Started playlist", + "studioSessionsPlaylist": "Studio Sessions playlist", + "discord": "Discord", + "github": "GitHub", "watch": "Watch", "studioSessionsDesc": "Join our to participate in the live sessions and ask questions. Sessions are uploaded to the playlist the following week.", "videos": { diff --git a/invokeai/frontend/web/public/locales/es.json b/invokeai/frontend/web/public/locales/es.json index 4c58ea87f54..8f68ea585c2 100644 --- a/invokeai/frontend/web/public/locales/es.json +++ b/invokeai/frontend/web/public/locales/es.json @@ -92,7 +92,9 @@ "toResolve": "Para resolver", "outpaint": "outpaint", "simple": "Sencillo", - "close": "Cerrar" + "close": "Cerrar", + "board": "Tablero", + "crop": "Cortar" }, "gallery": { "galleryImageSize": "Tamaño de la imagen", @@ -327,7 +329,7 @@ "movingImagesToBoard_one": "Moviendo {{count}} imagen al panel:", "movingImagesToBoard_many": "Moviendo {{count}} imágenes al panel:", "movingImagesToBoard_other": "Moviendo {{count}} imágenes al panel:", - "bottomMessage": "Al eliminar este panel y las imágenes que contiene, se restablecerán las funciones que los estén utilizando actualmente.", + "bottomMessage": "Al eliminarlas imágenes, se restablecerán las funcionalidades que actualmente las estén utilizando.", "deleteBoardAndImages": "Borrar el panel y las imágenes", "loading": "Cargando...", "deletedBoardsCannotbeRestored": "Los paneles eliminados no se pueden restaurar. Al Seleccionar 'Borrar solo el panel' transferirá las imágenes a un estado sin categorizar.", @@ -354,9 +356,21 @@ "unarchiveBoard": "Desarchivar el panel", "noBoards": "No hay paneles {{boardType}}", "shared": "Paneles compartidos", - "deletedPrivateBoardsCannotbeRestored": "Los paneles eliminados no se pueden restaurar. Al elegir \"Eliminar solo el panel\", las imágenes se colocan en un estado privado y sin categoría para el creador de la imagen.", + "deletedPrivateBoardsCannotbeRestored": "Los paneles eliminados no se pueden restaurar. Al elegir \"Eliminar solo el panel\", las imágenes se colocarán en un estado privado y sin categoría para el creador de la imagen.", "private": "Paneles privados", - "updateBoardError": "No se pudo actualizar el panel" + "updateBoardError": "No se pudo actualizar el panel", + "pause": "Pausa", + "resume": "Reanudar", + "restartFailed": "Reinicio fallido", + "restartFile": "Reiniciar archivo", + "restartRequired": "Reinicio requerido", + "resumeRefused": "Reanudación rechazada por el servidor. Reinicio requerido.", + "uncategorizedImages": "Imágenes sin categoría", + "deleteAllUncategorizedImages": "Eliminar todas las imágenes sin categoría", + "deletedImagesCannotBeRestored": "Las imágenes eliminadas no pueden ser restauradas.", + "hideBoards": "Ocultar tableros", + "locateInGalery": "Ubicar en galeria", + "viewBoards": "Ver paneles" }, "accordions": { "compositing": { @@ -867,5 +881,100 @@ "noModelsInstalled": "Parece que no tienes ningún modelo instalado", "gettingStartedSeries": "¿Desea más orientación? Consulte nuestra Serie de introducción para obtener consejos sobre cómo aprovechar todo el potencial de Invoke Studio.", "toGetStartedLocal": "Para empezar, asegúrate de descargar o importar los modelos necesarios para ejecutar Invoke. A continuación, introduzca un mensaje en el cuadro y haga clic en Invocar para generar su primera imagen. Seleccione una plantilla para mejorar los resultados. Puede elegir guardar sus imágenes directamente en Galería o editarlas en el Lienzo." + }, + "auth": { + "login": { + "title": "Iniciar sesión en InvokeAI", + "email": "Email", + "emailPlaceholder": "Email", + "password": "Contraseña", + "passwordPlaceholder": "Contraseña", + "rememberMe": "Recordarme por 7 días", + "signIn": "Iniciar sesión", + "signingIn": "Iniciando sesión...", + "loginFailed": "Inicio de sesión fallido. Por favor revise sus credenciales." + }, + "setup": { + "title": "Bienvenido a InvokeAI", + "subtitle": "Configure su cuenta de administrador para empezar", + "email": "Email", + "emailPlaceholder": "admin@example.com", + "emailHelper": "Este será su nombre de usuario para iniciar sesión", + "displayName": "Nombre para mostrar", + "displayNamePlaceholder": "Administrador", + "displayNameHelper": "Su nombre como se mostrará en la aplicación", + "password": "Contraseña", + "passwordPlaceholder": "Contraseña", + "passwordHelper": "Debe tener al menos 8 caracteres con mayúsculas, minúsculas y números", + "passwordTooShort": "La contraseña debe tener al menos 8 caracteres", + "passwordMissingRequirements": "La contraseña debe contener mayúsculas, minúsculas y numeros", + "confirmPassword": "Confirmar contraseña", + "confirmPasswordPlaceholder": "Confirmar contraseña", + "passwordsDoNotMatch": "Las contraseñas no coinciden", + "createAccount": "Crear cuenta de administrador", + "creatingAccount": "Configurando...", + "setupFailed": "Configuración fallida. Por favor intente nuevamente.", + "passwordHelperRelaxed": "Ingrese una contraseña (se mostrará la fortaleza)" + }, + "userMenu": "Menu de usuario", + "admin": "Administrador", + "logout": "Cerrar Sesión", + "adminOnlyFeature": "Esta funcionalidad solo esta disponible para administradores.", + "profile": { + "menuItem": "Mi perfil", + "title": "Mi perfil", + "email": "Email", + "emailReadOnly": "La dirección de email no puede ser cambiada", + "displayName": "Nombre para mostrar", + "displayNamePlaceholder": "Su nombre", + "changePassword": "Cambiar contraseña", + "currentPassword": "Contraseña Actual", + "currentPasswordPlaceholder": "Contraseña Actual", + "newPassword": "Nueva contraseña", + "newPasswordPlaceholder": "Nueva contraseña", + "confirmPassword": "Confirmar nueva contraseña", + "confirmPasswordPlaceholder": "Confirmar nueva contraseña", + "passwordsDoNotMatch": "Las contraseñas no coinciden", + "saveSuccess": "Perfil actualizado correctamente", + "saveFailed": "Falló el guardado del perfil. Por favor intente nuevamente." + }, + "userManagement": { + "menuItem": "Administración de usuario", + "title": "Administración de usuario", + "email": "Email", + "emailPlaceholder": "user@example.com", + "displayName": "Nombre para mostrar", + "displayNamePlaceholder": "Nombre para mostrar", + "password": "Contraseña", + "passwordPlaceholder": "Contraseña", + "newPassword": "Nueva contraseña", + "newPasswordPlaceholder": "Deje en blanco para conservar la contraseña actual", + "role": "Rol", + "status": "Estado", + "actions": "Acciones", + "isAdmin": "Administrador", + "user": "Usuario", + "you": "Tu", + "createUser": "Crear usuario", + "editUser": "Editar usuario", + "deleteUser": "Eliminar usuario", + "deleteConfirm": "Esta seguro que desea eliminar {{name}}? Esta accion no se podrá revertir.", + "generatePassword": "Generar contraseña robusta", + "showPassword": "Mostrar contraseña", + "hidePassword": "Ocultar contraseña", + "activate": "Activar", + "deactivate": "Desactivar", + "saveFailed": "Fallo al guardar usuario. Por favor intente nuevamente.", + "deleteFailed": "Fallo al borrar usuario. Por favor intente nuevamente.", + "loadFailed": "Fallo al cargar usuarios.", + "back": "Atras", + "cannotDeleteSelf": "Usted no puede eliminar su propia cuenta", + "cannotDeactivateSelf": "Usted no puede desactivar su propia cuenta" + }, + "passwordStrength": { + "weak": "Contraseña debil", + "moderate": "Contraseña moderada", + "strong": "Contraseña fuerte" + } } } diff --git a/invokeai/frontend/web/public/locales/it.json b/invokeai/frontend/web/public/locales/it.json index cafd91d2a46..98039951943 100644 --- a/invokeai/frontend/web/public/locales/it.json +++ b/invokeai/frontend/web/public/locales/it.json @@ -844,7 +844,16 @@ "settingsImportedPartial": "Impostazioni del modello parzialmente importate. Le impostazioni incompatibili sono state ignorate: {{fields}}", "settingsImportFailed": "Impossibile importare le impostazioni del modello", "settingsImportIncompatible": "Il file delle impostazioni non contiene impostazioni compatibili per questo tipo di modello", - "settingsImportInvalidFile": "File di impostazioni non valido" + "settingsImportInvalidFile": "File di impostazioni non valido", + "reidentifyModels": "Re-identificare i modelli", + "reidentifyModelsConfirm": "Sei sicuro di voler re-identificare {{count}} modello(i)? Questa operazione eseguirà una nuova scansione dei relativi file dei pesi per determinarne il formato e le impostazioni corrette.", + "reidentifyWarning": "Questa operazione ripristinerà tutte le impostazioni personalizzate che potresti aver applicato a questi modelli.", + "modelsReidentified": "{{count}} modello(i) re-identificato(i) con successo", + "modelsReidentifyFailed": "Impossibile re-identificare i modelli", + "someModelsFailedToReidentify": "Non è stato possibile re-identificare {{count}} modello(i)", + "modelsReidentifiedPartial": "Completato parzialmente", + "someModelsReidentified": "{{succeeded}} re-identificato(i), {{failed}} fallito(i)", + "modelsReidentifyError": "Errore nella re-identificazione dei modelli" }, "parameters": { "images": "Immagini", @@ -2585,7 +2594,9 @@ "saveToGallery": "Salva nella Galleria", "previous": "Precedente", "showResultsOn": "Visualizzare i risultati", - "showResultsOff": "Nascondere i risultati" + "showResultsOff": "Nascondere i risultati", + "hideThumbnails": "Nascondi le miniature", + "showThumbnails": "Mostra miniature" }, "HUD": { "bbox": "Riquadro di delimitazione", @@ -2826,6 +2837,27 @@ "alignLeft": "Allinea a sinistra", "alignCenter": "Allinea al centro", "alignRight": "Allinea a destra" + }, + "workflowIntegration": { + "title": "Eseguire il flusso di lavoro sula Tela", + "description": "Seleziona un flusso di lavoro con un nodo Output su tela e un parametro immagine da eseguire sul livello corrente della tela. Puoi regolare i parametri prima dell'esecuzione. Il risultato verrà aggiunto nuovamente alla tela.", + "execute": "Eseguire il flusso di lavoro", + "executing": "Esecuzione in corso...", + "runWorkflow": "Avvia il flusso di lavoro", + "filteringWorkflows": "Filtraggio dei flussi di lavoro...", + "loadingWorkflows": "Caricamento dei flussi di lavoro...", + "noWorkflowsFound": "Nessun flusso di lavoro trovato.", + "noWorkflowsWithImageField": "Nessun flusso di lavoro compatibile trovato. Un flusso di lavoro richiede un Generatore Modello con un campo di input immagine e un nodo Output su tela.", + "selectWorkflow": "Seleziona il flusso di lavoro", + "selectPlaceholder": "Scegli un flusso di lavoro...", + "unnamedWorkflow": "Flusso di lavoro senza nome", + "loadingParameters": "Caricamento dei parametri del flusso di lavoro in corso...", + "noFormBuilderError": "Questo flusso di lavoro non dispone di un generatore di moduli e non può essere utilizzato. Selezionare un flusso di lavoro diverso.", + "imageFieldSelected": "Questo campo riceverà l'immagine della tela", + "imageFieldNotSelected": "Fai clic su questo campo per usarlo per l'immagine sulla tela", + "executionStarted": "L'esecuzione del flusso di lavoro è stata avviata", + "executionStartedDescription": "Il risultato apparirà nell'area di lavoro una volta completata l'operazione.", + "executionFailed": "Impossibile eseguire il flusso di lavoro" } }, "ui": { @@ -3009,9 +3041,9 @@ "readReleaseNotes": "Leggi le note di rilascio", "watchRecentReleaseVideos": "Guarda i video su questa versione", "items": [ - "Supporto FLUX.2 Klein: InvokeAI ora supporta i nuovi modelli FLUX.2 Klein (varianti 4B e 9B) con formati GGUF, FP8 e Diffusers. Le funzionalità includono txt2img, img2img, inpainting e outpainting. Consultare la sezione \"Modelli di partenza\" per iniziare.", - "Il supporto DyPE per i modelli FLUX migliora le immagini ad alta risoluzione (da >1536 px fino a 4K). Vai alla sezione \"Opzioni avanzate\" per attivarlo.", - "Diversità Z-Image Turbo: attiva 'Seed Variance Enhancer' in 'Opzioni avanzate' per aggiungere diversità alle tue generazioni ZiT." + "La modalità multiutente supporta più utenti isolati sullo stesso server.", + "Supporto migliorato per i modelli Z-Image e FLUX.2.", + "Numerosi miglioramenti dell'interfaccia utente e nuove funzionalità Tela." ], "watchUiUpdatesOverview": "Guarda la panoramica degli aggiornamenti dell'interfaccia utente", "takeUserSurvey": "📣 Facci sapere cosa ne pensi di InvokeAI. Partecipa al nostro sondaggio sull'esperienza utente!" @@ -3072,7 +3104,8 @@ "rememberMe": "Ricordami per 7 giorni", "signIn": "Accedi", "signingIn": "Accesso in corso...", - "loginFailed": "Accesso non riuscito. Controlla le tue credenziali." + "loginFailed": "Accesso non riuscito. Controlla le tue credenziali.", + "sessionExpired": "Le tue credenziali sono scadute. Effettua nuovamente l'accesso per continuare." }, "setup": { "title": "Benvenuti a InvokeAI", diff --git a/invokeai/frontend/web/public/locales/ja.json b/invokeai/frontend/web/public/locales/ja.json index 7acf6c98cab..ed8e4386939 100644 --- a/invokeai/frontend/web/public/locales/ja.json +++ b/invokeai/frontend/web/public/locales/ja.json @@ -157,7 +157,7 @@ "noImageSelected": "画像が選択されていません", "deleteSelection": "選択中のものを削除", "downloadSelection": "選択中のものをダウンロード", - "starImage": "スターをつける", + "starImage": "スター", "viewerImage": "閲覧画像", "compareImage": "比較画像", "openInViewer": "ビューアで開く", @@ -190,7 +190,7 @@ "selectAllOnPage": "ページ上のすべてを選択", "images": "画像", "assetsTab": "プロジェクトで使用するためにアップロードされたファイル。", - "imagesTab": "Invoke内で作成および保存された画像。", + "imagesTab": "Invoke内であなたが作成および保存した画像。", "assets": "アセット", "useForPromptGeneration": "プロンプト生成に使用する", "jump": "ジャンプ", @@ -198,7 +198,8 @@ "unableToLoad": "ギャラリーを読み込めません", "selectAnImageToCompare": "比較する画像を選択", "openViewer": "ビューアーを開く", - "closeViewer": "ビューアーを閉じる" + "closeViewer": "ビューアーを閉じる", + "usePagedGalleryView": "ページ型ギャラリービューを使う" }, "hotkeys": { "searchHotkeys": "ホットキーを検索", @@ -515,6 +516,9 @@ "title": "生成タブを選択", "desc": "生成タブを選択。", "key": "1" + }, + "promptWeightUp": { + "title": "選択したプロンプトの重みを増加" } }, "hotkeys": "ホットキー", @@ -568,7 +572,30 @@ "title": "画像にスターを付ける/スターを外す", "desc": "選択した画像にスターを付けたり、スターを外したりします。" } - } + }, + "editMode": "編集モード", + "viewMode": "ビューモード", + "editHotkey": "ホットキーの編集", + "addHotkey": "ホットキーの追加", + "resetToDefault": "デフォルトにリセット", + "resetAll": "全てをデフォルトにリセット", + "resetAllConfirmation": "すべてのホットキーをデフォルトに戻してよろしいですか？この操作は取り消せません。", + "enterHotkeys": "カンマ区切りでホットキーを入力してください", + "save": "保存", + "cancel": "キャンセル", + "modifiers": "モディファイア", + "syntaxHelp": "構文のヘルプ", + "multipleHotkeys": "カンマで区切られた複数のホットキー", + "help": "ヘルプ", + "noHotkeysRecorded": "まだホットキーが記録されていません", + "pressKeys": "キーを押してください...", + "setHotkey": "セット", + "setAnother": "他をセット", + "removeLastHotkey": "最後のホットキーを削除", + "clearAll": "全てをクリア", + "duplicateWarning": "このホットキーはすでに記録済みです", + "conflictWarning": "はすでに \"{{hotkeyTitle}}\" で使われています", + "thisHotkey": "このホットキー" }, "modelManager": { "modelManager": "モデルマネージャ", @@ -1083,7 +1110,7 @@ "batchQueuedDesc_other": "{{count}} セッションをキューの{{direction}}に追加しました", "graphQueued": "グラフをキューに追加しました", "batch": "バッチ", - "clearQueueAlertDialog": "キューをクリアすると、処理中の項目は直ちにキャンセルされ、キューは完全にクリアされます。保留中のフィルターもキャンセルされます。", + "clearQueueAlertDialog": "キューをクリアすると、処理中の項目は直ちにキャンセルされ、キューは完全にクリアされます。保留中のフィルターもキャンセルされ、ステージングエリアもリセットされます。", "pending": "保留中", "resumeFailed": "処理の再開に問題があります", "clear": "クリア", @@ -1135,7 +1162,13 @@ "sortColumn": "列の並べ替え", "sortBy": "{{column}}で並べ替え", "sortOrderAscending": "昇順", - "sortOrderDescending": "降順" + "sortOrderDescending": "降順", + "cancelFailedAccessDenied": "アイテムのキャンセル中に問題が発生しました：アクセスが拒否されました", + "clearFailedAccessDenied": "キューのクリア中に問題が発生しました：アクセスが拒否されました", + "paused": "一時停止中", + "user": "ユーザー", + "fieldValuesHidden": "<非表示>", + "cannotViewDetails": "このキューアイテムを閲覧する権限がありません" }, "models": { "noMatchingModels": "一致するモデルがありません", @@ -1381,7 +1414,13 @@ "deletedImagesCannotBeRestored": "削除された画像は復元できません。", "hideBoards": "ボードを隠す", "locateInGalery": "ギャラリーで検索", - "viewBoards": "ボードを表示" + "viewBoards": "ボードを表示", + "pause": "一時停止", + "resume": "再開", + "restartFailed": "再起動に失敗しました", + "restartFile": "ファイルを再起動", + "restartRequired": "再起動が必要です", + "resumeRefused": "サーバーで再開が拒否されました。再起動が必要です。" }, "invocationCache": { "invocationCache": "呼び出しキャッシュ", @@ -2786,5 +2825,100 @@ }, "lora": { "weight": "重み" + }, + "auth": { + "login": { + "title": "Invokeにサインイン", + "email": "Eメール", + "emailPlaceholder": "Eメール", + "password": "パスワード", + "passwordPlaceholder": "パスワード", + "rememberMe": "7日間は記憶", + "signIn": "サインイン", + "signingIn": "サインイン中...", + "loginFailed": "ログインに失敗しました。正しい内容かを確認してください。" + }, + "setup": { + "title": "Invokeへようこそ", + "subtitle": "管理者アカウントをセットアップします", + "email": "Eメール", + "emailPlaceholder": "hoge@example.com", + "emailHelper": "これはサインインに使うユーザー名になります", + "displayName": "表示名", + "displayNamePlaceholder": "管理者", + "displayNameHelper": "アプリケーションの中で表示される名前です", + "password": "パスワード", + "passwordPlaceholder": "パスワード", + "passwordHelper": "大文字、小文字、数字を組み合わせた8文字以上", + "passwordTooShort": "パスワードは8文字以上である必要があります", + "passwordMissingRequirements": "パスワードは小文字、大文字、数字を含まなければなりません", + "confirmPassword": "パスワードの確認", + "confirmPasswordPlaceholder": "パスワードの確認", + "passwordsDoNotMatch": "パスワードが一致しません", + "createAccount": "管理者アカウントを作る", + "creatingAccount": "設定中...", + "setupFailed": "セットアップに失敗しました。もう一度試してください。", + "passwordHelperRelaxed": "パスワードを入力してください（強度が表示されます）" + }, + "userMenu": "ユーザーメニュー", + "admin": "管理", + "logout": "ログアウト", + "adminOnlyFeature": "この機能は管理者のみ使用できます。", + "profile": { + "menuItem": "プロフィール", + "title": "プロフィール", + "email": "Eメール", + "emailReadOnly": "Eメールアドレスは変更できません", + "displayName": "表示名", + "displayNamePlaceholder": "あなたの名前", + "changePassword": "パスワードの変更", + "currentPassword": "現在のパスワード", + "currentPasswordPlaceholder": "現在のパスワード", + "newPassword": "新しいパスワード", + "newPasswordPlaceholder": "新しいパスワード", + "confirmPassword": "新しいパスワードの確認", + "confirmPasswordPlaceholder": "新しいパスワードの確認", + "passwordsDoNotMatch": "パスワードが一致しません", + "saveSuccess": "プロフィールのアップデートに成功しました", + "saveFailed": "プロフィールの保存に失敗しました。もう一度試してください。" + }, + "userManagement": { + "menuItem": "ユーザー管理", + "title": "ユーザー管理", + "email": "Eメール", + "emailPlaceholder": "hoge@example.com", + "displayName": "表示名", + "displayNamePlaceholder": "表示名", + "password": "パスワード", + "passwordPlaceholder": "パスワード", + "newPassword": "新しいパスワード", + "newPasswordPlaceholder": "現在のパスワードを維持するには空白にしておいてください", + "role": "ロール", + "status": "ステータス", + "actions": "アクション", + "isAdmin": "管理者", + "user": "ユーザー", + "you": "あなた", + "createUser": "ユーザーの作成", + "editUser": "ユーザーの編集", + "deleteUser": "ユーザーの削除", + "deleteConfirm": "本当に \"{{name}}\" を削除しますか？このアクションは取り消せません。", + "generatePassword": "強力なパスワードを生成", + "showPassword": "パスワードの表示", + "hidePassword": "パスワードを隠す", + "activate": "有効化", + "deactivate": "非有効化", + "saveFailed": "ユーザーの保存に失敗しました。もう一度実行してください。", + "deleteFailed": "ユーザーの削除に失敗しました。もう一度実行してください。", + "loadFailed": "ユーザーのロードに失敗しました。", + "back": "戻る", + "cannotDeleteSelf": "あなた自身のアカウントを削除することはできません", + "cannotDeactivateSelf": "あなた自身のアカウントを非有効化することはできません" + }, + "passwordStrength": { + "weak": "弱いパスワード", + "moderate": "適切なパスワード", + "strong": "強力なパスワード" + } } } diff --git a/invokeai/frontend/web/scripts/typegen.js b/invokeai/frontend/web/scripts/typegen.js index 7b3a7472857..87c00a28833 100644 --- a/invokeai/frontend/web/scripts/typegen.js +++ b/invokeai/frontend/web/scripts/typegen.js @@ -34,7 +34,37 @@ async function generateTypes(schema) { }, defaultNonNullable: false, }); - fs.writeFileSync(OUTPUT_FILE, astToString(types)); + let output = astToString(types); + + // Post-process: openapi-typescript sometimes computes enum types from `const` + // usage in discriminated unions rather than from the enum definition itself, + // dropping values that only appear in some union members. Patch the generated + // output to match the OpenAPI schema's actual enum definitions. + // + // The `schema` parameter is a parsed JSON object when piped from stdin, or + // a URL/Buffer when passed as an argument. We only patch in the JSON case. + if (schema && typeof schema === 'object' && !Buffer.isBuffer(schema)) { + const schemas = schema.components?.schemas; + if (schemas) { + // Collect all string enum types and their expected values from the OpenAPI schema + for (const [typeName, typeDef] of Object.entries(schemas)) { + if (typeDef && typeDef.type === 'string' && Array.isArray(typeDef.enum)) { + const expectedUnion = typeDef.enum.map((v) => `"${v}"`).join(' | '); + // Match the type definition line. These appear as: + // `TypeName: "val1" | "val2" | ...;` + // Use word boundary to avoid matching types that contain this + // type name as a substring (e.g. ModelType vs BaseModelType). + const regex = new RegExp(`(\\b${typeName}: )"[^;]+(;)`); + const match = output.match(regex); + if (match) { + output = output.replace(regex, `$1${expectedUnion}$2`); + } + } + } + } + } + + fs.writeFileSync(OUTPUT_FILE, output); process.stdout.write(`\nOK!\r\n`); } diff --git a/invokeai/frontend/web/src/app/components/GlobalModalIsolator.tsx b/invokeai/frontend/web/src/app/components/GlobalModalIsolator.tsx index 5c1446662ef..ef0747707ff 100644 --- a/invokeai/frontend/web/src/app/components/GlobalModalIsolator.tsx +++ b/invokeai/frontend/web/src/app/components/GlobalModalIsolator.tsx @@ -1,6 +1,7 @@ import { GlobalImageHotkeys } from 'app/components/GlobalImageHotkeys'; import ChangeBoardModal from 'features/changeBoardModal/components/ChangeBoardModal'; import { CanvasPasteModal } from 'features/controlLayers/components/CanvasPasteModal'; +import { CanvasWorkflowIntegrationModal } from 'features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationModal'; import { CanvasManagerProviderGate } from 'features/controlLayers/contexts/CanvasManagerProviderGate'; import { CropImageModal } from 'features/cropper/components/CropImageModal'; import { DeleteImageModal } from 'features/deleteImageModal/components/DeleteImageModal'; @@ -51,6 +52,7 @@ export const GlobalModalIsolator = memo(() => { + diff --git a/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx b/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx index 437cecd4925..62b7114288d 100644 --- a/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx +++ b/invokeai/frontend/web/src/app/components/ThemeLocaleProvider.tsx @@ -2,6 +2,7 @@ import '@fontsource-variable/inter'; import 'overlayscrollbars/overlayscrollbars.css'; import '@xyflow/react/dist/base.css'; import 'common/components/OverlayScrollbars/overlayscrollbars.css'; +import 'app/components/touchDevice.css'; import { ChakraProvider, DarkMode, extendTheme, theme as baseTheme, TOAST_OPTIONS } from '@invoke-ai/ui-library'; import { useStore } from '@nanostores/react'; diff --git a/invokeai/frontend/web/src/app/components/touchDevice.css b/invokeai/frontend/web/src/app/components/touchDevice.css new file mode 100644 index 00000000000..4753e2a9a8c --- /dev/null +++ b/invokeai/frontend/web/src/app/components/touchDevice.css @@ -0,0 +1,7 @@ +/* Hide tooltips on touch devices where hover gets "stuck" */ +@media (hover: none) { + [role='tooltip'] { + visibility: hidden !important; + opacity: 0 !important; + } +} diff --git a/invokeai/frontend/web/src/app/logging/logger.ts b/invokeai/frontend/web/src/app/logging/logger.ts index 6c843068df3..d20ef77090f 100644 --- a/invokeai/frontend/web/src/app/logging/logger.ts +++ b/invokeai/frontend/web/src/app/logging/logger.ts @@ -16,6 +16,7 @@ const $logger = atom(Roarr.child(BASE_CONTEXT)); export const zLogNamespace = z.enum([ 'canvas', + 'canvas-workflow-integration', 'config', 'dnd', 'events', diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/bulkDownload.tsx b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/bulkDownload.tsx index e0e72d12ffd..fa4c29b8f42 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/bulkDownload.tsx +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/bulkDownload.tsx @@ -12,10 +12,14 @@ export const addBulkDownloadListeners = (startAppListening: AppStartListening) = effect: (action) => { log.debug(action.payload, 'Bulk download requested'); - // If we have an item name, we are processing the bulk download locally and should use it as the toast id to - // prevent multiple toasts for the same item. + // Use a "preparing:" prefix so this toast cannot collide with the + // "ready to download" toast that arrives via the bulk_download_complete + // socket event. The background task can complete in under 20ms, so the + // socket event may arrive *before* this Redux middleware runs — without + // distinct IDs the "preparing" toast would overwrite the "ready" toast. + const itemName = action.payload.bulk_download_item_name; toast({ - id: action.payload.bulk_download_item_name ?? undefined, + id: itemName ? `preparing:${itemName}` : undefined, title: t('gallery.bulkDownloadRequested'), status: 'success', // Show the response message if it exists, otherwise show the default message diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts index 20057472ca8..63b0d6469d9 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/modelSelected.ts @@ -4,9 +4,13 @@ import { bboxSyncedToOptimalDimension, rgRefImageModelChanged } from 'features/c import { buildSelectIsStaging, selectCanvasSessionId } from 'features/controlLayers/store/canvasStagingAreaSlice'; import { loraIsEnabledChanged } from 'features/controlLayers/store/lorasSlice'; import { + animaQwen3EncoderModelSelected, + animaT5EncoderModelSelected, + animaVaeModelSelected, kleinQwen3EncoderModelSelected, kleinVaeModelSelected, modelChanged, + qwenImageComponentSourceSelected, setZImageScheduler, syncedToOptimalDimension, vaeSelected, @@ -24,12 +28,17 @@ import { selectBboxModelBase, selectCanvasSlice, } from 'features/controlLayers/store/selectors'; -import { getEntityIdentifier, isFlux2ReferenceImageConfig } from 'features/controlLayers/store/types'; +import { + getEntityIdentifier, + isFlux2ReferenceImageConfig, + isQwenImageReferenceImageConfig, +} from 'features/controlLayers/store/types'; import { initialFlux2ReferenceImage, initialFluxKontextReferenceImage, initialFLUXRedux, initialIPAdapter, + initialQwenImageReferenceImage, } from 'features/controlLayers/store/util'; import { SUPPORTS_REF_IMAGES_BASE_MODELS } from 'features/modelManagerV2/models'; import { zModelIdentifierField } from 'features/nodes/types/common'; @@ -39,10 +48,14 @@ import { toast } from 'features/toast/toast'; import { t } from 'i18next'; import { modelConfigsAdapterSelectors, selectModelConfigsQuery } from 'services/api/endpoints/models'; import { + selectAnimaQwen3EncoderModels, + selectAnimaVAEModels, selectFluxVAEModels, selectGlobalRefImageModels, selectQwen3EncoderModels, + selectQwenImageDiffusersModels, selectRegionalRefImageModels, + selectT5EncoderModels, selectZImageDiffusersModels, } from 'services/api/hooks/modelsByType'; import type { FLUXKontextModelConfig, FLUXReduxModelConfig, IPAdapterModelConfig } from 'services/api/types'; @@ -155,6 +168,68 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = } } + // handle incompatible Anima models - clear if switching away from anima + const { animaVaeModel, animaQwen3EncoderModel, animaT5EncoderModel } = state.params; + if (newBase !== 'anima') { + if (animaVaeModel) { + dispatch(animaVaeModelSelected(null)); + modelsUpdatedDisabledOrCleared += 1; + } + if (animaQwen3EncoderModel) { + dispatch(animaQwen3EncoderModelSelected(null)); + modelsUpdatedDisabledOrCleared += 1; + } + if (animaT5EncoderModel) { + dispatch(animaT5EncoderModelSelected(null)); + modelsUpdatedDisabledOrCleared += 1; + } + } else { + // Switching to Anima - set defaults if no valid configuration exists + const hasValidConfig = animaVaeModel && animaQwen3EncoderModel && animaT5EncoderModel; + + if (!hasValidConfig) { + const availableQwen3Encoders = selectAnimaQwen3EncoderModels(state); + const availableAnimaVAEs = selectAnimaVAEModels(state); + const availableT5Encoders = selectT5EncoderModels(state); + + if (availableQwen3Encoders.length > 0 && availableAnimaVAEs.length > 0) { + const qwen3Encoder = availableQwen3Encoders[0]; + const fluxVAE = availableAnimaVAEs[0]; + + if (qwen3Encoder && !animaQwen3EncoderModel) { + dispatch( + animaQwen3EncoderModelSelected({ + key: qwen3Encoder.key, + name: qwen3Encoder.name, + base: qwen3Encoder.base, + }) + ); + } + if (fluxVAE && !animaVaeModel) { + dispatch( + animaVaeModelSelected({ + key: fluxVAE.key, + hash: fluxVAE.hash, + name: fluxVAE.name, + base: fluxVAE.base, + type: fluxVAE.type, + }) + ); + } + const t5Encoder = availableT5Encoders[0]; + if (t5Encoder && !animaT5EncoderModel) { + dispatch( + animaT5EncoderModelSelected({ + key: t5Encoder.key, + name: t5Encoder.name, + base: t5Encoder.base, + }) + ); + } + } + } + } + // handle incompatible FLUX.2 Klein models - clear if switching away from flux2 const { kleinVaeModel, kleinQwen3EncoderModel } = state.params; if (newBase !== 'flux2') { @@ -168,6 +243,44 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = } } + // handle incompatible Qwen Image Edit component source - clear if switching away + const { qwenImageComponentSource } = state.params; + if (newBase !== 'qwen-image') { + if (qwenImageComponentSource) { + dispatch(qwenImageComponentSourceSelected(null)); + modelsUpdatedDisabledOrCleared += 1; + } + } else { + // Switching to Qwen Image - auto-default component source to a matching diffusers model + if (!qwenImageComponentSource) { + const availableQwenImageDiffusers = selectQwenImageDiffusersModels(state); + + // Look up the new model's variant to match generate vs edit + const modelConfigsResult = selectModelConfigsQuery(state); + let selectedVariant: string | null = null; + if (modelConfigsResult.data) { + const newModelConfig = modelConfigsAdapterSelectors.selectById(modelConfigsResult.data, newModel.key); + if (newModelConfig && 'variant' in newModelConfig && typeof newModelConfig.variant === 'string') { + selectedVariant = newModelConfig.variant; + } + } + + // Find a diffusers model matching the variant; if no variant on denoiser, prefer "generate" then "edit" + const variantToMatch = selectedVariant ?? 'generate'; + const matchingModel = availableQwenImageDiffusers.find( + (m) => 'variant' in m && m.variant === variantToMatch + ); + const fallbackModel = availableQwenImageDiffusers.find( + (m) => 'variant' in m && m.variant !== variantToMatch + ); + const diffusersModel = matchingModel ?? fallbackModel ?? availableQwenImageDiffusers[0]; + + if (diffusersModel) { + dispatch(qwenImageComponentSourceSelected(zModelIdentifierField.parse(diffusersModel))); + } + } + } + if (SUPPORTS_REF_IMAGES_BASE_MODELS.includes(newModel.base)) { // Handle incompatible reference image models - switch to first compatible model, with some smart logic // to choose the best available model based on the new main model. @@ -210,6 +323,20 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = continue; } + if (newBase === 'qwen-image') { + // Switching TO Qwen Image Edit - convert any non-qwen configs to qwen_image_reference_image + if (!isQwenImageReferenceImageConfig(entity.config)) { + dispatch( + refImageConfigChanged({ + id: entity.id, + config: { ...initialQwenImageReferenceImage }, + }) + ); + modelsUpdatedDisabledOrCleared += 1; + } + continue; + } + if (isFlux2ReferenceImageConfig(entity.config)) { // Switching AWAY from FLUX.2 - convert flux2_reference_image to the appropriate config type let newConfig; @@ -234,6 +361,30 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = continue; } + if (isQwenImageReferenceImageConfig(entity.config)) { + // Switching AWAY from Qwen Image Edit - convert to the appropriate config type + let newConfig; + if (newGlobalRefImageModel) { + const parsedModel = zModelIdentifierField.parse(newGlobalRefImageModel); + if (newModel.base === 'flux' && newModel.name.toLowerCase().includes('kontext')) { + newConfig = { ...initialFluxKontextReferenceImage, model: parsedModel }; + } else if (newGlobalRefImageModel.type === 'flux_redux') { + newConfig = { ...initialFLUXRedux, model: parsedModel }; + } else { + newConfig = { ...initialIPAdapter, model: parsedModel }; + if (parsedModel.base === 'flux') { + newConfig.clipVisionModel = 'ViT-L'; + } + } + } else { + // No compatible model found - fall back to an empty IP adapter config + newConfig = { ...initialIPAdapter }; + } + dispatch(refImageConfigChanged({ id: entity.id, config: newConfig })); + modelsUpdatedDisabledOrCleared += 1; + continue; + } + // Standard handling for non-flux2 configs const shouldUpdateModel = (entity.config.model && entity.config.model.base !== newBase) || @@ -321,6 +472,32 @@ export const addModelSelectedListener = (startAppListening: AppStartListening) = } } + // Handle Qwen Image model changes within the same base (variant may change between generate/edit) + // Auto-update the component source diffusers model to match the new variant + if ( + newBase === 'qwen-image' && + state.params.model?.base === 'qwen-image' && + newModel.key !== state.params.model?.key + ) { + const modelConfigsResult = selectModelConfigsQuery(state); + if (modelConfigsResult.data) { + const newModelConfig = modelConfigsAdapterSelectors.selectById(modelConfigsResult.data, newModel.key); + const newVariant = + newModelConfig && 'variant' in newModelConfig && typeof newModelConfig.variant === 'string' + ? newModelConfig.variant + : 'generate'; + + const availableQwenImageDiffusers = selectQwenImageDiffusersModels(state); + const matchingModel = availableQwenImageDiffusers.find((m) => 'variant' in m && m.variant === newVariant); + const fallbackModel = availableQwenImageDiffusers.find((m) => 'variant' in m && m.variant !== newVariant); + const diffusersModel = matchingModel ?? fallbackModel ?? availableQwenImageDiffusers[0]; + + if (diffusersModel) { + dispatch(qwenImageComponentSourceSelected(zModelIdentifierField.parse(diffusersModel))); + } + } + } + // Handle Z-Image scheduler when switching to Z-Image Base (zbase) model // LCM is not supported for undistilled models, so reset to euler if (newBase === 'z-image' && state.params.zImageScheduler === 'lcm') { diff --git a/invokeai/frontend/web/src/app/store/store.ts b/invokeai/frontend/web/src/app/store/store.ts index 8f077baaea5..f24d2d0105c 100644 --- a/invokeai/frontend/web/src/app/store/store.ts +++ b/invokeai/frontend/web/src/app/store/store.ts @@ -25,6 +25,7 @@ import { canvasSettingsSliceConfig } from 'features/controlLayers/store/canvasSe import { canvasSliceConfig } from 'features/controlLayers/store/canvasSlice'; import { canvasSessionSliceConfig } from 'features/controlLayers/store/canvasStagingAreaSlice'; import { canvasTextSliceConfig } from 'features/controlLayers/store/canvasTextSlice'; +import { canvasWorkflowIntegrationSliceConfig } from 'features/controlLayers/store/canvasWorkflowIntegrationSlice'; import { lorasSliceConfig } from 'features/controlLayers/store/lorasSlice'; import { paramsSliceConfig } from 'features/controlLayers/store/paramsSlice'; import { refImagesSliceConfig } from 'features/controlLayers/store/refImagesSlice'; @@ -67,6 +68,7 @@ const SLICE_CONFIGS = { [canvasSettingsSliceConfig.slice.reducerPath]: canvasSettingsSliceConfig, [canvasTextSliceConfig.slice.reducerPath]: canvasTextSliceConfig, [canvasSliceConfig.slice.reducerPath]: canvasSliceConfig, + [canvasWorkflowIntegrationSliceConfig.slice.reducerPath]: canvasWorkflowIntegrationSliceConfig, [changeBoardModalSliceConfig.slice.reducerPath]: changeBoardModalSliceConfig, [dynamicPromptsSliceConfig.slice.reducerPath]: dynamicPromptsSliceConfig, [gallerySliceConfig.slice.reducerPath]: gallerySliceConfig, @@ -98,6 +100,7 @@ const ALL_REDUCERS = { canvasSliceConfig.slice.reducer, canvasSliceConfig.undoableConfig?.reduxUndoOptions ), + [canvasWorkflowIntegrationSliceConfig.slice.reducerPath]: canvasWorkflowIntegrationSliceConfig.slice.reducer, [changeBoardModalSliceConfig.slice.reducerPath]: changeBoardModalSliceConfig.slice.reducer, [dynamicPromptsSliceConfig.slice.reducerPath]: dynamicPromptsSliceConfig.slice.reducer, [gallerySliceConfig.slice.reducerPath]: gallerySliceConfig.slice.reducer, diff --git a/invokeai/frontend/web/src/common/components/ColorPicker/RgbaColorPicker.tsx b/invokeai/frontend/web/src/common/components/ColorPicker/RgbaColorPicker.tsx index 6bcf546fd34..38fcb2696ef 100644 --- a/invokeai/frontend/web/src/common/components/ColorPicker/RgbaColorPicker.tsx +++ b/invokeai/frontend/web/src/common/components/ColorPicker/RgbaColorPicker.tsx @@ -85,7 +85,7 @@ const RgbaColorPicker = (props: Props) => { h={10} whiteSpace="nowrap" onClick={onToggleMode} - aria-label="Toggle RGB/HEX" + aria-label={t('common.toggleRgbHex')} > RGB @@ -144,12 +144,12 @@ const RgbaColorPicker = (props: Props) => { h={10} whiteSpace="nowrap" onClick={onToggleMode} - aria-label="Toggle RGB/HEX" + aria-label={t('common.toggleRgbHex')} > HEX - {t('common.hex', { defaultValue: 'Hex' })} + {t('common.hex')} diff --git a/invokeai/frontend/web/src/common/components/Picker/Picker.tsx b/invokeai/frontend/web/src/common/components/Picker/Picker.tsx index ffd0b30242a..b70e44dd649 100644 --- a/invokeai/frontend/web/src/common/components/Picker/Picker.tsx +++ b/invokeai/frontend/web/src/common/components/Picker/Picker.tsx @@ -867,7 +867,7 @@ const GroupToggleButtons = typedMemo(() => { } return ( - + {groups.map((group) => ( ))} @@ -927,6 +927,7 @@ const GroupToggleButton = typedMemo(({ group }: { group: Group size="xs" variant="solid" userSelect="none" + flexShrink={0} bg={bg} color={color} borderColor={groupColor} diff --git a/invokeai/frontend/web/src/common/hooks/focus.test.ts b/invokeai/frontend/web/src/common/hooks/focus.test.ts new file mode 100644 index 00000000000..c106fe1cec4 --- /dev/null +++ b/invokeai/frontend/web/src/common/hooks/focus.test.ts @@ -0,0 +1,13 @@ +import { describe, expect, it } from 'vitest'; + +import { getFocusedRegion, setFocusedRegion } from './focus'; + +describe('focus regions', () => { + it('supports the workflows region', () => { + setFocusedRegion('workflows'); + expect(getFocusedRegion()).toBe('workflows'); + + setFocusedRegion(null); + expect(getFocusedRegion()).toBe(null); + }); +}); diff --git a/invokeai/frontend/web/src/common/hooks/useImageUploadButton.tsx b/invokeai/frontend/web/src/common/hooks/useImageUploadButton.tsx index 445f58c9a66..fc173de979f 100644 --- a/invokeai/frontend/web/src/common/hooks/useImageUploadButton.tsx +++ b/invokeai/frontend/web/src/common/hooks/useImageUploadButton.tsx @@ -187,11 +187,12 @@ export const UploadImageIconButton = memo( onUpload?: (imageDTO: ImageDTO) => void; isError?: boolean; } & SetOptional) => { + const { t } = useTranslation(); const uploadApi = useImageUploadButton({ isDisabled, allowMultiple: false, onUpload }); return ( <> { + const { t } = useTranslation(); const { children, isDisabled = false, onUpload, isError = false, ...rest } = props; const uploadApi = useImageUploadButton({ isDisabled, allowMultiple: false, onUpload }); return ( <> + + + + + + + ); +}); + +CanvasWorkflowIntegrationModal.displayName = 'CanvasWorkflowIntegrationModal'; diff --git a/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationParameterPanel.tsx b/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationParameterPanel.tsx new file mode 100644 index 00000000000..f59a6c45edb --- /dev/null +++ b/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationParameterPanel.tsx @@ -0,0 +1,13 @@ +import { Box } from '@invoke-ai/ui-library'; +import { WorkflowFormPreview } from 'features/controlLayers/components/CanvasWorkflowIntegration/WorkflowFormPreview'; +import { memo } from 'react'; + +export const CanvasWorkflowIntegrationParameterPanel = memo(() => { + return ( + + + + ); +}); + +CanvasWorkflowIntegrationParameterPanel.displayName = 'CanvasWorkflowIntegrationParameterPanel'; diff --git a/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationWorkflowSelector.tsx b/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationWorkflowSelector.tsx new file mode 100644 index 00000000000..30bc60605c6 --- /dev/null +++ b/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/CanvasWorkflowIntegrationWorkflowSelector.tsx @@ -0,0 +1,92 @@ +import { Flex, FormControl, FormLabel, Select, Spinner, Text } from '@invoke-ai/ui-library'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { + canvasWorkflowIntegrationWorkflowSelected, + selectCanvasWorkflowIntegrationSelectedWorkflowId, +} from 'features/controlLayers/store/canvasWorkflowIntegrationSlice'; +import type { ChangeEvent } from 'react'; +import { memo, useCallback, useMemo } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useListWorkflowsInfiniteInfiniteQuery } from 'services/api/endpoints/workflows'; + +import { useFilteredWorkflows } from './useFilteredWorkflows'; + +export const CanvasWorkflowIntegrationWorkflowSelector = memo(() => { + const { t } = useTranslation(); + const dispatch = useAppDispatch(); + + const selectedWorkflowId = useAppSelector(selectCanvasWorkflowIntegrationSelectedWorkflowId); + const { data: workflowsData, isLoading } = useListWorkflowsInfiniteInfiniteQuery( + { + per_page: 100, // Get a reasonable number of workflows + page: 0, + }, + { + selectFromResult: ({ data, isLoading }) => ({ + data, + isLoading, + }), + } + ); + + const workflows = useMemo(() => { + if (!workflowsData) { + return []; + } + // Flatten all pages into a single list + return workflowsData.pages.flatMap((page) => page.items); + }, [workflowsData]); + + // Filter workflows to only show those with ImageFields + const { filteredWorkflows, isFiltering } = useFilteredWorkflows(workflows); + + const onChange = useCallback( + (e: ChangeEvent) => { + const workflowId = e.target.value || null; + dispatch(canvasWorkflowIntegrationWorkflowSelected({ workflowId })); + }, + [dispatch] + ); + + if (isLoading || isFiltering) { + return ( + + + + {isFiltering + ? t('controlLayers.workflowIntegration.filteringWorkflows') + : t('controlLayers.workflowIntegration.loadingWorkflows')} + + + ); + } + + if (filteredWorkflows.length === 0) { + return ( + + {workflows.length === 0 + ? t('controlLayers.workflowIntegration.noWorkflowsFound') + : t('controlLayers.workflowIntegration.noWorkflowsWithImageField')} + + ); + } + + return ( + + {t('controlLayers.workflowIntegration.selectWorkflow')} + + + ); +}); + +CanvasWorkflowIntegrationWorkflowSelector.displayName = 'CanvasWorkflowIntegrationWorkflowSelector'; diff --git a/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/WorkflowFieldRenderer.tsx b/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/WorkflowFieldRenderer.tsx new file mode 100644 index 00000000000..2d91be13bfa --- /dev/null +++ b/invokeai/frontend/web/src/features/controlLayers/components/CanvasWorkflowIntegration/WorkflowFieldRenderer.tsx @@ -0,0 +1,548 @@ +import type { ComboboxOnChange, ComboboxOption } from '@invoke-ai/ui-library'; +import { + Combobox, + Flex, + FormControl, + FormLabel, + IconButton, + Input, + Radio, + Select, + Switch, + Text, + Textarea, +} from '@invoke-ai/ui-library'; +import { useStore } from '@nanostores/react'; +import { skipToken } from '@reduxjs/toolkit/query'; +import { logger } from 'app/logging/logger'; +import { EMPTY_ARRAY } from 'app/store/constants'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { UploadImageIconButton } from 'common/hooks/useImageUploadButton'; +import { + canvasWorkflowIntegrationFieldValueChanged, + canvasWorkflowIntegrationImageFieldSelected, + selectCanvasWorkflowIntegrationFieldValues, + selectCanvasWorkflowIntegrationSelectedImageFieldKey, + selectCanvasWorkflowIntegrationSelectedWorkflowId, +} from 'features/controlLayers/store/canvasWorkflowIntegrationSlice'; +import { DndImage } from 'features/dnd/DndImage'; +import { ModelFieldCombobox } from 'features/nodes/components/flow/nodes/Invocation/fields/inputs/ModelFieldCombobox'; +import { $templates } from 'features/nodes/store/nodesSlice'; +import type { NodeFieldElement } from 'features/nodes/types/workflow'; +import { SCHEDULER_OPTIONS } from 'features/parameters/types/constants'; +import { isParameterScheduler } from 'features/parameters/types/parameterSchemas'; +import type { ChangeEvent } from 'react'; +import { memo, useCallback, useMemo } from 'react'; +import { useTranslation } from 'react-i18next'; +import { PiTrashSimpleBold } from 'react-icons/pi'; +import { useListAllBoardsQuery } from 'services/api/endpoints/boards'; +import { useGetImageDTOQuery } from 'services/api/endpoints/images'; +import { modelConfigsAdapterSelectors, useGetModelConfigsQuery } from 'services/api/endpoints/models'; +import { useGetWorkflowQuery } from 'services/api/endpoints/workflows'; +import type { AnyModelConfig, ImageDTO } from 'services/api/types'; + +const log = logger('canvas-workflow-integration'); + +interface WorkflowFieldRendererProps { + el: NodeFieldElement; +} + +export const WorkflowFieldRenderer = memo(({ el }: WorkflowFieldRendererProps) => { + const dispatch = useAppDispatch(); + const { t } = useTranslation(); + const selectedWorkflowId = useAppSelector(selectCanvasWorkflowIntegrationSelectedWorkflowId); + const fieldValues = useAppSelector(selectCanvasWorkflowIntegrationFieldValues); + const selectedImageFieldKey = useAppSelector(selectCanvasWorkflowIntegrationSelectedImageFieldKey); + const templates = useStore($templates); + + const { data: workflow } = useGetWorkflowQuery(selectedWorkflowId!, { + skip: !selectedWorkflowId, + }); + + // Load boards and models for BoardField and ModelIdentifierField + const { data: boardsData } = useListAllBoardsQuery({ include_archived: true }); + const { data: modelsData, isLoading: isLoadingModels } = useGetModelConfigsQuery(); + + const { fieldIdentifier } = el.data; + const fieldKey = `${fieldIdentifier.nodeId}.${fieldIdentifier.fieldName}`; + + log.debug({ fieldIdentifier, fieldKey }, 'Rendering workflow field'); + + // Get the node, field instance, and field template + const { field, fieldTemplate } = useMemo(() => { + if (!workflow?.workflow.nodes) { + log.warn('No workflow nodes found'); + return { field: null, fieldTemplate: null }; + } + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const foundNode = workflow.workflow.nodes.find((n: any) => n.data.id === fieldIdentifier.nodeId); + if (!foundNode) { + log.warn({ nodeId: fieldIdentifier.nodeId }, 'Node not found'); + return { field: null, fieldTemplate: null }; + } + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const foundField = (foundNode.data as any).inputs[fieldIdentifier.fieldName]; + if (!foundField) { + log.warn({ nodeId: fieldIdentifier.nodeId, fieldName: fieldIdentifier.fieldName }, 'Field not found in node'); + return { field: null, fieldTemplate: null }; + } + + // Get the field template from the invocation templates + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const nodeType = (foundNode.data as any).type; + const template = templates[nodeType]; + if (!template) { + log.warn({ nodeType }, 'No template found for node type'); + return { field: foundField, fieldTemplate: null }; + } + + const foundFieldTemplate = template.inputs[fieldIdentifier.fieldName]; + if (!foundFieldTemplate) { + log.warn({ nodeType, fieldName: fieldIdentifier.fieldName }, 'Field template not found'); + return { field: foundField, fieldTemplate: null }; + } + + return { field: foundField, fieldTemplate: foundFieldTemplate }; + }, [workflow, fieldIdentifier, templates]); + + // Get the current value from Redux or fallback to field default + const currentValue = useMemo(() => { + if (fieldValues && fieldKey in fieldValues) { + return fieldValues[fieldKey]; + } + + return field?.value ?? fieldTemplate?.default ?? ''; + }, [fieldValues, fieldKey, field, fieldTemplate]); + + // Get field type from the template + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const fieldType = fieldTemplate ? (fieldTemplate as any).type?.name : null; + + const handleChange = useCallback( + (value: unknown) => { + dispatch(canvasWorkflowIntegrationFieldValueChanged({ fieldName: fieldKey, value })); + }, + [dispatch, fieldKey] + ); + + const handleStringChange = useCallback( + (e: ChangeEvent) => { + handleChange(e.target.value); + }, + [handleChange] + ); + + const handleNumberChange = useCallback( + (e: ChangeEvent) => { + const val = fieldType === 'IntegerField' ? parseInt(e.target.value, 10) : parseFloat(e.target.value); + handleChange(isNaN(val) ? 0 : val); + }, + [handleChange, fieldType] + ); + + const handleBooleanChange = useCallback( + (e: ChangeEvent) => { + handleChange(e.target.checked); + }, + [handleChange] + ); + + const handleSelectChange = useCallback( + (e: ChangeEvent) => { + handleChange(e.target.value); + }, + [handleChange] + ); + + // SchedulerField handlers + const handleSchedulerChange = useCallback( + (v) => { + if (!isParameterScheduler(v?.value)) { + return; + } + handleChange(v.value); + }, + [handleChange] + ); + + const schedulerValue = useMemo(() => SCHEDULER_OPTIONS.find((o) => o.value === currentValue), [currentValue]); + + // BoardField handlers + const handleBoardChange = useCallback( + (v) => { + if (!v) { + return; + } + const value = v.value === 'auto' || v.value === 'none' ? v.value : { board_id: v.value }; + handleChange(value); + }, + [handleChange] + ); + + const boardOptions = useMemo(() => { + const _options: ComboboxOption[] = [ + { label: t('common.auto'), value: 'auto' }, + { label: `${t('common.none')} (${t('boards.uncategorized')})`, value: 'none' }, + ]; + if (boardsData) { + for (const board of boardsData) { + _options.push({ + label: board.board_name, + value: board.board_id, + }); + } + } + return _options; + }, [boardsData, t]); + + const boardValue = useMemo(() => { + const _value = currentValue; + const autoOption = boardOptions[0]; + const noneOption = boardOptions[1]; + if (!_value || _value === 'auto') { + return autoOption; + } + if (_value === 'none') { + return noneOption; + } + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const boardId = typeof _value === 'object' ? (_value as any).board_id : _value; + const boardOption = boardOptions.find((o) => o.value === boardId); + return boardOption ?? autoOption; + }, [currentValue, boardOptions]); + + const noOptionsMessage = useCallback(() => t('boards.noMatching'), [t]); + + // ModelIdentifierField handlers + const handleModelChange = useCallback( + (value: AnyModelConfig | null) => { + if (!value) { + return; + } + handleChange(value); + }, + [handleChange] + ); + + const modelConfigs = useMemo(() => { + if (!modelsData) { + return EMPTY_ARRAY; + } + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const ui_model_base = fieldTemplate ? (fieldTemplate as any)?.ui_model_base : null; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const ui_model_type = fieldTemplate ? (fieldTemplate as any)?.ui_model_type : null; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const ui_model_variant = fieldTemplate ? (fieldTemplate as any)?.ui_model_variant : null; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const ui_model_format = fieldTemplate ? (fieldTemplate as any)?.ui_model_format : null; + + if (!ui_model_base && !ui_model_type) { + return modelConfigsAdapterSelectors.selectAll(modelsData); + } + + return modelConfigsAdapterSelectors.selectAll(modelsData).filter((config) => { + if (ui_model_base && !ui_model_base.includes(config.base)) { + return false; + } + if (ui_model_type && !ui_model_type.includes(config.type)) { + return false; + } + if (ui_model_variant && 'variant' in config && config.variant && !ui_model_variant.includes(config.variant)) { + return false; + } + if (ui_model_format && !ui_model_format.includes(config.format)) { + return false; + } + return true; + }); + }, [modelsData, fieldTemplate]); + + // ImageField handler + const handleImageFieldSelect = useCallback(() => { + dispatch(canvasWorkflowIntegrationImageFieldSelected({ fieldKey })); + }, [dispatch, fieldKey]); + + if (!field || !fieldTemplate) { + log.warn({ fieldIdentifier }, 'Field or template is null - not rendering'); + return null; + } + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const label = (field as any)?.label || (fieldTemplate as any)?.title || fieldIdentifier.fieldName; + + // Log the entire field structure to understand its shape + log.debug( + { fieldType, label, currentValue, fieldStructure: field, fieldTemplateStructure: fieldTemplate }, + 'Field info' + ); + + // ImageField - allow user to select which one receives the canvas image + if (fieldType === 'ImageField') { + return ( + + ); + } + + // Render different input types based on field type + if (fieldType === 'StringField') { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const isTextarea = (fieldTemplate as any)?.ui_component === 'textarea'; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const isRequired = (fieldTemplate as any)?.required ?? false; + + if (isTextarea) { + return ( + + {label} +