From 6644a4e0bd1580bc9d024ad99984ddba7f351eb1 Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Mon, 15 Dec 2025 15:33:59 -0500 Subject: [PATCH 1/7] Personalive implementation in scope. --- .../src/components/InputAndControlsPanel.tsx | 193 ++- frontend/src/data/pipelines.ts | 18 + frontend/src/lib/api.ts | 38 + frontend/src/pages/StreamPage.tsx | 94 +- frontend/src/types/index.ts | 4 +- pyproject.toml | 6 + src/scope/core/pipelines/__init__.py | 10 + .../core/pipelines/personalive/__init__.py | 5 + .../pipelines/personalive/blocks/__init__.py | 1 + .../personalive/components/__init__.py | 7 + .../personalive/components/face_detector.py | 199 +++ .../core/pipelines/personalive/docs/README.md | 124 ++ .../core/pipelines/personalive/model.yaml | 70 + .../modules/FAN_feature_extractor.py | 334 ++++ .../modules/FAN_temporal_feature_extractor.py | 231 +++ .../pipelines/personalive/modules/__init__.py | 30 + .../personalive/modules/attention.py | 463 ++++++ .../pipelines/personalive/modules/camera.py | 73 + .../personalive/modules/convnextv2.py | 215 +++ .../personalive/modules/liveportrait_util.py | 492 ++++++ .../personalive/modules/motion_encoder.py | 42 + .../personalive/modules/motion_extractor.py | 212 +++ .../personalive/modules/motion_module.py | 472 ++++++ .../modules/mutual_self_attention.py | 463 ++++++ .../personalive/modules/pose_guider.py | 57 + .../pipelines/personalive/modules/resnet.py | 255 +++ .../personalive/modules/scheduler_ddim.py | 550 ++++++ .../personalive/modules/transformer_2d.py | 395 +++++ .../personalive/modules/transformer_3d.py | 182 ++ .../personalive/modules/unet_2d_blocks.py | 1073 ++++++++++++ .../personalive/modules/unet_2d_condition.py | 1311 +++++++++++++++ .../pipelines/personalive/modules/unet_3d.py | 717 ++++++++ .../personalive/modules/unet_3d_blocks.py | 1119 +++++++++++++ .../pipelines/personalive/modules/util.py | 313 ++++ .../core/pipelines/personalive/pipeline.py | 656 ++++++++ src/scope/core/pipelines/registry.py | 2 + src/scope/core/pipelines/schema.py | 44 + src/scope/server/app.py | 65 + src/scope/server/download_models.py | 77 +- src/scope/server/models_config.py | 13 + src/scope/server/patch_xformers.py | 85 + src/scope/server/pipeline_manager.py | 32 + src/scope/server/schema.py | 36 + uv.lock | 1469 ++++++++++++++--- 44 files changed, 11969 insertions(+), 278 deletions(-) create mode 100644 src/scope/core/pipelines/personalive/__init__.py create mode 100644 src/scope/core/pipelines/personalive/blocks/__init__.py create mode 100644 src/scope/core/pipelines/personalive/components/__init__.py create mode 100644 src/scope/core/pipelines/personalive/components/face_detector.py create mode 100644 src/scope/core/pipelines/personalive/docs/README.md create mode 100644 src/scope/core/pipelines/personalive/model.yaml create mode 100644 src/scope/core/pipelines/personalive/modules/FAN_feature_extractor.py create mode 100644 src/scope/core/pipelines/personalive/modules/FAN_temporal_feature_extractor.py create mode 100644 src/scope/core/pipelines/personalive/modules/__init__.py create mode 100644 src/scope/core/pipelines/personalive/modules/attention.py create mode 100644 src/scope/core/pipelines/personalive/modules/camera.py create mode 100644 src/scope/core/pipelines/personalive/modules/convnextv2.py create mode 100644 src/scope/core/pipelines/personalive/modules/liveportrait_util.py create mode 100644 src/scope/core/pipelines/personalive/modules/motion_encoder.py create mode 100644 src/scope/core/pipelines/personalive/modules/motion_extractor.py create mode 100644 src/scope/core/pipelines/personalive/modules/motion_module.py create mode 100644 src/scope/core/pipelines/personalive/modules/mutual_self_attention.py create mode 100644 src/scope/core/pipelines/personalive/modules/pose_guider.py create mode 100644 src/scope/core/pipelines/personalive/modules/resnet.py create mode 100644 src/scope/core/pipelines/personalive/modules/scheduler_ddim.py create mode 100644 src/scope/core/pipelines/personalive/modules/transformer_2d.py create mode 100644 src/scope/core/pipelines/personalive/modules/transformer_3d.py create mode 100644 src/scope/core/pipelines/personalive/modules/unet_2d_blocks.py create mode 100644 src/scope/core/pipelines/personalive/modules/unet_2d_condition.py create mode 100644 src/scope/core/pipelines/personalive/modules/unet_3d.py create mode 100644 src/scope/core/pipelines/personalive/modules/unet_3d_blocks.py create mode 100644 src/scope/core/pipelines/personalive/modules/util.py create mode 100644 src/scope/core/pipelines/personalive/pipeline.py create mode 100644 src/scope/server/patch_xformers.py diff --git a/frontend/src/components/InputAndControlsPanel.tsx b/frontend/src/components/InputAndControlsPanel.tsx index a6cbbcb1..176bdef2 100644 --- a/frontend/src/components/InputAndControlsPanel.tsx +++ b/frontend/src/components/InputAndControlsPanel.tsx @@ -15,10 +15,11 @@ import { LabelWithTooltip } from "./ui/label-with-tooltip"; import type { VideoSourceMode } from "../hooks/useVideoSource"; import type { PromptItem, PromptTransition } from "../lib/api"; import type { InputMode } from "../types"; -import { pipelineIsMultiMode } from "../data/pipelines"; +import { pipelineIsMultiMode, pipelineRequiresReferenceImage } from "../data/pipelines"; import { PromptInput } from "./PromptInput"; import { TimelinePromptEditor } from "./TimelinePromptEditor"; import type { TimelinePrompt } from "./PromptTimeline"; +import { ImageIcon } from "lucide-react"; interface InputAndControlsPanelProps { className?: string; @@ -61,6 +62,10 @@ interface InputAndControlsPanelProps { onInputModeChange: (mode: InputMode) => void; // Whether Spout is available (server-side detection for native Windows, not WSL) spoutAvailable?: boolean; + // PersonaLive reference image + referenceImageUrl?: string | null; + onReferenceImageUpload?: (file: File) => void; + isUploadingReference?: boolean; } export function InputAndControlsPanel({ @@ -101,6 +106,9 @@ export function InputAndControlsPanel({ inputMode, onInputModeChange, spoutAvailable = false, + referenceImageUrl = null, + onReferenceImageUpload, + isUploadingReference = false, }: InputAndControlsPanelProps) { // Helper function to determine if playhead is at the end of timeline const isAtEndOfTimeline = () => { @@ -117,6 +125,20 @@ export function InputAndControlsPanel({ // Check if this pipeline supports multiple input modes const isMultiMode = pipelineIsMultiMode(pipelineId); + // Check if this pipeline requires a reference image (PersonaLive) + const needsReferenceImage = pipelineRequiresReferenceImage(pipelineId); + + const handleReferenceImageUpload = ( + event: React.ChangeEvent + ) => { + const file = event.target.files?.[0]; + if (file && onReferenceImageUpload) { + onReferenceImageUpload(file); + } + // Reset the input value so the same file can be selected again + event.target.value = ""; + }; + useEffect(() => { if (videoRef.current && localStream) { videoRef.current.srcObject = localStream; @@ -170,6 +192,56 @@ export function InputAndControlsPanel({ )} + {/* Reference Image upload - only show for PersonaLive */} + {needsReferenceImage && ( +
+

Reference Portrait

+
+ {referenceImageUrl ? ( + Reference portrait + ) : ( +
+ + Upload a portrait image +
+ )} + + +
+ {isUploadingReference && ( +

+ Uploading reference image... +

+ )} + {!referenceImageUrl && ( +

+ This pipeline animates this portrait using your webcam/video as the + driving source. +

+ )} +
+ )} + {/* Video Source toggle - only show when in video input mode */} {inputMode === "video" && (
@@ -280,67 +352,70 @@ export function InputAndControlsPanel({
)} -
- {(() => { - // The Input can have two states: Append (default) and Edit (when a prompt is selected and the video is paused) - const isEditMode = selectedTimelinePrompt && isVideoPaused; + {/* Prompts section - hide for PersonaLive (it uses image conditioning, not text) */} + {!needsReferenceImage && ( +
+ {(() => { + // The Input can have two states: Append (default) and Edit (when a prompt is selected and the video is paused) + const isEditMode = selectedTimelinePrompt && isVideoPaused; + + return ( +
+
+

Prompts

+ {isEditMode && ( + + Editing + + )} +
- return ( -
-
-

Prompts

- {isEditMode && ( - - Editing - + {selectedTimelinePrompt ? ( + p.id === selectedTimelinePrompt.id + )} + /> + ) : ( + )}
- - {selectedTimelinePrompt ? ( - p.id === selectedTimelinePrompt.id - )} - /> - ) : ( - - )} -
- ); - })()} -
+ ); + })()} +
+ )} ); diff --git a/frontend/src/data/pipelines.ts b/frontend/src/data/pipelines.ts index 8ab68cde..e380950e 100644 --- a/frontend/src/data/pipelines.ts +++ b/frontend/src/data/pipelines.ts @@ -18,6 +18,7 @@ export interface PipelineInfo { defaultTemporalInterpolationMethod?: "linear" | "slerp"; // Default method for temporal interpolation defaultTemporalInterpolationSteps?: number; // Default number of steps for temporal interpolation supportsLoRA?: boolean; // Whether this pipeline supports LoRA adapters + requiresReferenceImage?: boolean; // Whether this pipeline requires a reference image (e.g. PersonaLive) // Multi-mode support supportedModes: InputMode[]; @@ -98,6 +99,19 @@ export const PIPELINES: Record = { supportedModes: ["video"], defaultMode: "video", }, + personalive: { + name: "PersonaLive", + docsUrl: "https://github.com/GVCLab/PersonaLive", + about: + "Real-time portrait animation pipeline from GVCLab. Animates a reference portrait image using driving video frames to transfer expressions and head movements.", + estimatedVram: 12, + requiresModels: true, + // Video-only pipeline - requires reference image + driving video + supportedModes: ["video"], + defaultMode: "video", + // PersonaLive requires a reference image to be uploaded first + requiresReferenceImage: true, + }, }; export function pipelineSupportsLoRA(pipelineId: string): boolean { @@ -123,3 +137,7 @@ export function getPipelineDefaultMode(pipelineId: string): InputMode { export function getDefaultPromptForMode(mode: InputMode): string { return DEFAULT_PROMPTS[mode]; } + +export function pipelineRequiresReferenceImage(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.requiresReferenceImage === true; +} diff --git a/frontend/src/lib/api.ts b/frontend/src/lib/api.ts index 9e657d44..9c1305e7 100644 --- a/frontend/src/lib/api.ts +++ b/frontend/src/lib/api.ts @@ -63,6 +63,12 @@ export interface KreaRealtimeVideoLoadParams extends PipelineLoadParams { lora_merge_mode?: "permanent_merge" | "runtime_peft"; } +export interface PersonaLiveLoadParams extends PipelineLoadParams { + height?: number; + width?: number; + seed?: number; +} + export interface PipelineLoadRequest { pipeline_id?: string; load_params?: @@ -70,6 +76,7 @@ export interface PipelineLoadRequest { | StreamDiffusionV2LoadParams | LongLiveLoadParams | KreaRealtimeVideoLoadParams + | PersonaLiveLoadParams | null; } @@ -298,6 +305,37 @@ export const listLoRAFiles = async (): Promise => { return result; }; +// PersonaLive reference image upload +export interface PersonaLiveReferenceResponse { + success: boolean; + message: string; +} + +export const uploadPersonaLiveReference = async ( + imageFile: File | Blob +): Promise => { + // Read file as array buffer + const arrayBuffer = await imageFile.arrayBuffer(); + + const response = await fetch("/api/v1/personalive/reference", { + method: "POST", + headers: { + "Content-Type": imageFile.type || "image/jpeg", + }, + body: arrayBuffer, + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `PersonaLive reference upload failed: ${response.status} ${response.statusText}: ${errorText}` + ); + } + + const result = await response.json(); + return result; +}; + // Pipeline schema types - matches output of get_schema_with_metadata() export interface PipelineSchemaProperty { type?: string; diff --git a/frontend/src/pages/StreamPage.tsx b/frontend/src/pages/StreamPage.tsx index 212c07b3..0996119b 100644 --- a/frontend/src/pages/StreamPage.tsx +++ b/frontend/src/pages/StreamPage.tsx @@ -16,6 +16,7 @@ import { PIPELINES, getPipelineDefaultMode, getDefaultPromptForMode, + pipelineRequiresReferenceImage, } from "../data/pipelines"; import type { InputMode, @@ -24,7 +25,11 @@ import type { LoraMergeStrategy, } from "../types"; import type { PromptItem, PromptTransition } from "../lib/api"; -import { checkModelStatus, downloadPipelineModels } from "../lib/api"; +import { + checkModelStatus, + downloadPipelineModels, + uploadPersonaLiveReference, +} from "../lib/api"; import { sendLoRAScaleUpdates } from "../utils/loraHelpers"; // Delay before resetting video reinitialization flag (ms) @@ -102,6 +107,13 @@ export function StreamPage() { null ); + // PersonaLive reference image state + const [referenceImage, setReferenceImage] = useState(null); + const [referenceImageUrl, setReferenceImageUrl] = useState( + null + ); + const [isUploadingReference, setIsUploadingReference] = useState(false); + // Ref to access timeline functions const timelineRef = useRef<{ getCurrentTimelinePrompt: () => string; @@ -231,6 +243,9 @@ export function StreamPage() { stopStream(); } + // Clear reference image when switching pipelines + clearReferenceImage(); + const newPipeline = PIPELINES[pipelineId]; const modeToUse = newPipeline?.defaultMode || "text"; const currentMode = settings.inputMode || "text"; @@ -473,6 +488,25 @@ export function StreamPage() { }); }; + // Handle reference image upload for PersonaLive + const handleReferenceImageUpload = (file: File) => { + // Clean up old URL + if (referenceImageUrl) { + URL.revokeObjectURL(referenceImageUrl); + } + setReferenceImage(file); + setReferenceImageUrl(URL.createObjectURL(file)); + }; + + // Clear reference image (used when switching pipelines) + const clearReferenceImage = () => { + if (referenceImageUrl) { + URL.revokeObjectURL(referenceImageUrl); + } + setReferenceImage(null); + setReferenceImageUrl(null); + }; + // Sync spoutReceiver.enabled with mode changes const handleModeChange = (newMode: typeof mode) => { // When switching to spout mode, enable spout input @@ -686,6 +720,22 @@ export function StreamPage() { console.log( `Loading with resolution: ${resolution.width}x${resolution.height}, seed: ${loadParams.seed}, quantization: ${loadParams.quantization}, lora_merge_mode: ${loadParams.lora_merge_mode}` ); + } else if (pipelineIdToUse === "personalive" && resolution) { + // PersonaLive requires a reference image + if (!referenceImage) { + console.error( + "PersonaLive requires a reference image. Please upload one first." + ); + return false; + } + loadParams = { + height: resolution.height, + width: resolution.width, + seed: settings.seed ?? 42, + }; + console.log( + `Loading PersonaLive with resolution: ${resolution.width}x${resolution.height}, seed: ${loadParams.seed}` + ); } const loadSuccess = await loadPipeline( @@ -697,6 +747,22 @@ export function StreamPage() { return false; } + // For PersonaLive, upload reference image after pipeline is loaded + if (pipelineIdToUse === "personalive" && referenceImage) { + try { + setIsUploadingReference(true); + console.log("Uploading PersonaLive reference image..."); + const result = await uploadPersonaLiveReference(referenceImage); + console.log("Reference image uploaded:", result.message); + } catch (error) { + console.error("Failed to upload reference image:", error); + setIsUploadingReference(false); + return false; + } finally { + setIsUploadingReference(false); + } + } + // Check video requirements based on input mode const needsVideoInput = currentMode === "video"; const isSpoutMode = mode === "spout" && settings.spoutReceiver?.enabled; @@ -729,7 +795,11 @@ export function StreamPage() { }; // Common parameters for pipelines that support prompts - if (pipelineIdToUse !== "passthrough") { + // PersonaLive and Passthrough don't use text prompts + if ( + pipelineIdToUse !== "passthrough" && + pipelineIdToUse !== "personalive" + ) { initialParameters.prompts = promptItems; initialParameters.prompt_interpolation_method = interpolationMethod; initialParameters.denoising_step_list = settings.denoisingSteps || [ @@ -799,11 +869,17 @@ export function StreamPage() { isConnecting={isConnecting} isPipelineLoading={isPipelineLoading} canStartStream={ - settings.inputMode === "text" - ? !isInitializing - : mode === "spout" - ? !isInitializing // Spout mode doesn't need local stream - : !!localStream && !isInitializing + // PersonaLive requires reference image + pipelineRequiresReferenceImage(settings.pipelineId) + ? !!referenceImage && + (mode === "spout" + ? !isInitializing + : !!localStream && !isInitializing) + : settings.inputMode === "text" + ? !isInitializing + : mode === "spout" + ? !isInitializing // Spout mode doesn't need local stream + : !!localStream && !isInitializing } onStartStream={handleStartStream} onStopStream={stopStream} @@ -834,6 +910,9 @@ export function StreamPage() { } onInputModeChange={handleInputModeChange} spoutAvailable={spoutAvailable} + referenceImageUrl={referenceImageUrl} + onReferenceImageUpload={handleReferenceImageUpload} + isUploadingReference={isUploadingReference} />
@@ -947,6 +1026,7 @@ export function StreamPage() { }} disabled={ settings.pipelineId === "passthrough" || + settings.pipelineId === "personalive" || isPipelineLoading || isConnecting || showDownloadDialog diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 95bf3e1a..ac4c0c74 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -3,7 +3,8 @@ export type PipelineId = | "passthrough" | "longlive" | "krea-realtime-video" - | "reward-forcing"; + | "reward-forcing" + | "personalive"; // Input mode for pipeline operation export type InputMode = "text" | "video"; @@ -87,6 +88,7 @@ export interface PipelineInfo { defaultTemporalInterpolationMethod?: "linear" | "slerp"; defaultTemporalInterpolationSteps?: number; supportsLoRA?: boolean; + requiresReferenceImage?: boolean; // Whether this pipeline requires a reference image (e.g. PersonaLive) // Multi-mode support supportedModes: InputMode[]; diff --git a/pyproject.toml b/pyproject.toml index 307642f9..edcd08b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,12 +55,18 @@ dependencies = [ "triton-windows==3.4.0.post21; sys_platform == 'win32'", "SpoutGL>=0.1.1; sys_platform == 'win32'", "PyOpenGL>=3.1.10; sys_platform == 'win32'", + # PersonaLive dependencies + "mediapipe>=0.10.11", + "decord>=0.6.0", + "xformers>=0.0.32.post2", + "scikit-image>=0.22.0", ] [project.scripts] daydream-scope = "scope.server.app:main" build = "scope.server.build:main" download_models = "scope.server.download_models:main" +patch-xformers = "scope.server.patch_xformers:main" [project.urls] Homepage = "https://github.com/daydreamlive/scope" diff --git a/src/scope/core/pipelines/__init__.py b/src/scope/core/pipelines/__init__.py index 9ea63303..02ecf5c1 100644 --- a/src/scope/core/pipelines/__init__.py +++ b/src/scope/core/pipelines/__init__.py @@ -24,6 +24,10 @@ def __getattr__(name): from .passthrough.pipeline import PassthroughPipeline return PassthroughPipeline + elif name == "PersonaLivePipeline": + from .personalive.pipeline import PersonaLivePipeline + + return PersonaLivePipeline # Config classes elif name == "BasePipelineConfig": from .schema import BasePipelineConfig @@ -45,6 +49,10 @@ def __getattr__(name): from .schema import PassthroughConfig return PassthroughConfig + elif name == "PersonaLiveConfig": + from .schema import PersonaLiveConfig + + return PersonaLiveConfig raise AttributeError(f"module {__name__!r} has no attribute {name!r}") @@ -55,10 +63,12 @@ def __getattr__(name): "RewardForcingPipeline", "StreamDiffusionV2Pipeline", "PassthroughPipeline", + "PersonaLivePipeline", # Config classes "BasePipelineConfig", "LongLiveConfig", "StreamDiffusionV2Config", "KreaRealtimeVideoConfig", "PassthroughConfig", + "PersonaLiveConfig", ] diff --git a/src/scope/core/pipelines/personalive/__init__.py b/src/scope/core/pipelines/personalive/__init__.py new file mode 100644 index 00000000..f3995ef9 --- /dev/null +++ b/src/scope/core/pipelines/personalive/__init__.py @@ -0,0 +1,5 @@ +"""PersonaLive pipeline for real-time portrait animation.""" + +from .pipeline import PersonaLivePipeline + +__all__ = ["PersonaLivePipeline"] diff --git a/src/scope/core/pipelines/personalive/blocks/__init__.py b/src/scope/core/pipelines/personalive/blocks/__init__.py new file mode 100644 index 00000000..fdb7b942 --- /dev/null +++ b/src/scope/core/pipelines/personalive/blocks/__init__.py @@ -0,0 +1 @@ +"""PersonaLive modular blocks for pipeline composition.""" diff --git a/src/scope/core/pipelines/personalive/components/__init__.py b/src/scope/core/pipelines/personalive/components/__init__.py new file mode 100644 index 00000000..941b0908 --- /dev/null +++ b/src/scope/core/pipelines/personalive/components/__init__.py @@ -0,0 +1,7 @@ +"""PersonaLive model component wrappers.""" + +from .face_detector import FaceDetector + +__all__ = [ + "FaceDetector", +] diff --git a/src/scope/core/pipelines/personalive/components/face_detector.py b/src/scope/core/pipelines/personalive/components/face_detector.py new file mode 100644 index 00000000..9724c95c --- /dev/null +++ b/src/scope/core/pipelines/personalive/components/face_detector.py @@ -0,0 +1,199 @@ +"""Face detection component using MediaPipe for PersonaLive pipeline.""" + +import logging +from typing import Tuple + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + +try: + import mediapipe as mp + MEDIAPIPE_AVAILABLE = True +except ImportError: + MEDIAPIPE_AVAILABLE = False + logger.warning("MediaPipe not available. Face detection will be disabled.") + + +class FaceDetector: + """Face detector using MediaPipe Face Mesh. + + Extracts face regions from images for motion encoding. + """ + + def __init__(self, static_image_mode: bool = True, max_num_faces: int = 1): + """Initialize the face detector. + + Args: + static_image_mode: If True, treats input as unrelated images (no tracking). + max_num_faces: Maximum number of faces to detect. + """ + if not MEDIAPIPE_AVAILABLE: + raise RuntimeError( + "MediaPipe is required for PersonaLive face detection. " + "Install with: pip install mediapipe" + ) + + self.mp_face_mesh = mp.solutions.face_mesh + self.face_mesh = self.mp_face_mesh.FaceMesh( + static_image_mode=static_image_mode, + max_num_faces=max_num_faces, + min_detection_confidence=0.5, + min_tracking_confidence=0.5, + ) + + def detect_face_landmarks(self, image: np.ndarray) -> list | None: + """Detect face landmarks in an image. + + Args: + image: RGB image as numpy array (H, W, 3), uint8. + + Returns: + List of face landmarks or None if no face detected. + """ + results = self.face_mesh.process(image) + if results.multi_face_landmarks: + return results.multi_face_landmarks + return None + + def get_face_bounding_box( + self, + image: np.ndarray, + padding_ratio: float = 0.2 + ) -> Tuple[int, int, int, int] | None: + """Get bounding box for the detected face. + + Args: + image: RGB image as numpy array (H, W, 3). + padding_ratio: Extra padding around the face box. + + Returns: + Tuple of (left, top, right, bottom) or None if no face. + """ + landmarks = self.detect_face_landmarks(image) + if landmarks is None: + return None + + h, w = image.shape[:2] + face_landmarks = landmarks[0] + + # Get all landmark coordinates + xs = [lm.x * w for lm in face_landmarks.landmark] + ys = [lm.y * h for lm in face_landmarks.landmark] + + # Compute bounding box + left = min(xs) + right = max(xs) + top = min(ys) + bottom = max(ys) + + # Add padding + box_w = right - left + box_h = bottom - top + pad_w = box_w * padding_ratio + pad_h = box_h * padding_ratio + + left = max(0, int(left - pad_w)) + top = max(0, int(top - pad_h)) + right = min(w, int(right + pad_w)) + bottom = min(h, int(bottom + pad_h)) + + return (left, top, right, bottom) + + def crop_face( + self, + image: np.ndarray, + target_size: Tuple[int, int] = (512, 512), + padding_ratio: float = 0.2, + ) -> np.ndarray | None: + """Crop and resize face from image. + + Args: + image: RGB image as numpy array (H, W, 3). + target_size: Size to resize the cropped face to (W, H). + padding_ratio: Extra padding around the face box. + + Returns: + Cropped and resized face image, or None if no face detected. + """ + from PIL import Image + + bbox = self.get_face_bounding_box(image, padding_ratio) + if bbox is None: + return None + + left, top, right, bottom = bbox + face_patch = image[top:bottom, left:right] + + # Resize to target size + face_pil = Image.fromarray(face_patch) + face_pil = face_pil.resize(target_size, Image.BILINEAR) + + return np.array(face_pil) + + def crop_face_from_pil(self, pil_image, target_size: Tuple[int, int] = (512, 512)): + """Crop face from a PIL image. + + Args: + pil_image: PIL Image in RGB mode. + target_size: Size to resize the cropped face to (W, H). + + Returns: + PIL Image of cropped face, or None if no face detected. + """ + from PIL import Image + + image_np = np.array(pil_image) + face_np = self.crop_face(image_np, target_size) + + if face_np is None: + return None + + return Image.fromarray(face_np) + + def crop_face_tensor( + self, + image_tensor: torch.Tensor, + boxes: Tuple[int, int, int, int], + target_size: Tuple[int, int] = (224, 224), + ) -> torch.Tensor: + """Crop face region from a tensor using precomputed box. + + Args: + image_tensor: Image tensor of shape (C, H, W) or (B, C, H, W). + boxes: Bounding box as (left, top, right, bottom). + target_size: Output size (H, W). + + Returns: + Cropped and resized face tensor. + """ + import torch.nn.functional as F + + left, top, right, bottom = map(int, boxes) + + if image_tensor.dim() == 3: + # (C, H, W) -> crop + face_patch = image_tensor[:, top:bottom, left:right] + face_patch = face_patch.unsqueeze(0) # Add batch dim for interpolate + else: + # (B, C, H, W) + face_patch = image_tensor[:, :, top:bottom, left:right] + + # Resize + face_patch = F.interpolate( + face_patch, + size=target_size, + mode="bilinear", + align_corners=False, + ) + + return face_patch + + def close(self): + """Release resources.""" + if hasattr(self, 'face_mesh') and self.face_mesh: + self.face_mesh.close() + + def __del__(self): + self.close() diff --git a/src/scope/core/pipelines/personalive/docs/README.md b/src/scope/core/pipelines/personalive/docs/README.md new file mode 100644 index 00000000..da470be1 --- /dev/null +++ b/src/scope/core/pipelines/personalive/docs/README.md @@ -0,0 +1,124 @@ +# PersonaLive Pipeline + +PersonaLive is a real-time portrait animation pipeline that animates a reference portrait image using driving video frames. + +## Overview + +PersonaLive takes: +1. A **reference image** - A portrait photograph of the person to animate +2. **Driving video frames** - Video frames that provide the motion/expression to transfer + +And outputs animated video frames of the reference person following the expressions and movements from the driving video. + +## Installation + +### 1. Download Models + +```bash +# Download PersonaLive models +download_models --pipeline personalive +``` + +This downloads from [huggingface.co/huaichang/PersonaLive](https://huggingface.co/huaichang/PersonaLive), which includes: +- PersonaLive custom weights (denoising_unet, reference_unet, motion_encoder, etc.) +- sd-image-variations-diffusers base model +- sd-vae-ft-mse VAE + +### 2. Install Dependencies + +Make sure you have the required dependencies installed: +- `mediapipe>=0.10.11` - For face detection and keypoint extraction +- `decord>=0.6.0` - For video decoding +- `xformers>=0.0.28` - For memory-efficient attention + +## Usage + +### API Usage + +1. **Load the pipeline**: + +```bash +POST /api/v1/pipeline/load +Content-Type: application/json + +{ + "pipeline_id": "personalive", + "load_params": { + "height": 512, + "width": 512, + "seed": 42 + } +} +``` + +2. **Upload reference image**: + +```bash +POST /api/v1/personalive/reference +Content-Type: image/jpeg + + +``` + +3. **Start WebRTC stream**: + +Connect via WebRTC and send driving video frames. The output will be animated frames of the reference portrait. + +### Workflow + +``` +1. Load PersonaLive pipeline +2. Upload reference portrait image +3. Connect WebRTC stream with camera/video input +4. Receive animated output frames +``` + +## Technical Details + +### Architecture + +PersonaLive uses a multi-component architecture: + +- **Reference UNet (2D)**: Extracts features from the reference image +- **Denoising UNet (3D)**: Generates animated frames with temporal consistency +- **Motion Encoder**: Extracts motion features from driving video faces +- **Pose Encoder (MotionExtractor)**: Extracts keypoints from faces +- **Pose Guider**: Converts keypoints to conditioning features +- **VAE**: Encodes/decodes latent representations +- **Image Encoder (CLIP)**: Extracts semantic features from reference image + +### Processing Flow + +1. **Reference Fusion** (once at start): + - Encode reference image with CLIP + - Encode reference image with VAE + - Run Reference UNet to cache attention features + - Extract reference face keypoints + +2. **Frame Processing** (for each driving frame batch): + - Extract face keypoints from driving frames + - Compute motion features from cropped face regions + - Generate pose condition embeddings + - Denoise latents with temporal attention + - Decode output frames from latents + +### Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `height` | 512 | Output video height | +| `width` | 512 | Output video width | +| `seed` | 42 | Random seed for reproducibility | +| `temporal_window_size` | 4 | Number of frames processed together | +| `temporal_adaptive_step` | 4 | Temporal denoising step | +| `num_inference_steps` | 4 | Denoising steps (typically 4 for real-time) | + +## References + +- [PersonaLive Paper](https://arxiv.org/abs/2506.09887) +- [GitHub Repository](https://github.com/GVCLab/PersonaLive) +- [HuggingFace Models](https://huggingface.co/huaichang/PersonaLive) + +## Credits + +Based on the PersonaLive implementation by GVCLab. diff --git a/src/scope/core/pipelines/personalive/model.yaml b/src/scope/core/pipelines/personalive/model.yaml new file mode 100644 index 00000000..520642fb --- /dev/null +++ b/src/scope/core/pipelines/personalive/model.yaml @@ -0,0 +1,70 @@ +# PersonaLive model configuration +# Based on https://github.com/GVCLab/PersonaLive + +# Base model paths (relative to models_dir) +pretrained_base_model_path: "sd-image-variations-diffusers" +image_encoder_path: "sd-image-variations-diffusers/image_encoder" +vae_model_path: "sd-vae-ft-mse" + +# PersonaLive custom weights +personalive_weights_path: "PersonaLive/pretrained_weights/personalive" + +# Inference configuration +dtype: "fp16" +batch_size: 1 +temporal_window_size: 4 +temporal_adaptive_step: 4 +num_inference_steps: 4 +seed: 42 + +# Resolution +reference_image_height: 512 +reference_image_width: 512 +output_height: 512 +output_width: 512 + +# Scheduler configuration +noise_scheduler_kwargs: + beta_start: 0.00085 + beta_end: 0.02 + beta_schedule: "scaled_linear" + clip_sample: false + steps_offset: 1 + prediction_type: "epsilon" + timestep_spacing: "trailing" + +# UNet additional kwargs for 3D UNet +unet_additional_kwargs: + use_inflated_groupnorm: true + unet_use_cross_frame_attention: false + unet_use_temporal_attention: false + use_motion_module: true + motion_module_resolutions: + - 1 + - 2 + - 4 + - 8 + motion_module_mid_block: true + motion_module_decoder_only: false + motion_module_type: "Vanilla" + motion_module_kwargs: + num_attention_heads: 8 + num_transformer_block: 1 + cross_attention_dim: 16 + attention_block_types: + - "Spatial_Cross" + - "Spatial_Cross" + temporal_position_encoding: false + temporal_position_encoding_max_len: 32 + temporal_attention_dim_div: 1 + use_temporal_module: true + temporal_module_type: "Vanilla" + temporal_module_kwargs: + num_attention_heads: 8 + num_transformer_block: 1 + attention_block_types: + - "Temporal_Self" + - "Temporal_Self" + temporal_position_encoding: true + temporal_position_encoding_max_len: 32 + temporal_attention_dim_div: 1 diff --git a/src/scope/core/pipelines/personalive/modules/FAN_feature_extractor.py b/src/scope/core/pipelines/personalive/modules/FAN_feature_extractor.py new file mode 100644 index 00000000..1fca87de --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/FAN_feature_extractor.py @@ -0,0 +1,334 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import Dict, Union +from diffusers.models.attention_processor import AttentionProcessor + + +def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, + stride=strd, padding=padding, bias=bias) + + +class ConvBlock(nn.Module): + def __init__(self, in_planes, out_planes): + super(ConvBlock, self).__init__() + self.bn1 = nn.BatchNorm2d(in_planes) + self.conv1 = conv3x3(in_planes, int(out_planes / 2)) + self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) + self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4)) + self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) + self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4)) + + if in_planes != out_planes: + self.downsample = nn.Sequential( + nn.BatchNorm2d(in_planes, eps=1e-4), + nn.ReLU(True), + nn.Conv2d(in_planes, out_planes, + kernel_size=1, stride=1, bias=False), + ) + else: + self.downsample = None + + def forward(self, x): + residual = x + + out1 = self.bn1(x) + out1 = F.relu(out1, True) + out1 = self.conv1(out1) + + out2 = self.bn2(out1) + out2 = F.relu(out2, True) + out2 = self.conv2(out2) + + out3 = self.bn3(out2) + out3 = F.relu(out3, True) + out3 = self.conv3(out3) + + out3 = torch.cat((out1, out2, out3), 1) + + if self.downsample is not None: + residual = self.downsample(residual) + + out3 += residual + + return out3 + + +class HourGlass(nn.Module): + def __init__(self, num_modules, depth, num_features): + super(HourGlass, self).__init__() + self.num_modules = num_modules + self.depth = depth + self.features = num_features + self.dropout = nn.Dropout(0.5) + + self._generate_network(self.depth) + + def _generate_network(self, level): + self.add_module('b1_' + str(level), ConvBlock(256, 256)) + + self.add_module('b2_' + str(level), ConvBlock(256, 256)) + + if level > 1: + self._generate_network(level - 1) + else: + self.add_module('b2_plus_' + str(level), ConvBlock(256, 256)) + + self.add_module('b3_' + str(level), ConvBlock(256, 256)) + + def _forward(self, level, inp): + # Upper branch + up1 = inp + up1 = self._modules['b1_' + str(level)](up1) + up1 = self.dropout(up1) + # Lower branch + low1 = F.max_pool2d(inp, 2, stride=2) + low1 = self._modules['b2_' + str(level)](low1) + + if level > 1: + low2 = self._forward(level - 1, low1) + else: + low2 = low1 + low2 = self._modules['b2_plus_' + str(level)](low2) + + low3 = low2 + low3 = self._modules['b3_' + str(level)](low3) + up1size = up1.size() + rescale_size = (up1size[2], up1size[3]) + up2 = F.upsample(low3, size=rescale_size, mode='bilinear') + + return up1 + up2 + + def forward(self, x): + return self._forward(self.depth, x) + + +class FAN_use(nn.Module): + def __init__(self): + super(FAN_use, self).__init__() + self.num_modules = 1 + + # Base part + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) + self.bn1 = nn.BatchNorm2d(64) + self.conv2 = ConvBlock(64, 128) + self.conv3 = ConvBlock(128, 128) + self.conv4 = ConvBlock(128, 256) + + # Stacking part + hg_module = 0 + self.add_module('m' + str(hg_module), HourGlass(1, 4, 256)) + self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256)) + self.add_module('conv_last' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module('l' + str(hg_module), nn.Conv2d(256, 68, kernel_size=1, stride=1, padding=0)) + self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) + + if hg_module < self.num_modules - 1: + self.add_module('bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module('al' + str(hg_module), nn.Conv2d(68, 256, kernel_size=1, stride=1, padding=0)) + + self.avgpool = nn.MaxPool2d((2, 2), 2) + self.conv6 = nn.Conv2d(68, 1, 3, 2, 1) + self.fc = nn.Linear(28 * 28, 512) + self.bn5 = nn.BatchNorm2d(68) + self.relu = nn.ReLU(True) + + def forward(self, x, return_featmap=False): + x = F.relu(self.bn1(self.conv1(x)), True) # 112 + x = F.max_pool2d(self.conv2(x), 2) # 56 # [B, 128, 112, 112] + x = self.conv3(x) + x = self.conv4(x) # [B, 256, 56, 56] + + previous = x + + i = 0 + hg = self._modules['m' + str(i)](previous) + + ll = hg + ll = self._modules['top_m_' + str(i)](ll) + + ll = self._modules['bn_end' + str(i)](self._modules['conv_last' + str(i)](ll)) # [B, 256, 56, 56] + if return_featmap: + return ll + tmp_out = self._modules['l' + str(i)](F.relu(ll)) + + net = self.relu(self.bn5(tmp_out)) # [B, 68, 56, 56] + net = self.conv6(net) # 28 # [B, 1, 28, 28] + net = net.view(-1, net.shape[-2] * net.shape[-1]) + net = self.relu(net) + net = self.fc(net) + return net + +from .FAN_temporal_feature_extractor import TemporalTransformer3DModel +from einops import rearrange + + +class FAN_SA(nn.Module): + def __init__(self): + super(FAN_SA, self).__init__() + self.num_modules = 1 + + # Base part + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) + self.bn1 = nn.BatchNorm2d(64) + self.conv2 = ConvBlock(64, 128) + self.conv3 = ConvBlock(128, 128) + self.conv4 = ConvBlock(128, 256) + + # Stacking part + hg_module = 0 + self.add_module('m' + str(hg_module), HourGlass(1, 4, 256)) + self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256)) + self.add_module( + 'conv_last' + str(hg_module), + nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0), + ) + self.add_module( + 'l' + str(hg_module), nn.Conv2d(256, 68, kernel_size=1, stride=1, padding=0) + ) + self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) + + if hg_module < self.num_modules - 1: + self.add_module( + 'bl' + str(hg_module), + nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0), + ) + self.add_module( + 'al' + str(hg_module), + nn.Conv2d(68, 256, kernel_size=1, stride=1, padding=0), + ) + + self.avgpool = nn.MaxPool2d((2, 2), 2) + self.conv6 = nn.Conv2d(68, 1, 3, 2, 1) + self.fc = nn.Linear(28 * 28, 512) + # self.conv6 = nn.Conv2d(68, 2, 3, 2, 1) + # self.fc = nn.Linear(28 * 28 * 2, 1024) + self.bn5 = nn.BatchNorm2d(68) + self.relu = nn.ReLU(True) + + # Add by zxc + self.att_1 = TemporalTransformer3DModel( + in_channels=128, + sample_size=112, + patch_size=4, + attention_block_types=("Spatial_Self",), + zero_initialize=True, + ) + self.att_2 = TemporalTransformer3DModel( + in_channels=256, + sample_size=56, + patch_size=2, + attention_block_types=("Spatial_Self",), + zero_initialize=True, + ) + self.att_3 = TemporalTransformer3DModel( + in_channels=256, + sample_size=56, + patch_size=2, + attention_block_types=("Spatial_Self",), + zero_initialize=True, + ) + + @property + # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors(name: str, module: torch.nn.Module, processors: Dict[str, AttentionProcessor]): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor(return_deprecated_lora=True) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor( + self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x)), True) # 112 + x = self.conv2(x) # [B, 128, 112, 112] + + # Temp Self-Att: [B*h*w, T, 128*p*p] [B*28*28, T, 1024] + x = rearrange(x, "(b f) c h w -> b c f h w", f=1) + x = self.att_1(x, skip=True)[:, :, 0] + + x = F.max_pool2d(x, 2) # 56 + x = self.conv3(x) + x = self.conv4(x) # [B, 256, 56, 56] + + # Temp Self-Att: [B*h*w, T, 256*p*p] [B*28*28, T, 1024] + x = rearrange(x, "(b f) c h w -> b c f h w", f=1) + x = self.att_2(x, skip=True)[:, :, 0] + + previous = x + + i = 0 + hg = self._modules['m' + str(i)](previous) + + ll = hg + ll = self._modules['top_m_' + str(i)](ll) + + ll = self._modules['bn_end' + str(i)]( + self._modules['conv_last' + str(i)](ll) + ) # [B, 256, 56, 56] + # Temp Cross-Att: [B*28*28, 1, 1024]*[B*28*28, T, 1024] + ll = rearrange(ll, "(b f) c h w -> b c f h w", f=1) + ll = self.att_3(ll, skip=True)[:, :, 0] # "b c 1 h w -> b c h w" + # print('att3', torch.abs(ll).mean().item()) + + tmp_out = self._modules['l' + str(i)](F.relu(ll)) + + net = self.relu(self.bn5(tmp_out)) # [B, 68, 56, 56] + net = self.conv6(net) # 28 # [B, 1, 28, 28] + net = net.view(-1, net.shape[-2] * net.shape[-1] * net.shape[1]) + net = self.relu(net) + net = self.fc(net) + return net diff --git a/src/scope/core/pipelines/personalive/modules/FAN_temporal_feature_extractor.py b/src/scope/core/pipelines/personalive/modules/FAN_temporal_feature_extractor.py new file mode 100644 index 00000000..d1968803 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/FAN_temporal_feature_extractor.py @@ -0,0 +1,231 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import PatchEmbed + +from einops import rearrange +from .motion_module import TemporalTransformerBlock, zero_module, random_module + +def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, + stride=strd, padding=padding, bias=bias) + + +class ConvBlock(nn.Module): + def __init__(self, in_planes, out_planes): + super(ConvBlock, self).__init__() + self.bn1 = nn.BatchNorm2d(in_planes) + self.conv1 = conv3x3(in_planes, int(out_planes / 2)) + self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) + self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4)) + self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) + self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4)) + + if in_planes != out_planes: + self.downsample = nn.Sequential( + nn.BatchNorm2d(in_planes), + nn.ReLU(True), + nn.Conv2d(in_planes, out_planes, + kernel_size=1, stride=1, bias=False), + ) + else: + self.downsample = None + + def forward(self, x): + residual = x + + out1 = self.bn1(x) + out1 = F.relu(out1, True) + out1 = self.conv1(out1) + + out2 = self.bn2(out1) + out2 = F.relu(out2, True) + out2 = self.conv2(out2) + + out3 = self.bn3(out2) + out3 = F.relu(out3, True) + out3 = self.conv3(out3) + + out3 = torch.cat((out1, out2, out3), 1) + + if self.downsample is not None: + residual = self.downsample(residual) + + out3 += residual + + return out3 + + +class HourGlass(nn.Module): + def __init__(self, num_modules, depth, num_features): + super(HourGlass, self).__init__() + self.num_modules = num_modules + self.depth = depth + self.features = num_features + self.dropout = nn.Dropout(0.5) + + self._generate_network(self.depth) + + def _generate_network(self, level): + self.add_module('b1_' + str(level), ConvBlock(256, 256)) + + self.add_module('b2_' + str(level), ConvBlock(256, 256)) + + if level > 1: + self._generate_network(level - 1) + else: + self.add_module('b2_plus_' + str(level), ConvBlock(256, 256)) + + self.add_module('b3_' + str(level), ConvBlock(256, 256)) + + def _forward(self, level, inp): + # Upper branch + up1 = inp + up1 = self._modules['b1_' + str(level)](up1) + up1 = self.dropout(up1) + # Lower branch + low1 = F.max_pool2d(inp, 2, stride=2) + low1 = self._modules['b2_' + str(level)](low1) + + if level > 1: + low2 = self._forward(level - 1, low1) + else: + low2 = low1 + low2 = self._modules['b2_plus_' + str(level)](low2) + + low3 = low2 + low3 = self._modules['b3_' + str(level)](low3) + up1size = up1.size() + rescale_size = (up1size[2], up1size[3]) + up2 = F.upsample(low3, size=rescale_size, mode='bilinear') + + return up1 + up2 + + def forward(self, x): + return self._forward(self.depth, x) + +class TemporalTransformer3DModel(nn.Module): + def __init__( + self, + in_channels, + num_attention_heads=8, + attention_head_dim=128, + num_layers=1, + attention_block_types=("Temporal_Self",), + sample_size=56, + patch_size=2, + dropout=0.0, + norm_num_groups=32, + cross_attention_dim=0, + activation_fn="geglu", + attention_bias=False, + upcast_attention=False, + cross_frame_attention_mode=None, + temporal_position_encoding=True, + temporal_position_encoding_max_len=24, + zero_initialize=True + ): + super().__init__() + self.in_channels = in_channels + inner_dim = num_attention_heads * attention_head_dim + + self.norm = torch.nn.GroupNorm( + num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) + # self.proj_in = nn.Linear(in_channels, inner_dim) # patchEmbed中有 + if cross_attention_dim == 0 and any([block_name.endswith("_Cross") for block_name in attention_block_types]): + cross_attention_dim = inner_dim + self.cross_frame_attention_mode = cross_frame_attention_mode + if cross_frame_attention_mode is not None: + assert cross_frame_attention_mode in ['Spatial', 'Temporal'], cross_attention_dim + assert any([block_name.endswith("_Cross") for block_name in attention_block_types]) + self.transformer_blocks = nn.ModuleList( + [ + TemporalTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + attention_block_types=attention_block_types, + dropout=dropout, + norm_num_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + attention_bias=attention_bias, + upcast_attention=upcast_attention, + cross_frame_attention_mode=None, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ) + for _ in range(num_layers) + ] + ) + self.proj_out = nn.Linear(inner_dim // (patch_size**2), in_channels) + + self.patch_size = patch_size + self.sample_size = sample_size + self.pos_embed = PatchEmbed( + height=sample_size, + width=sample_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=inner_dim, + interpolation_scale=1, + ) + # 在转onnx模型时,需要设置为False + self.pos_embed.pos_embed.requires_grad = False + self.proj_out = zero_module(self.proj_out) if zero_initialize else random_module(self.proj_out) + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, skip=True): + assert hidden_states.dim() == 5, f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}." + batch, _, video_length = hidden_states.shape[:3] + residual = hidden_states + hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w") + + _, _, height, width = hidden_states.shape + + hidden_states = self.norm(hidden_states) # 需要先norm,再加posEmb,否则Emb偏执太大了 + grid_h, grid_w = height // self.patch_size, width // self.patch_size, + + hidden_states = self.pos_embed(hidden_states) # [(bf), l, c] # 偏执太大了 + # print( + # round(torch.abs(hidden_states).mean().item(), 6), + # round(torch.abs(norm_hidden_states).mean().item(), 6), + # ) + if self.cross_frame_attention_mode is not None: + assert encoder_hidden_states is None + encoder_hidden_states = rearrange(hidden_states, "(b f) d c -> b f d c", b=batch) + + hidden_states = encoder_hidden_states[:, -1] # [b, d, c] + residual = residual[:, :, -1:] + video_length = 1 + if self.cross_frame_attention_mode == 'Temporal': # 用所有帧来condition最后一帧(当前帧) + encoder_hidden_states = rearrange(encoder_hidden_states, "b f d c -> (b d) f c") + elif self.cross_frame_attention_mode == 'Spatial': + encoder_hidden_states = rearrange(encoder_hidden_states, "b f d c -> b (f d) c") + + # Transformer Blocks + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + video_length=video_length, + att_flag=False + ) + + hidden_states = hidden_states.reshape( + shape=(batch * video_length, grid_h, grid_w, self.patch_size, self.patch_size, -1) + ) + hidden_states = self.proj_out(hidden_states) + hidden_states = torch.einsum("nhwpqc->nchpwq", hidden_states) + hidden_states = hidden_states.reshape(shape=(batch * video_length, self.in_channels, height, width)) + + hidden_states = rearrange(hidden_states, "(b f) c h w -> b c f h w", b=batch) + # print( + # 'out', + # round(torch.abs(hidden_states).mean().item(), 6), + # round(torch.abs(residual).mean().item(), 6), + # ) + output = (hidden_states + residual) if skip else hidden_states + + return output diff --git a/src/scope/core/pipelines/personalive/modules/__init__.py b/src/scope/core/pipelines/personalive/modules/__init__.py new file mode 100644 index 00000000..b00bbbeb --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/__init__.py @@ -0,0 +1,30 @@ +"""PersonaLive neural network modules. + +These are ported from the original PersonaLive implementation at +https://github.com/GVCLab/PersonaLive +""" + +from .motion_module import zero_module, get_motion_module +from .mutual_self_attention import ReferenceAttentionControl +from .unet_2d_condition import UNet2DConditionModel +from .unet_3d import UNet3DConditionModel +from .motion_encoder import MotEncoder +from .pose_guider import PoseGuider +from .motion_extractor import MotionExtractor +from .scheduler_ddim import DDIMScheduler +from .util import draw_keypoints, get_boxes, crop_face + +__all__ = [ + "zero_module", + "get_motion_module", + "ReferenceAttentionControl", + "UNet2DConditionModel", + "UNet3DConditionModel", + "MotEncoder", + "PoseGuider", + "MotionExtractor", + "DDIMScheduler", + "draw_keypoints", + "get_boxes", + "crop_face", +] diff --git a/src/scope/core/pipelines/personalive/modules/attention.py b/src/scope/core/pipelines/personalive/modules/attention.py new file mode 100644 index 00000000..fd8b4359 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/attention.py @@ -0,0 +1,463 @@ +from typing import Any, Dict, Optional + +import torch +from diffusers.models.attention import AdaLayerNorm, Attention, FeedForward +from diffusers.models.embeddings import SinusoidalPositionalEmbedding +from einops import rearrange +from torch import nn + +from diffusers.models.attention import * +from diffusers.models.attention_processor import * + +class BasicTransformerBlock(nn.Module): + r""" + A basic Transformer block. + + Parameters: + dim (`int`): The number of channels in the input and output. + num_attention_heads (`int`): The number of heads to use for multi-head attention. + attention_head_dim (`int`): The number of channels in each head. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward. + num_embeds_ada_norm (: + obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`. + attention_bias (: + obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter. + only_cross_attention (`bool`, *optional*): + Whether to use only cross-attention layers. In this case two cross attention layers are used. + double_self_attention (`bool`, *optional*): + Whether to use two self-attention layers. In this case no cross attention layers are used. + upcast_attention (`bool`, *optional*): + Whether to upcast the attention computation to float32. This is useful for mixed precision training. + norm_elementwise_affine (`bool`, *optional*, defaults to `True`): + Whether to use learnable elementwise affine parameters for normalization. + norm_type (`str`, *optional*, defaults to `"layer_norm"`): + The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`. + final_dropout (`bool` *optional*, defaults to False): + Whether to apply a final dropout after the last feed-forward layer. + attention_type (`str`, *optional*, defaults to `"default"`): + The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`. + positional_embeddings (`str`, *optional*, defaults to `None`): + The type of positional embeddings to apply to. + num_positional_embeddings (`int`, *optional*, defaults to `None`): + The maximum number of positional embeddings to apply. + """ + + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single' + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: Optional[str] = None, + num_positional_embeddings: Optional[int] = None, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + + self.use_ada_layer_norm_zero = ( + num_embeds_ada_norm is not None + ) and norm_type == "ada_norm_zero" + self.use_ada_layer_norm = ( + num_embeds_ada_norm is not None + ) and norm_type == "ada_norm" + self.use_ada_layer_norm_single = norm_type == "ada_norm_single" + self.use_layer_norm = norm_type == "layer_norm" + + if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None: + raise ValueError( + f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to" + f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}." + ) + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding( + dim, max_seq_length=num_positional_embeddings + ) + else: + self.pos_embed = None + + # Define 3 blocks. Each block has its own normalization layer. + # 1. Self-Attn + if self.use_ada_layer_norm: + self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm) + elif self.use_ada_layer_norm_zero: + self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm) + else: + self.norm1 = nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps + ) + + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim if only_cross_attention else None, + upcast_attention=upcast_attention, + ) + + # 2. Cross-Attn + if cross_attention_dim is not None or double_self_attention: + # We currently only use AdaLayerNormZero for self attention where there will only be one attention block. + # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during + # the second cross attention block. + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps + ) + ) + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim + if not double_self_attention + else None, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) # is self-attn if encoder_hidden_states is none + else: + self.norm2 = None + self.attn2 = None + + # 3. Feed-forward + if not self.use_ada_layer_norm_single: + self.norm3 = nn.LayerNorm( + dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps + ) + + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + ) + + # 4. Fuser + if attention_type == "gated" or attention_type == "gated-text-image": + self.fuser = GatedSelfAttentionDense( + dim, cross_attention_dim, num_attention_heads, attention_head_dim + ) + + # 5. Scale-shift for PixArt-Alpha. + if self.use_ada_layer_norm_single: + self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + + # let chunk size default to None + self._chunk_size = None + self._chunk_dim = 0 + + def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0): + # Sets chunk feed-forward + self._chunk_size = chunk_size + self._chunk_dim = dim + + def forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + # Notice that normalization is always applied before the real computation in the following blocks. + # 0. Self-Attention + batch_size = hidden_states.shape[0] + + if self.use_ada_layer_norm: + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1( + hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype + ) + elif self.use_layer_norm: + norm_hidden_states = self.norm1(hidden_states) + elif self.use_ada_layer_norm_single: + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1) + ).chunk(6, dim=1) + norm_hidden_states = self.norm1(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa + norm_hidden_states = norm_hidden_states.squeeze(1) + else: + raise ValueError("Incorrect norm used") + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + # 1. Retrieve lora scale. + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + + # 2. Prepare GLIGEN inputs + cross_attention_kwargs = ( + cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {} + ) + gligen_kwargs = cross_attention_kwargs.pop("gligen", None) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states + if self.only_cross_attention + else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + elif self.use_ada_layer_norm_single: + attn_output = gate_msa * attn_output + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + # 2.5 GLIGEN Control + if gligen_kwargs is not None: + hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"]) + + # 3. Cross-Attention + if self.attn2 is not None: + if self.use_ada_layer_norm: + norm_hidden_states = self.norm2(hidden_states, timestep) + elif self.use_ada_layer_norm_zero or self.use_layer_norm: + norm_hidden_states = self.norm2(hidden_states) + elif self.use_ada_layer_norm_single: + # For PixArt norm2 isn't applied here: + # https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L70C1-L76C103 + norm_hidden_states = hidden_states + else: + raise ValueError("Incorrect norm") + + if self.pos_embed is not None and self.use_ada_layer_norm_single is False: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 4. Feed-forward + if not self.use_ada_layer_norm_single: + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = ( + norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ) + + if self.use_ada_layer_norm_single: + norm_hidden_states = self.norm2(hidden_states) + norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp + + ff_output = self.ff(norm_hidden_states, scale=lora_scale) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + elif self.use_ada_layer_norm_single: + ff_output = gate_mlp * ff_output + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + return hidden_states + + +class TemporalBasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + attention_bias: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + ): + super().__init__() + self.only_cross_attention = only_cross_attention + self.use_ada_layer_norm = num_embeds_ada_norm is not None + self.unet_use_cross_frame_attention = unet_use_cross_frame_attention + self.unet_use_temporal_attention = unet_use_temporal_attention + + # SC-Attn + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + self.norm1 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm(dim) + ) + + # Cross-Attn + if cross_attention_dim is not None: + self.attn2 = Attention( + query_dim=dim, + cross_attention_dim=cross_attention_dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + else: + self.attn2 = None + + if cross_attention_dim is not None: + self.norm2 = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm(dim) + ) + else: + self.norm2 = None + + # Feed-forward + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn) + self.norm3 = nn.LayerNorm(dim) + self.use_ada_layer_norm_zero = False + + # Temp-Attn + assert unet_use_temporal_attention is not None + if unet_use_temporal_attention: + self.attn_temp = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + ) + nn.init.zeros_(self.attn_temp.to_out[0].weight.data) + self.norm_temp = ( + AdaLayerNorm(dim, num_embeds_ada_norm) + if self.use_ada_layer_norm + else nn.LayerNorm(dim) + ) + + def set_use_cross_frame_attention(self, value): + self.unet_use_cross_frame_attention = value + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + timestep=None, + attention_mask=None, + video_length=None, + reference=None, + ): + norm_hidden_states = ( + self.norm1(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm1(hidden_states) + ) + + if self.unet_use_cross_frame_attention: + hidden_states = ( + self.attn1( + norm_hidden_states, + attention_mask=attention_mask, + video_length=video_length, + ) + + hidden_states + ) + else: + if reference is not None: + bank_fea = [ + rearrange( + reference.unsqueeze(1).repeat(1, video_length, 1, 1), + "b t l c -> (b t) l c", + ) + ] + modify_norm_hidden_states = torch.cat( + [norm_hidden_states] + bank_fea, dim=1 + ) + + hidden_states = ( + self.attn1(norm_hidden_states, encoder_hidden_states=modify_norm_hidden_states, attention_mask=attention_mask) + + hidden_states + ) + else: + hidden_states = ( + self.attn1(norm_hidden_states, attention_mask=attention_mask) + + hidden_states + ) + + if self.attn2 is not None: + # Cross-Attention + norm_hidden_states = ( + self.norm2(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm2(hidden_states) + ) + hidden_states = ( + self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + ) + + hidden_states + ) + + # Feed-forward + hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states + + # Temporal-Attention + if self.unet_use_temporal_attention: + d = hidden_states.shape[1] + hidden_states = rearrange( + hidden_states, "(b f) d c -> (b d) f c", f=video_length + ) + norm_hidden_states = ( + self.norm_temp(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm_temp(hidden_states) + ) + hidden_states = self.attn_temp(norm_hidden_states) + hidden_states + hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d) + + return hidden_states diff --git a/src/scope/core/pipelines/personalive/modules/camera.py b/src/scope/core/pipelines/personalive/modules/camera.py new file mode 100644 index 00000000..a3dd9426 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/camera.py @@ -0,0 +1,73 @@ +# coding: utf-8 + +""" +functions for processing and transforming 3D facial keypoints +""" + +import numpy as np +import torch +import torch.nn.functional as F + +PI = np.pi + + +def headpose_pred_to_degree(pred): + """ + pred: (bs, 66) or (bs, 1) or others + """ + if pred.ndim > 1 and pred.shape[1] == 66: + # NOTE: note that the average is modified to 97.5 + device = pred.device + idx_tensor = [idx for idx in range(0, 66)] + idx_tensor = torch.FloatTensor(idx_tensor).to(device) + pred = F.softmax(pred, dim=1) + degree = torch.sum(pred*idx_tensor, axis=1) * 3 - 97.5 + + return degree + + return pred + + +def get_rotation_matrix(pitch_, yaw_, roll_): + """ the input is in degree + """ + # transform to radian + pitch = pitch_ / 180 * PI + yaw = yaw_ / 180 * PI + roll = roll_ / 180 * PI + + device = pitch.device + + if pitch.ndim == 1: + pitch = pitch.unsqueeze(1) + if yaw.ndim == 1: + yaw = yaw.unsqueeze(1) + if roll.ndim == 1: + roll = roll.unsqueeze(1) + + # calculate the euler matrix + bs = pitch.shape[0] + ones = torch.ones([bs, 1]).to(device) + zeros = torch.zeros([bs, 1]).to(device) + x, y, z = pitch, yaw, roll + + rot_x = torch.cat([ + ones, zeros, zeros, + zeros, torch.cos(x), -torch.sin(x), + zeros, torch.sin(x), torch.cos(x) + ], dim=1).reshape([bs, 3, 3]) + + rot_y = torch.cat([ + torch.cos(y), zeros, torch.sin(y), + zeros, ones, zeros, + -torch.sin(y), zeros, torch.cos(y) + ], dim=1).reshape([bs, 3, 3]) + + rot_z = torch.cat([ + torch.cos(z), -torch.sin(z), zeros, + torch.sin(z), torch.cos(z), zeros, + zeros, zeros, ones + ], dim=1).reshape([bs, 3, 3]) + + rot = rot_z @ rot_y @ rot_x + return rot.permute(0, 2, 1) # transpose diff --git a/src/scope/core/pipelines/personalive/modules/convnextv2.py b/src/scope/core/pipelines/personalive/modules/convnextv2.py new file mode 100644 index 00000000..910c4313 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/convnextv2.py @@ -0,0 +1,215 @@ +# coding: utf-8 + +""" +This moudle is adapted to the ConvNeXtV2 version for the extraction of implicit keypoints, poses, and expression deformation. +""" + +import torch +import torch.nn as nn +from .liveportrait_util import LayerNorm, DropPath, trunc_normal_, GRN +from einops import rearrange + +__all__ = ['convnextv2_tiny'] + + +class Block(nn.Module): + """ ConvNeXtV2 Block. + + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + """ + + def __init__(self, dim, drop_path=0.): + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.norm = LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(4 * dim) + self.pwconv2 = nn.Linear(4 * dim, dim) + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + x = input + self.drop_path(x) + return x + + +class ConvNeXtV2(nn.Module): + """ ConvNeXt V2 + + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + + def __init__( + self, + in_chans=3, + depths=[3, 3, 9, 3], + dims=[96, 192, 384, 768], + drop_path_rate=0., + **kwargs + ): + super().__init__() + self.depths = depths + self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format="channels_first") + ) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(4): + stage = nn.Sequential( + *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + + # NOTE: the output semantic items + num_bins = kwargs.get('num_bins', 66) + num_kp = kwargs.get('num_kp', 24) # the number of implicit keypoints + self.fc_kp = nn.Linear(dims[-1], 3 * num_kp) # implicit keypoints + + # print('dims[-1]: ', dims[-1]) + self.fc_scale = nn.Linear(dims[-1], 1) # scale + self.fc_pitch = nn.Linear(dims[-1], num_bins) # pitch bins + self.fc_yaw = nn.Linear(dims[-1], num_bins) # yaw bins + self.fc_roll = nn.Linear(dims[-1], num_bins) # roll bins + self.fc_t = nn.Linear(dims[-1], 3) # translation + self.fc_exp = nn.Linear(dims[-1], 3 * num_kp) # expression / delta + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x) + x = self.stages[i](x) + return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + + # implicit keypoints + kp = self.fc_kp(x) + + # pose and expression deformation + pitch = self.fc_pitch(x) + yaw = self.fc_yaw(x) + roll = self.fc_roll(x) + t = self.fc_t(x) + # exp = self.fc_exp(x) + scale = self.fc_scale(x) + + ret_dct = { + 'pitch': pitch, + 'yaw': yaw, + 'roll': roll, + 't': t, + # 'exp': exp, + 'scale': scale, + + 'kp': kp, # canonical keypoint + } + + return ret_dct + + +def convnextv2_tiny(**kwargs): + model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs) + return model + +class ConvNeXt(nn.Module): + """ ConvNeXt V2 + + Args: + in_chans (int): Number of input image channels. Default: 3 + num_classes (int): Number of classes for classification head. Default: 1000 + depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3] + dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768] + drop_path_rate (float): Stochastic depth rate. Default: 0. + head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1. + """ + + def __init__( + self, + in_chans=3, + depths=[3, 3, 9, 3], + dims=[96, 192, 384, 768], + drop_path_rate=0., + **kwargs + ): + super().__init__() + self.depths = depths + self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers + stem = nn.Sequential( + nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4), + LayerNorm(dims[0], eps=1e-6, data_format="channels_first") + ) + self.downsample_layers.append(stem) + for i in range(3): + downsample_layer = nn.Sequential( + LayerNorm(dims[i], eps=1e-6, data_format="channels_first"), + nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2), + ) + self.downsample_layers.append(downsample_layer) + + self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks + dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + cur = 0 + for i in range(4): + stage = nn.Sequential( + *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])] + ) + self.stages.append(stage) + cur += depths[i] + + self.norm = nn.LayerNorm(dims[-1], eps=1e-6) # final norm layer + + def _init_weights(self, m): + if isinstance(m, (nn.Conv2d, nn.Linear)): + trunc_normal_(m.weight, std=.02) + nn.init.constant_(m.bias, 0) + + def forward_features(self, x): + for i in range(4): + x = self.downsample_layers[i](x) + x = self.stages[i](x) + return self.norm(x.mean([-2, -1])) # global average pooling, (N, C, H, W) -> (N, C) + + def forward(self, x): + x = self.forward_features(x) + return x + + +def convnextv2(**kwargs): + model = ConvNeXt(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs) + return model diff --git a/src/scope/core/pipelines/personalive/modules/liveportrait_util.py b/src/scope/core/pipelines/personalive/modules/liveportrait_util.py new file mode 100644 index 00000000..a58ce587 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/liveportrait_util.py @@ -0,0 +1,492 @@ +# coding: utf-8 + +""" +This file defines various neural network modules and utility functions, including convolutional and residual blocks, +normalizations, and functions for spatial transformation and tensor manipulation. +""" + +from torch import nn +import torch.nn.functional as F +import torch +import torch.nn.utils.spectral_norm as spectral_norm +import math +import warnings +import collections.abc +from itertools import repeat + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp + + coordinate_grid = make_coordinate_grid(spatial_size, mean) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance) + + return out + + +def make_coordinate_grid(spatial_size, ref, **kwargs): + d, h, w = spatial_size + x = torch.arange(w).type(ref.dtype).to(ref.device) + y = torch.arange(h).type(ref.dtype).to(ref.device) + z = torch.arange(d).type(ref.dtype).to(ref.device) + + # NOTE: must be right-down-in + x = (2 * (x / (w - 1)) - 1) # the x axis faces to the right + y = (2 * (y / (h - 1)) - 1) # the y axis faces to the bottom + z = (2 * (z / (d - 1)) - 1) # the z axis faces to the inner + + yy = y.view(1, -1, 1).repeat(d, 1, w) + xx = x.view(1, 1, -1).repeat(d, h, 1) + zz = z.view(-1, 1, 1).repeat(1, h, w) + + meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3) + + return meshed + + +class ConvT2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1): + super(ConvT2d, self).__init__() + + self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride, + padding=padding, output_padding=output_padding) + self.norm = nn.InstanceNorm2d(out_features) + + def forward(self, x): + out = self.convT(x) + out = self.norm(out) + out = F.leaky_relu(out) + return out + + +class ResBlock3d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock3d, self).__init__() + self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding) + self.norm1 = nn.BatchNorm3d(in_features, affine=True) + self.norm2 = nn.BatchNorm3d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock3d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock3d, self).__init__() + + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=(1, 2, 2)) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class DownBlock3d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock3d, self).__init__() + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups, stride=(1, 2, 2)) + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = nn.BatchNorm3d(out_features, affine=True) + self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = nn.BatchNorm2d(out_features, affine=True) + if lrelu: + self.ac = nn.LeakyReLU() + else: + self.ac = nn.ReLU() + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = self.ac(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + self.out_filters = block_expansion + in_features + + self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1) + self.norm = nn.BatchNorm3d(self.out_filters, affine=True) + + def forward(self, x): + out = x.pop() + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class SPADE(nn.Module): + def __init__(self, norm_nc, label_nc): + super().__init__() + + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + nhidden = 128 + + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1), + nn.ReLU()) + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + + def forward(self, x, segmap): + normalized = self.param_free_norm(x) + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + actv = self.mlp_shared(segmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + out = normalized * (1 + gamma) + beta + return out + + +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1): + super().__init__() + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + self.use_se = use_se + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + # apply spectral norm if specified + if 'spectral' in norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + # define normalization layers + self.norm_0 = SPADE(fin, label_nc) + self.norm_1 = SPADE(fmiddle, label_nc) + if self.learned_shortcut: + self.norm_s = SPADE(fin, label_nc) + + def forward(self, x, seg1): + x_s = self.shortcut(x, seg1) + dx = self.conv_0(self.actvn(self.norm_0(x, seg1))) + dx = self.conv_1(self.actvn(self.norm_1(dx, seg1))) + out = x_s + dx + return out + + def shortcut(self, x, seg1): + if self.learned_shortcut: + x_s = self.conv_s(self.norm_s(x, seg1)) + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + + +def filter_state_dict(state_dict, remove_name='fc'): + new_state_dict = {} + for key in state_dict: + if remove_name in key: + continue + new_state_dict[key] = state_dict[key] + return new_state_dict + + +class GRN(nn.Module): + """ GRN (Global Response Normalization) layer + """ + + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +class LayerNorm(nn.Module): + r""" LayerNorm that supports two data formats: channels_last (default) or channels_first. + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + """ + + def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"): + super().__init__() + self.weight = nn.Parameter(torch.ones(normalized_shape, dtype=torch.float32)) + self.bias = nn.Parameter(torch.zeros(normalized_shape, dtype=torch.float32)) + self.eps = float(eps) + self.data_format = data_format + if self.data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + self.normalized_shape = (normalized_shape, ) + + def _apply(self, fn): + """ + 重写 _apply,完全接管参数的转换逻辑。 + 拦截所有 .cuda(), .cpu(), .half(), .to() 操作。 + """ + + for name, param in self._parameters.items(): + if param is not None: + dummy_probe = param.data.view(-1)[:1] + + try: + target_tensor = fn(dummy_probe) + + target_device = target_tensor.device + target_dtype = target_tensor.dtype + except: + target_device = param.device + target_dtype = param.dtype + + if name in ['weight', 'bias']: + # 核心逻辑:如果是 weight/bias,且目标是半精度,则强制保持 FP32 + if target_dtype in [torch.float16, torch.bfloat16]: + new_data = param.data.to(device=target_device, dtype=torch.float32) + else: + new_data = fn(param.data) + else: + new_data = fn(param.data) + + param.data = new_data + + if param.grad is not None: + param.grad.data = param.grad.data.to(device=new_data.device, dtype=new_data.dtype) + + for name, buf in self._buffers.items(): + if buf is not None: + self._buffers[name] = fn(buf) + + return self + + def forward(self, x): + dtype = x.dtype + x = x.float() + if self.data_format == "channels_last": + x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + elif self.data_format == "channels_first": + x = x.permute(0, 2, 3, 1) # BCHW → BHWC + x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps) + x = x.permute(0, 3, 1, 2) # BHWC → BCHW + return x.to(dtype) + + +def _no_grad_trunc_normal_(tensor, mean, std, a, b): + # Cut & paste from PyTorch official master until it's in a few official releases - RW + # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + def norm_cdf(x): + # Computes standard normal cumulative distribution function + return (1. + math.erf(x / math.sqrt(2.))) / 2. + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2) + + with torch.no_grad(): + # Values are generated by using a truncated uniform distribution and + # then using the inverse CDF for the normal distribution. + # Get upper and lower cdf values + l = norm_cdf((a - mean) / std) + u = norm_cdf((b - mean) / std) + + # Uniformly fill tensor with values from [l, u], then translate to + # [2l-1, 2u-1]. + tensor.uniform_(2 * l - 1, 2 * u - 1) + + # Use inverse cdf transform for normal distribution to get truncated + # standard normal + tensor.erfinv_() + + # Transform to proper mean, std + tensor.mul_(std * math.sqrt(2.)) + tensor.add_(mean) + + # Clamp to ensure it's in the proper range + tensor.clamp_(min=a, max=b) + return tensor + + +def drop_path(x, drop_prob=0., training=False, scale_by_keep=True): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + """ + if drop_prob == 0. or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets + random_tensor = x.new_empty(shape).bernoulli_(keep_prob) + if keep_prob > 0.0 and scale_by_keep: + random_tensor.div_(keep_prob) + return x * random_tensor + + +class DropPath(nn.Module): + """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + """ + + def __init__(self, drop_prob=None, scale_by_keep=True): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + self.scale_by_keep = scale_by_keep + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training, self.scale_by_keep) + + +def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.): + return _no_grad_trunc_normal_(tensor, mean, std, a, b) + +# From PyTorch internals +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + return parse + +to_2tuple = _ntuple(2) diff --git a/src/scope/core/pipelines/personalive/modules/motion_encoder.py b/src/scope/core/pipelines/personalive/modules/motion_encoder.py new file mode 100644 index 00000000..740b1adb --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/motion_encoder.py @@ -0,0 +1,42 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from .FAN_feature_extractor import FAN_SA +from einops import rearrange +from diffusers.models.embeddings import get_1d_sincos_pos_embed_from_grid +from diffusers.models.modeling_utils import ModelMixin + +def zero_module(module): + # Zero out the parameters of a module and return it. + assert isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear), type(module) + for p in module.parameters(): + p.detach().zero_() + return module + +class MotEncoder(ModelMixin): + def __init__(self, out_ch=16): + super().__init__() + self.model = FAN_SA() + self.out_drop = None #nn.Dropout(p=0.4) + self.out_ch = out_ch + expr_dim = 512 + # Use torch.arange and output_type='pt' for diffusers >= 0.34.0 compatibility + pos_grid = torch.arange(expr_dim // out_ch, dtype=torch.float32) + extra_pos_embed = get_1d_sincos_pos_embed_from_grid(out_ch, pos_grid, output_type='pt') + self.register_buffer("pe", extra_pos_embed.float().unsqueeze(0)) + self.final_proj = nn.Linear(expr_dim, expr_dim) + self.out_bn = None + + def change_out_dim(self, out_ch): + self.out_proj = nn.Linear(self.out_ch, out_ch) + + def set_attn_processor(self, processor): + self.model.set_attn_processor(processor) + + def forward(self, x): + x = x.to(self.dtype) + latent = self.model(rearrange(x, "b c f h w -> (b f) c h w")) + latent = self.final_proj(latent) + latent = rearrange(latent, "b (l c) -> b l c", c=self.out_ch) + self.pe + latent = rearrange(latent, "(b f) l c -> b f l c", f=x.shape[2]) + return latent diff --git a/src/scope/core/pipelines/personalive/modules/motion_extractor.py b/src/scope/core/pipelines/personalive/modules/motion_extractor.py new file mode 100644 index 00000000..bdd23080 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/motion_extractor.py @@ -0,0 +1,212 @@ +# coding: utf-8 + +""" +Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image +""" + +from torch import nn +import torch +from diffusers.models.modeling_utils import ModelMixin +from .convnextv2 import convnextv2_tiny +from .liveportrait_util import filter_state_dict +from .camera import headpose_pred_to_degree, get_rotation_matrix + +model_dict = { + 'convnextv2_tiny': convnextv2_tiny, +} + + +class MotionExtractor(ModelMixin): + def __init__(self, **kwargs): + super(MotionExtractor, self).__init__() + + # default is convnextv2_base + backbone = kwargs.get('backbone', 'convnextv2_tiny') + self.detector = model_dict.get(backbone)(**kwargs) + self.register_buffer('idx_tensor', torch.arange(66, dtype=torch.float32)) + + def headpose_pred_to_degree(self, pred): + """ + pred: (bs, 66) or (bs, 1) or others + """ + if pred.ndim > 1 and pred.shape[1] == 66: + # NOTE: note that the average is modified to 97.5 + prob = torch.nn.functional.softmax(pred, dim=1) + degree = torch.matmul(prob, self.idx_tensor) + degree = degree * 3 - 97.5 + + return degree + + return pred + + def load_pretrained(self, init_path: str): + if init_path not in (None, ''): + state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model'] + state_dict = filter_state_dict(state_dict, remove_name='head') + ret = self.detector.load_state_dict(state_dict, strict=False) + print(f'Load pretrained model from {init_path}, ret: {ret}') + + def forward(self, x): + kp_info = self.detector(x) + return self.get_kp(kp_info) + + def get_kp(self, kp_info): + bs = kp_info['kp'].shape[0] + + angles_raw = torch.cat([kp_info['pitch'], kp_info['yaw'], kp_info['roll']], dim=0) # (3, 66) + angles_deg = self.headpose_pred_to_degree(angles_raw)[:, None] # (B, 3) + pitch, yaw, roll = torch.chunk(angles_deg, chunks=3, dim=0) + + + kp = kp_info['kp'].reshape(bs, -1, 3) # BxNx3 + t, scale = kp_info['t'], kp_info['scale'] + + rot_mat = get_rotation_matrix(pitch, yaw, roll).to(self.dtype) # (bs, 3, 3) + + if kp.ndim == 2: + num_kp = kp.shape[1] // 3 # Bx(num_kpx3) + else: + num_kp = kp.shape[1] # Bxnum_kpx3 + + # Eqn.2: s * (R * x_c,s) + t + kp_transformed = kp.view(bs, num_kp, 3) @ rot_mat# + exp.view(bs, num_kp, 3) + kp_transformed *= scale[..., None] # (bs, k, 3) * (bs, 1, 1) = (bs, k, 3) + kp_transformed[:, :, 0:2] += t[:, None, 0:2] # remove z, only apply tx ty + + return kp_transformed + + def interpolate_tensors(self, a: torch.Tensor, b: torch.Tensor, num: int = 10) -> torch.Tensor: + if a.shape != b.shape: + raise ValueError(f"Shape mismatch: a.shape={a.shape}, b.shape={b.shape}") + + B, *rest = a.shape + alphas = torch.linspace(0, 1, num, device=a.device, dtype=a.dtype) + view_shape = (num,) + (1,) * len(rest) + alphas = alphas.view(view_shape) # (1, num, 1, 1, ...) + + result = (1 - alphas) * a + alphas * b + return result[:-1] + + def interpolate_kps(self, ref, motion, num_interp, t_scale=0.5, s_scale=0): + kp1 = self.detector(ref.to(self.dtype)) + kp2_list = [] + for i in range(0, motion.shape[0], 256): + motion_chunk = motion[i:i+256] + kp2_chunk = self.detector(motion_chunk.to(self.dtype)) + kp2_list.append(kp2_chunk) + kp2 = {} + for key in kp2_list[0].keys(): + kp2[key] = torch.cat([kp2_chunk[key] for kp2_chunk in kp2_list], dim=0) + + angles_raw = torch.cat([kp1['pitch'], kp1['yaw'], kp1['roll']], dim=0) # (3, 66) + angles_deg = self.headpose_pred_to_degree(angles_raw) # (B, 3) + pitch_1, yaw_1, roll_1 = torch.chunk(angles_deg, chunks=3, dim=0) + + angles_raw = torch.cat([kp2['pitch'], kp2['yaw'], kp2['roll']], dim=0) # (3, 66) + angles_deg = self.headpose_pred_to_degree(angles_raw) # (B, 3) + pitch_2, yaw_2, roll_2 = torch.chunk(angles_deg, chunks=3, dim=0) + + pitch_interp = self.interpolate_tensors(pitch_1, pitch_2[:1], num_interp) # Bx(num_interp)x1 + yaw_interp = self.interpolate_tensors(yaw_1, yaw_2[:1], num_interp) # Bx(num_interp)x1 + roll_interp = self.interpolate_tensors(roll_1, roll_2[:1], num_interp) # Bx(num_interp)x1 + + t_1 = kp1['t'] + t_2 = kp2['t'] + t_2 = (t_2 - t_2[0]) * t_scale + t_1 + t_interp = self.interpolate_tensors(t_1, t_2[:1], num_interp) + + s_1 = kp1['scale'] + s_2 = kp2['scale'] + s_2 = s_2 * s_scale + s_1 + s_interp = self.interpolate_tensors(s_1, s_2[:1], num_interp) + + kp = kp1['kp'].repeat(num_interp+motion.shape[0]-1, 1) + + kps_interp = { + 'pitch': torch.cat([pitch_interp, pitch_2], dim=0), + 'yaw': torch.cat([yaw_interp, yaw_2], dim=0), + 'roll': torch.cat([roll_interp, roll_2], dim=0), + 't': torch.cat([t_interp, t_2], dim=0), + 'scale': torch.cat([s_interp, s_2], dim=0), + 'kp': kp + } + + kp_intrep = self.get_kp(kps_interp) + + return kp_intrep + + + def interpolate_kps_online(self, ref, motion, num_interp, t_scale=0.5, s_scale=0): + kp1 = self.detector(ref.to(self.dtype)) + kp_frame1 = self.detector(motion[:1].to(self.dtype)) + kp2 = self.detector(motion.to(self.dtype)) + + angles_raw = torch.cat([kp1['pitch'], kp1['yaw'], kp1['roll']], dim=0) # (3, 66) + angles_deg = self.headpose_pred_to_degree(angles_raw) # (B, 3) + pitch_1, yaw_1, roll_1 = torch.chunk(angles_deg, chunks=3, dim=0) + + angles_raw = torch.cat([kp2['pitch'], kp2['yaw'], kp2['roll']], dim=0) # (3, 66) + angles_deg = self.headpose_pred_to_degree(angles_raw) # (B, 3) + pitch_2, yaw_2, roll_2 = torch.chunk(angles_deg, chunks=3, dim=0) + + pitch_interp = self.interpolate_tensors(pitch_1, pitch_2[:1], num_interp) # Bx(num_interp)x1 + yaw_interp = self.interpolate_tensors(yaw_1, yaw_2[:1], num_interp) # Bx(num_interp)x1 + roll_interp = self.interpolate_tensors(roll_1, roll_2[:1], num_interp) # Bx(num_interp)x1 + + t_1 = kp1['t'] + t_2 = kp2['t'] + t_2 = (t_2 - t_2[0]) * t_scale + t_1 + t_interp = self.interpolate_tensors(t_1, t_2[:1], num_interp) + + s_1 = kp1['scale'] + s_2 = kp2['scale'] + s_2 = s_2 * s_scale + s_1 + s_interp = self.interpolate_tensors(s_1, s_2[:1], num_interp) + + kp = kp1['kp'].repeat(num_interp+motion.shape[0]-1, 1) + + kps_interp = { + 'pitch': torch.cat([pitch_interp, pitch_2], dim=0), + 'yaw': torch.cat([yaw_interp, yaw_2], dim=0), + 'roll': torch.cat([roll_interp, roll_2], dim=0), + 't': torch.cat([t_interp, t_2], dim=0), + 'scale': torch.cat([s_interp, s_2], dim=0), + 'kp': kp + } + + kp_intrep = self.get_kp(kps_interp) + + kp_dri = self.get_kp(kp2) + + return kp_intrep, kp1, kp_frame1, kp_dri + + def get_kps(self, kp_ref, kp_frame1, motion, t_scale=0.5, s_scale=0): + kps_motion = self.detector(motion.to(self.dtype)) + + kps_dri = self.get_kp(kps_motion) + + t_ref = kp_ref['t'] + t_frame1 = kp_frame1['t'] + t_motion = kps_motion['t'] + kps_motion['t'] = (t_motion - t_frame1) * t_scale + t_ref + + s_ref = kp_ref['scale'] + s_motion = kps_motion['scale'] + kps_motion['scale'] = s_motion * s_scale + s_ref + + + kps_motion['kp'] = kp_ref['kp'].repeat(motion.shape[0], 1) + + kps_motion = self.get_kp(kps_motion) + + return kps_motion, kps_dri + + def inference(self, ref, motion): + kps_ref = self.detector(ref.to(self.dtype)) + kps_motion = self.detector(motion.to(self.dtype)) + kps_motion['kp'] = kps_ref['kp'] + + kp_s = self.get_kp(kps_ref) + kp_d = self.get_kp(kps_motion) + + return kp_s, kp_d diff --git a/src/scope/core/pipelines/personalive/modules/motion_module.py b/src/scope/core/pipelines/personalive/modules/motion_module.py new file mode 100644 index 00000000..f3f0ba74 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/motion_module.py @@ -0,0 +1,472 @@ +import math +from dataclasses import dataclass +from typing import Callable, Optional + +import torch +from diffusers.models.attention import FeedForward +from diffusers.models.attention_processor import Attention, AttnProcessor +from diffusers.utils import BaseOutput +from diffusers.utils.import_utils import is_xformers_available +from einops import rearrange, repeat +from torch import nn + + +def zero_module(module): + # Zero out the parameters of a module and return it. + assert isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear), type(module) + for p in module.parameters(): + p.detach().zero_() + return module + +def random_module(m): + assert isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear), type(m) + # Initialize weights with He initialization and zero out the biases + n = (m.kernel_size[0] * m.kernel_size[1] * m.in_channels) if isinstance(m, nn.Conv2d) else m.in_features + nn.init.normal_(m.weight, mean=0.0, std=math.sqrt(2. / n)) + if m.bias is not None: + nn.init.zeros_(m.bias) + return m + + +@dataclass +class TemporalTransformer3DModelOutput(BaseOutput): + sample: torch.FloatTensor + + +if is_xformers_available(): + import xformers + import xformers.ops +else: + xformers = None + + +def get_motion_module(in_channels, motion_module_type: str, motion_module_kwargs: dict): + if motion_module_type == "Vanilla": + return VanillaTemporalModule( + in_channels=in_channels, + **motion_module_kwargs, + ) + elif motion_module_type == "RefImage_Vanilla": + return VanillaTemporalModule( + in_channels=in_channels, + skip_ref_image=True, + **motion_module_kwargs, + ) + elif motion_module_type == "RefImageCond_Vanilla": + return VanillaTemporalModule( + in_channels=in_channels, + cond_ref_image=True, + **motion_module_kwargs, + ) + else: + raise ValueError + + +class VanillaTemporalModule(nn.Module): + + def __init__( + self, + in_channels, + num_attention_heads=8, + num_transformer_block=2, + attention_block_types=("Temporal_Self", "Temporal_Self"), + cross_attention_dim=768, + cross_frame_attention_mode=None, + temporal_position_encoding=False, + temporal_position_encoding_max_len=24, + temporal_attention_dim_div=1, + zero_initialize=True, + skip_ref_image=False, + cond_ref_image=False, + ): + super().__init__() + self.skip_ref_image = skip_ref_image + self.cond_ref_image = cond_ref_image + + self.temporal_transformer = TemporalTransformer3DModel( + in_channels=in_channels, + num_attention_heads=num_attention_heads, + attention_head_dim=in_channels + // num_attention_heads + // temporal_attention_dim_div, + num_layers=num_transformer_block, + attention_block_types=attention_block_types, + cross_attention_dim=cross_attention_dim, + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ) + + if zero_initialize: + self.temporal_transformer.proj_out = zero_module( + self.temporal_transformer.proj_out + ) + + def set_use_cross_frame_attention(self, value): + self.skip_ref_image = value + + def forward( + self, + input_tensor, + temb, + encoder_hidden_states, + attention_mask=None, + anchor_frame_idx=None, + debug=False + ): + hidden_states = input_tensor + if self.skip_ref_image: + # if input_tensor.shape[2] > 1: + hidden_states, ref_hidden_states = input_tensor[:, :, :-1], input_tensor[:, :, -1:] + + hidden_states = self.temporal_transformer( + hidden_states, encoder_hidden_states, attention_mask, debug=debug + ) + + output = hidden_states + if self.skip_ref_image: + # if input_tensor.shape[2] > 1: + output = torch.cat([output, ref_hidden_states], dim=2) + elif self.cond_ref_image: + output = torch.cat([output[:, :, :-1], input_tensor[:, :, -1:]], dim=2) + return output + + +class TemporalTransformer3DModel(nn.Module): + def __init__( + self, + in_channels, + num_attention_heads, + attention_head_dim, + num_layers, + attention_block_types=( + "Temporal_Self", + "Temporal_Self", + ), + dropout=0.0, + norm_num_groups=32, + cross_attention_dim=768, + activation_fn="geglu", + attention_bias=False, + upcast_attention=False, + cross_frame_attention_mode=None, + temporal_position_encoding=False, + temporal_position_encoding_max_len=24, + ): + super().__init__() + + inner_dim = num_attention_heads * attention_head_dim + + self.norm = torch.nn.GroupNorm( + num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) + self.proj_in = nn.Linear(in_channels, inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + TemporalTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + attention_block_types=attention_block_types, + dropout=dropout, + norm_num_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + attention_bias=attention_bias, + upcast_attention=upcast_attention, + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ) + for d in range(num_layers) + ] + ) + self.proj_out = nn.Linear(inner_dim, in_channels) + + def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None, debug=False): + assert ( + hidden_states.dim() == 5 + ), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}." + video_length = hidden_states.shape[2] + hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w") + + if encoder_hidden_states is not None and encoder_hidden_states.ndim == 4: + assert encoder_hidden_states.shape[1] == video_length, (video_length, encoder_hidden_states.shape) + encoder_hidden_states = rearrange(encoder_hidden_states, "b d n c -> (b d) n c",) + + batch, channel, height, weight = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * weight, inner_dim + ) + hidden_states = self.proj_in(hidden_states) + + # Transformer Blocks + for block in self.transformer_blocks: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + video_length=video_length, + ) + + # output + hidden_states = self.proj_out(hidden_states) + hidden_states = ( + hidden_states.reshape(batch, height, weight, inner_dim) + .permute(0, 3, 1, 2) + .contiguous() + ) + if False: + print( + 'TemporalModule', + hidden_states.shape, + # round(torch.abs(residual).mean().item(), 6), + # round(torch.abs(residual).max().item(), 6), + # round(torch.abs(hidden_states).mean().item(), 6), + # round(torch.abs(hidden_states).max().item(), 6), + ) + # hidden_states *= 0 + output = hidden_states + residual + output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length) + + return output + + +class TemporalTransformerBlock(nn.Module): + + def __init__( + self, + dim, + num_attention_heads, + attention_head_dim, + attention_block_types=( + "Temporal_Self", + "Temporal_Self", + ), + dropout=0.0, + norm_num_groups=32, + cross_attention_dim=768, + activation_fn="geglu", + attention_bias=False, + upcast_attention=False, + cross_frame_attention_mode=None, + temporal_position_encoding=False, + temporal_position_encoding_max_len=24, + proj_out_dim=None, + ): + super().__init__() + + attention_blocks = [] + norms = [] + + for block_name in attention_block_types: + attention_blocks.append( + VersatileAttention( + attention_mode=block_name.split("_")[0], + cross_attention_dim=cross_attention_dim + if block_name.endswith("_Cross") + else None, + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + upcast_attention=upcast_attention, + cross_frame_attention_mode=cross_frame_attention_mode, + temporal_position_encoding=temporal_position_encoding, + temporal_position_encoding_max_len=temporal_position_encoding_max_len, + ) + ) + norms.append(nn.LayerNorm(dim)) + + self.attention_blocks = nn.ModuleList(attention_blocks) + self.norms = nn.ModuleList(norms) + + self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn) + self.ff_norm = nn.LayerNorm(dim) + + self.proj_out = nn.Linear(dim, proj_out_dim) if proj_out_dim is not None else None + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + video_length=None, + att_flag=False + ): + for attention_block, norm in zip(self.attention_blocks, self.norms): + norm_hidden_states = norm(hidden_states) + if att_flag: + print( + 'block', + round(torch.abs(hidden_states).mean().item(), 6), + round(torch.abs(norm_hidden_states).mean().item(), 6), + ) + hidden_states = ( + attention_block( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states + if attention_block.is_cross_attention + else None, + video_length=video_length, + att_flag=att_flag + ) + + hidden_states + ) + + hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states + + output = hidden_states if self.proj_out is None else self.proj_out(hidden_states) + return output + + +class PositionalEncoding(nn.Module): + def __init__(self, d_model, dropout=0.0, max_len=24): + super().__init__() + self.dropout = nn.Dropout(p=dropout) + position = torch.arange(max_len).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model) + ) + pe = torch.zeros(1, max_len, d_model) + pe[0, :, 0::2] = torch.sin(position * div_term) + pe[0, :, 1::2] = torch.cos(position * div_term) + self.register_buffer("pe", pe) + + def forward(self, x): + x = x + self.pe[:, : x.size(1)] + return self.dropout(x) + + +class VersatileAttention(Attention): + def __init__( + self, + attention_mode=None, + cross_frame_attention_mode=None, + temporal_position_encoding=False, + temporal_position_encoding_max_len=24, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + assert attention_mode in ["Temporal", "Spatial"], attention_mode + + self.attention_mode = attention_mode + self.is_cross_attention = kwargs["cross_attention_dim"] is not None + + self.pos_encoder = ( + PositionalEncoding( + kwargs["query_dim"], + dropout=0.0, + max_len=temporal_position_encoding_max_len, + ) + if (temporal_position_encoding and attention_mode == "Temporal") + else None + ) + + def extra_repr(self): + return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}" + + def set_use_memory_efficient_attention_xformers( + self, + use_memory_efficient_attention_xformers: bool, + attention_op: Optional[Callable] = None, + ): + if use_memory_efficient_attention_xformers: + if not is_xformers_available(): + raise ModuleNotFoundError( + ( + "Refer to https://github.com/facebookresearch/xformers for more information on how to install" + " xformers" + ), + name="xformers", + ) + elif not torch.cuda.is_available(): + raise ValueError( + "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is" + " only available for GPU " + ) + else: + try: + # Make sure we can run the memory efficient attention + _ = xformers.ops.memory_efficient_attention( + torch.randn((1, 2, 40), device="cuda"), + torch.randn((1, 2, 40), device="cuda"), + torch.randn((1, 2, 40), device="cuda"), + ) + except Exception as e: + raise e + + # XFormersAttnProcessor corrupts video generation and work with Pytorch 1.13. + # Pytorch 2.0.1 AttnProcessor works the same as XFormersAttnProcessor in Pytorch 1.13. + # You don't need XFormersAttnProcessor here. + # processor = XFormersAttnProcessor( + # attention_op=attention_op, + # ) + processor = AttnProcessor() + else: + processor = AttnProcessor() + + self.set_processor(processor) + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + attention_mask=None, + video_length=None, + bank=None, + att_flag=False, + **cross_attention_kwargs, + ): + if self.attention_mode == "Temporal": + d = hidden_states.shape[1] # d means HxW + hidden_states = rearrange(hidden_states, "(b f) d c -> (b d) f c", f=video_length) + + if encoder_hidden_states is not None: + if not encoder_hidden_states.shape[0] == hidden_states.shape[0]: + encoder_hidden_states = repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d) + + if bank is not None and self.attention_mode == "Temporal" and not self.is_cross_attention: + # motion_frames作为之前的帧,引入motion module进行condition + modify_norm_hidden_states = torch.cat(bank + [hidden_states], dim=1) + + if self.pos_encoder is not None: + modify_norm_hidden_states = self.pos_encoder(modify_norm_hidden_states) + + hidden_states = self.processor( + self, + hidden_states, + encoder_hidden_states=modify_norm_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) # 改为cross-att + + else: + if self.pos_encoder is not None: + hidden_states = self.pos_encoder(hidden_states) + inp = hidden_states + hidden_states = self.processor( + self, + hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + if att_flag: + print( + 'ver_att', + round(torch.abs(inp).mean().item(), 6), + round(torch.abs(encoder_hidden_states).mean().item(), 6), + round(torch.abs(hidden_states).mean().item(), 6), + ) + + if self.attention_mode == "Temporal": + hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d) + + return hidden_states \ No newline at end of file diff --git a/src/scope/core/pipelines/personalive/modules/mutual_self_attention.py b/src/scope/core/pipelines/personalive/modules/mutual_self_attention.py new file mode 100644 index 00000000..a89a63cd --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/mutual_self_attention.py @@ -0,0 +1,463 @@ +from typing import Any, Dict, Optional + +import torch +from einops import rearrange + +from .attention import TemporalBasicTransformerBlock, BasicTransformerBlock + + +def torch_dfs(model: torch.nn.Module): + result = [model] + for child in model.children(): + result += torch_dfs(child) + return result + + +class ReferenceAttentionControl: + def __init__( + self, + unet, + mode="write", + do_classifier_free_guidance=False, + attention_auto_machine_weight=float("inf"), + gn_auto_machine_weight=1.0, + style_fidelity=1.0, + reference_attn=True, + reference_adain=False, + fusion_blocks="midup", + batch_size=1, + cache_kv=False, + ) -> None: + # 10. Modify self attention and group norm + self.unet = unet + assert mode in ["read", "write"] + assert fusion_blocks in ["midup", "full"] + self.reference_attn = reference_attn + self.reference_adain = reference_adain + self.fusion_blocks = fusion_blocks + self.cache_kv = cache_kv + self.register_reference_hooks( + mode, + do_classifier_free_guidance, + attention_auto_machine_weight, + gn_auto_machine_weight, + style_fidelity, + reference_attn, + reference_adain, + fusion_blocks, + batch_size=batch_size, + cache_kv=self.cache_kv, + ) + + def register_reference_hooks( + self, + mode, + do_classifier_free_guidance, + attention_auto_machine_weight, + gn_auto_machine_weight, + style_fidelity, + reference_attn, + reference_adain, + dtype=torch.float16, + batch_size=1, + num_images_per_prompt=1, + device=torch.device("cpu"), + fusion_blocks="midup", + cache_kv=False, + ): + MODE = mode + do_classifier_free_guidance = do_classifier_free_guidance + attention_auto_machine_weight = attention_auto_machine_weight + gn_auto_machine_weight = gn_auto_machine_weight + style_fidelity = style_fidelity + reference_attn = reference_attn + reference_adain = reference_adain + fusion_blocks = fusion_blocks + num_images_per_prompt = num_images_per_prompt + cache_kv = cache_kv + dtype = dtype + if do_classifier_free_guidance: + uc_mask = ( + torch.Tensor( + [1] * batch_size * num_images_per_prompt * 16 + + [0] * batch_size * num_images_per_prompt * 16 + ) + .to(device) + .bool() + ) + else: + uc_mask = ( + torch.Tensor([0] * batch_size * num_images_per_prompt * 2) + .to(device) + .bool() + ) + + def hacked_basic_transformer_inner_forward( + self, + hidden_states: torch.FloatTensor, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + timestep: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + class_labels: Optional[torch.LongTensor] = None, + video_length=None, + ): + if self.use_ada_layer_norm: # False + norm_hidden_states = self.norm1(hidden_states, timestep) + elif self.use_ada_layer_norm_zero: + ( + norm_hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + ) = self.norm1( + hidden_states, + timestep, + class_labels, + hidden_dtype=hidden_states.dtype, + ) + else: + norm_hidden_states = self.norm1(hidden_states) + + # 1. Self-Attention + # self.only_cross_attention = False + cross_attention_kwargs = ( + cross_attention_kwargs if cross_attention_kwargs is not None else {} + ) + if self.only_cross_attention: + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states + if self.only_cross_attention + else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + else: + if MODE == "write": + self.bank.append(norm_hidden_states.clone()) + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states + if self.only_cross_attention + else None, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + if MODE == "read": + kv_cache = norm_hidden_states.clone() + kv_cache = rearrange( + kv_cache, "(b t) l c -> b t l c", t=video_length + ) + self.kv_cache = kv_cache[:,:2,:,:] + + bank_fea = [ + rearrange( + d.unsqueeze(1).repeat(1, video_length, 1, 1), + "b t l c -> (b t) l c", + ) + for d in self.bank + ] + + if self.kv_bank is not None and cache_kv: + ahead_fea = self.kv_bank.unsqueeze(1).repeat(1, video_length, 1, 1, 1) + ahead_fea = rearrange(ahead_fea, "b t n l c -> (b t) (n l) c") + bank_fea.append(ahead_fea) + + modify_norm_hidden_states = torch.cat( + [norm_hidden_states] + bank_fea, dim=1 + ) + + hidden_states_uc = ( + self.attn1( + norm_hidden_states, + encoder_hidden_states=modify_norm_hidden_states, + attention_mask=attention_mask, + ) + # self.attn1( + # modify_norm_hidden_states, + # encoder_hidden_states=modify_norm_hidden_states, + # attention_mask=attention_mask, + # )[:, : hidden_states.shape[-2], :] + + hidden_states + ) + if do_classifier_free_guidance: + hidden_states_c = hidden_states_uc.clone() + _uc_mask = uc_mask.clone() + if hidden_states.shape[0] != _uc_mask.shape[0]: + _uc_mask = ( + torch.Tensor( + [1] * (hidden_states.shape[0] // 2) + + [0] * (hidden_states.shape[0] // 2) + ) + .to(device) + .bool() + ) + if self.kv_bank is not None: + # if False: + modify_norm_hidden_states = torch.cat( + [norm_hidden_states, ahead_fea], dim=1 + ) + else: + modify_norm_hidden_states = norm_hidden_states + hidden_states_c[_uc_mask] = ( + self.attn1( + norm_hidden_states[_uc_mask], + encoder_hidden_states=modify_norm_hidden_states[_uc_mask], + attention_mask=attention_mask, + ) + + hidden_states[_uc_mask] + ) + hidden_states = hidden_states_c.clone() + else: + hidden_states = hidden_states_uc + + # self.bank.clear() + if self.attn2 is not None: + # Cross-Attention + norm_hidden_states = ( + self.norm2(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm2(hidden_states) + ) + hidden_states = ( + self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + ) + + hidden_states + ) + + # Feed-forward + hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states + + # Temporal-Attention + if self.unet_use_temporal_attention: + d = hidden_states.shape[1] + hidden_states = rearrange( + hidden_states, "(b f) d c -> (b d) f c", f=video_length + ) + norm_hidden_states = ( + self.norm_temp(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm_temp(hidden_states) + ) + hidden_states = ( + self.attn_temp(norm_hidden_states) + hidden_states + ) + hidden_states = rearrange( + hidden_states, "(b d) f c -> (b f) d c", d=d + ) + + return hidden_states + + if self.use_ada_layer_norm_zero: + attn_output = gate_msa.unsqueeze(1) * attn_output + hidden_states = attn_output + hidden_states + + if self.attn2 is not None: + norm_hidden_states = ( + self.norm2(hidden_states, timestep) + if self.use_ada_layer_norm + else self.norm2(hidden_states) + ) + + # 2. Cross-Attention + attn_output = self.attn2( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + **cross_attention_kwargs, + ) + hidden_states = attn_output + hidden_states + + # 3. Feed-forward + norm_hidden_states = self.norm3(hidden_states) + + if self.use_ada_layer_norm_zero: + norm_hidden_states = ( + norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + ) + + ff_output = self.ff(norm_hidden_states) + + if self.use_ada_layer_norm_zero: + ff_output = gate_mlp.unsqueeze(1) * ff_output + + hidden_states = ff_output + hidden_states + + return hidden_states + + if self.reference_attn: + if self.fusion_blocks == "midup": + attn_modules = [ + module + for module in ( + torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks) + ) + if isinstance(module, BasicTransformerBlock) + or isinstance(module, TemporalBasicTransformerBlock) + ] + elif self.fusion_blocks == "full": + attn_modules = [ + module + for module in torch_dfs(self.unet) + if isinstance(module, BasicTransformerBlock) + or isinstance(module, TemporalBasicTransformerBlock) + ] + attn_modules = sorted( + attn_modules, key=lambda x: -x.norm1.normalized_shape[0] + ) + + for i, module in enumerate(attn_modules): + module._original_inner_forward = module.forward + if isinstance(module, BasicTransformerBlock): + module.forward = hacked_basic_transformer_inner_forward.__get__( + module, BasicTransformerBlock + ) + if isinstance(module, TemporalBasicTransformerBlock): + module.forward = hacked_basic_transformer_inner_forward.__get__( + module, TemporalBasicTransformerBlock + ) + + module.bank = [] + if(self.cache_kv): + module.kv_bank = None + module.kv_cache = None + module.attn_weight = float(i) / float(len(attn_modules)) + + def update(self, writer, dtype=torch.float16, drop_ratio=0.): + if self.reference_attn: + if self.fusion_blocks == "midup": + reader_attn_modules = [ + module + for module in ( + torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks) + ) + if isinstance(module, TemporalBasicTransformerBlock) + ] + writer_attn_modules = [ + module + for module in ( + torch_dfs(writer.unet.mid_block) + + torch_dfs(writer.unet.up_blocks) + ) + if isinstance(module, BasicTransformerBlock) + ] + elif self.fusion_blocks == "full": + reader_attn_modules = [ + module + for module in torch_dfs(self.unet) + if isinstance(module, TemporalBasicTransformerBlock) + ] + writer_attn_modules = [ + module + for module in torch_dfs(writer.unet) + if isinstance(module, BasicTransformerBlock) + ] + + reader_attn_modules = sorted( + reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0] + ) + writer_attn_modules = sorted( + writer_attn_modules, key=lambda x: -x.norm1.normalized_shape[0] + ) + for r, w in zip(reader_attn_modules, writer_attn_modules): + if drop_ratio > 0: + r.bank = [] + for v in w.bank: + N, L, D = v.shape # batch, length, dim + len_keep = int(L * (1 - drop_ratio)) + + noise = torch.rand(N, L) # noise in [0, 1] + ids_shuffle = torch.argsort(noise, dim=1) # ascend: small is keep, large is remove + ids_keep = ids_shuffle[:, :len_keep].to(v.device) + visible_tokens = torch.gather(v.clone(), dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D)) + r.bank.append((visible_tokens).to(dtype)) + + else: + r.bank = [v.clone().to(dtype) for v in w.bank] + # w.bank.clear() + + def update_hkf(self, writer, dtype=torch.float16, drop_ratio=0.): + if self.reference_attn: + if self.fusion_blocks == "midup": + reader_attn_modules = [ + module + for module in ( + torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks) + ) + if isinstance(module, TemporalBasicTransformerBlock) + ] + writer_attn_modules = [ + module + for module in ( + torch_dfs(writer.unet.mid_block) + + torch_dfs(writer.unet.up_blocks) + ) + if isinstance(module, BasicTransformerBlock) + ] + elif self.fusion_blocks == "full": + reader_attn_modules = [ + module + for module in torch_dfs(self.unet) + if isinstance(module, TemporalBasicTransformerBlock) + ] + writer_attn_modules = [ + module + for module in torch_dfs(writer.unet) + if isinstance(module, BasicTransformerBlock) + ] + + reader_attn_modules = sorted( + reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0] + ) + writer_attn_modules = sorted( + writer_attn_modules, key=lambda x: -x.norm1.normalized_shape[0] + ) + for r, w in zip(reader_attn_modules, writer_attn_modules): + if r.kv_bank is None: + r.kv_bank = torch.cat([v.clone().unsqueeze(1).to(dtype) for v in w.bank], dim=1) + else: + r.kv_bank = torch.cat([r.kv_bank] + [v.clone().unsqueeze(1).to(dtype) for v in w.bank], dim=1).to(dtype) + + def output(self, dtype=torch.float16): + res = {} + for i in range(3): + for j in range(2): + res[f"d{i}{j}"] = torch.cat([v.clone().to(dtype=dtype, device=self.unet.device) for v in self.unet.down_blocks[i].attentions[j].transformer_blocks[0].bank], dim=1) + res["m"] = torch.cat([v.clone().to(dtype=dtype, device=self.unet.device) for v in self.unet.mid_block.attentions[0].transformer_blocks[0].bank], dim=1) + for i in range(1, 4): + for j in range(3): + res[f"u{i}{j}"] = torch.cat([v.clone().to(dtype=dtype, device=self.unet.device) for v in self.unet.up_blocks[i].attentions[j].transformer_blocks[0].bank], dim=1) + return res + + def clear(self): + if self.reference_attn: + if self.fusion_blocks == "midup": + reader_attn_modules = [ + module + for module in ( + torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks) + ) + if isinstance(module, BasicTransformerBlock) + or isinstance(module, TemporalBasicTransformerBlock) + ] + elif self.fusion_blocks == "full": + reader_attn_modules = [ + module + for module in torch_dfs(self.unet) + if isinstance(module, BasicTransformerBlock) + or isinstance(module, TemporalBasicTransformerBlock) + ] + reader_attn_modules = sorted( + reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0] + ) + for r in reader_attn_modules: + r.bank.clear() + if self.cache_kv: + r.kv_bank=None + r.kv_cache=None diff --git a/src/scope/core/pipelines/personalive/modules/pose_guider.py b/src/scope/core/pipelines/personalive/modules/pose_guider.py new file mode 100644 index 00000000..445fc495 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/pose_guider.py @@ -0,0 +1,57 @@ +from typing import Tuple + +import torch.nn as nn +import torch.nn.functional as F +import torch.nn.init as init +from diffusers.models.modeling_utils import ModelMixin + +from .motion_module import zero_module +from .resnet import InflatedConv3d + + +class PoseGuider(ModelMixin): + def __init__( + self, + conditioning_embedding_channels: int = 320, + conditioning_channels: int = 3, + block_out_channels: Tuple[int] = (16, 32, 96, 256), + ): + super().__init__() + self.conv_in = InflatedConv3d( + conditioning_channels, block_out_channels[0], kernel_size=3, padding=1 + ) + + self.blocks = nn.ModuleList([]) + + for i in range(len(block_out_channels) - 1): + channel_in = block_out_channels[i] + channel_out = block_out_channels[i + 1] + self.blocks.append( + InflatedConv3d(channel_in, channel_in, kernel_size=3, padding=1) + ) + self.blocks.append( + InflatedConv3d( + channel_in, channel_out, kernel_size=3, padding=1, stride=2 + ) + ) + + self.conv_out_modify = zero_module( + InflatedConv3d( + block_out_channels[-1], + conditioning_embedding_channels, + kernel_size=3, + padding=1, + ) + ) + + def forward(self, conditioning): + embedding = self.conv_in(conditioning) + embedding = F.silu(embedding) + + for block in self.blocks: + embedding = block(embedding) + embedding = F.silu(embedding) + + embedding = self.conv_out_modify(embedding) + + return embedding diff --git a/src/scope/core/pipelines/personalive/modules/resnet.py b/src/scope/core/pipelines/personalive/modules/resnet.py new file mode 100644 index 00000000..ea731445 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/resnet.py @@ -0,0 +1,255 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from typing import Dict, Optional + + +class InflatedConv3d(nn.Conv2d): + def forward(self, x): + video_length = x.shape[2] + + x = rearrange(x, "b c f h w -> (b f) c h w") + x = super().forward(x) + x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length) + + return x + + +class InflatedGroupNorm(nn.GroupNorm): + def forward(self, x): + video_length = x.shape[2] + + x = rearrange(x, "b c f h w -> (b f) c h w") + x = super().forward(x) + x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length) + + return x + + +class Upsample3D(nn.Module): + def __init__( + self, + channels, + use_conv=False, + use_conv_transpose=False, + out_channels=None, + name="conv", + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.use_conv_transpose = use_conv_transpose + self.name = name + + conv = None + if use_conv_transpose: + raise NotImplementedError + elif use_conv: + self.conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1) + + def forward(self, hidden_states, output_size=None): + assert hidden_states.shape[1] == self.channels + + if self.use_conv_transpose: + raise NotImplementedError + + # Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16 + dtype = hidden_states.dtype + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(torch.float32) + + # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + hidden_states = hidden_states.contiguous() + + # if `output_size` is passed we force the interpolation output + # size and do not make use of `scale_factor=2` + if output_size is None: + hidden_states = F.interpolate( + hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest" + ) + else: + hidden_states = F.interpolate( + hidden_states, size=output_size, mode="nearest" + ) + + # If the input is bfloat16, we cast back to bfloat16 + if dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype) + + # if self.use_conv: + # if self.name == "conv": + # hidden_states = self.conv(hidden_states) + # else: + # hidden_states = self.Conv2d_0(hidden_states) + hidden_states = self.conv(hidden_states) + + return hidden_states + + +class Downsample3D(nn.Module): + def __init__( + self, channels, use_conv=False, out_channels=None, padding=1, name="conv" + ): + super().__init__() + self.channels = channels + self.out_channels = out_channels or channels + self.use_conv = use_conv + self.padding = padding + stride = 2 + self.name = name + + if use_conv: + self.conv = InflatedConv3d( + self.channels, self.out_channels, 3, stride=stride, padding=padding + ) + else: + raise NotImplementedError + + def forward(self, hidden_states): + assert hidden_states.shape[1] == self.channels + if self.use_conv and self.padding == 0: + raise NotImplementedError + + assert hidden_states.shape[1] == self.channels + hidden_states = self.conv(hidden_states) + + return hidden_states + + +class ResnetBlock3D(nn.Module): + def __init__( + self, + *, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout=0.0, + temb_channels=512, + groups=32, + groups_out=None, + pre_norm=True, + eps=1e-6, + non_linearity="swish", + time_embedding_norm="default", + output_scale_factor=1.0, + use_in_shortcut=None, + use_inflated_groupnorm=None, + ): + super().__init__() + self.pre_norm = pre_norm + self.pre_norm = True + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + self.time_embedding_norm = time_embedding_norm + self.output_scale_factor = output_scale_factor + + if groups_out is None: + groups_out = groups + + assert use_inflated_groupnorm != None + if use_inflated_groupnorm: + self.norm1 = InflatedGroupNorm( + num_groups=groups, num_channels=in_channels, eps=eps, affine=True + ) + else: + self.norm1 = torch.nn.GroupNorm( + num_groups=groups, num_channels=in_channels, eps=eps, affine=True + ) + + self.conv1 = InflatedConv3d( + in_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + + if temb_channels is not None: + if self.time_embedding_norm == "default": + time_emb_proj_out_channels = out_channels + elif self.time_embedding_norm == "scale_shift": + time_emb_proj_out_channels = out_channels * 2 + else: + raise ValueError( + f"unknown time_embedding_norm : {self.time_embedding_norm} " + ) + + self.time_emb_proj = torch.nn.Linear( + temb_channels, time_emb_proj_out_channels + ) + else: + self.time_emb_proj = None + + if use_inflated_groupnorm: + self.norm2 = InflatedGroupNorm( + num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True + ) + else: + self.norm2 = torch.nn.GroupNorm( + num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True + ) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = InflatedConv3d( + out_channels, out_channels, kernel_size=3, stride=1, padding=1 + ) + + if non_linearity == "swish": + self.nonlinearity = lambda x: F.silu(x) + elif non_linearity == "mish": + self.nonlinearity = Mish() + elif non_linearity == "silu": + self.nonlinearity = nn.SiLU() + + self.use_in_shortcut = ( + self.in_channels != self.out_channels + if use_in_shortcut is None + else use_in_shortcut + ) + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = InflatedConv3d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, input_tensor, temb): + hidden_states = input_tensor + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.conv1(hidden_states) + + if temb is not None: + video_length = hidden_states.shape[2] + temb = self.time_emb_proj(self.nonlinearity(temb)) + temb = rearrange(temb, "(b f) c -> b c f", f=video_length)[:, :, :, None, None] + + # if temb is not None: + # temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None] + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + hidden_states = self.norm2(hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + +class Mish(torch.nn.Module): + def forward(self, hidden_states): + return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states)) diff --git a/src/scope/core/pipelines/personalive/modules/scheduler_ddim.py b/src/scope/core/pipelines/personalive/modules/scheduler_ddim.py new file mode 100644 index 00000000..cc526d8f --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/scheduler_ddim.py @@ -0,0 +1,550 @@ +# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion +# and https://github.com/hojonathanho/diffusion + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union + +from einops import rearrange +import numpy as np +import torch + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.utils import BaseOutput +from diffusers.utils.torch_utils import randn_tensor +from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin + + +@dataclass +# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM +class DDIMSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + The predicted denoised sample `(x_{0})` based on the model output from the current timestep. + `pred_original_sample` can be used to preview progress or for guidance. + """ + + prev_sample: torch.FloatTensor + pred_original_sample: Optional[torch.FloatTensor] = None + + +# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar +def betas_for_alpha_bar( + num_diffusion_timesteps, + max_beta=0.999, + alpha_transform_type="cosine", +): + """ + Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of + (1-beta) over time from t = [0,1]. + + Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up + to that part of the diffusion process. + + + Args: + num_diffusion_timesteps (`int`): the number of betas to produce. + max_beta (`float`): the maximum beta to use; use values lower than 1 to + prevent singularities. + alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar. + Choose from `cosine` or `exp` + + Returns: + betas (`np.ndarray`): the betas used by the scheduler to step the model outputs + """ + if alpha_transform_type == "cosine": + + def alpha_bar_fn(t): + return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + + elif alpha_transform_type == "exp": + + def alpha_bar_fn(t): + return math.exp(t * -12.0) + + else: + raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}") + + betas = [] + for i in range(num_diffusion_timesteps): + t1 = i / num_diffusion_timesteps + t2 = (i + 1) / num_diffusion_timesteps + betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta)) + return torch.tensor(betas, dtype=torch.float32) + + +def rescale_zero_terminal_snr(betas): + """ + Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1) + + + Args: + betas (`torch.FloatTensor`): + the betas that the scheduler is being initialized with. + + Returns: + `torch.FloatTensor`: rescaled betas with zero terminal SNR + """ + # Convert betas to alphas_bar_sqrt + alphas = 1.0 - betas + alphas_cumprod = torch.cumprod(alphas, dim=0) + alphas_bar_sqrt = alphas_cumprod.sqrt() + + # Store old values. + alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone() + alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone() + + # Shift so the last timestep is zero. + alphas_bar_sqrt -= alphas_bar_sqrt_T + + # Scale so the first timestep is back to the old value. + alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T) + + # Convert alphas_bar_sqrt to betas + alphas_bar = alphas_bar_sqrt**2 # Revert sqrt + alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod + alphas = torch.cat([alphas_bar[0:1], alphas]) + betas = 1 - alphas + + return betas + + +class DDIMScheduler(SchedulerMixin, ConfigMixin): + """ + `DDIMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with + non-Markovian guidance. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + beta_start (`float`, defaults to 0.0001): + The starting `beta` value of inference. + beta_end (`float`, defaults to 0.02): + The final `beta` value. + beta_schedule (`str`, defaults to `"linear"`): + The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from + `linear`, `scaled_linear`, or `squaredcos_cap_v2`. + trained_betas (`np.ndarray`, *optional*): + Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`. + clip_sample (`bool`, defaults to `True`): + Clip the predicted sample for numerical stability. + clip_sample_range (`float`, defaults to 1.0): + The maximum magnitude for sample clipping. Valid only when `clip_sample=True`. + set_alpha_to_one (`bool`, defaults to `True`): + Each diffusion step uses the alphas product value at that step and at the previous one. For the final step + there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`, + otherwise it uses the alpha value at step 0. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps. You can use a combination of `offset=1` and + `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable + Diffusion. + prediction_type (`str`, defaults to `epsilon`, *optional*): + Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process), + `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen + Video](https://imagen.research.google/video/paper.pdf) paper). + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True`. + timestep_spacing (`str`, defaults to `"leading"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + rescale_betas_zero_snr (`bool`, defaults to `False`): + Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and + dark samples instead of limiting it to samples with medium brightness. Loosely related to + [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506). + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + beta_start: float = 0.0001, + beta_end: float = 0.02, + beta_schedule: str = "linear", + trained_betas: Optional[Union[np.ndarray, List[float]]] = None, + clip_sample: bool = True, + set_alpha_to_one: bool = True, + steps_offset: int = 0, + prediction_type: str = "epsilon", + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + clip_sample_range: float = 1.0, + sample_max_value: float = 1.0, + timestep_spacing: str = "leading", + rescale_betas_zero_snr: bool = False, + ): + if trained_betas is not None: + self.betas = torch.tensor(trained_betas, dtype=torch.float32) + elif beta_schedule == "linear": + self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32) + elif beta_schedule == "scaled_linear": + # this schedule is very specific to the latent diffusion model. + self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2 + elif beta_schedule == "squaredcos_cap_v2": + # Glide cosine schedule + self.betas = betas_for_alpha_bar(num_train_timesteps) + else: + raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}") + + # Rescale for zero SNR + if rescale_betas_zero_snr: + self.betas = rescale_zero_terminal_snr(self.betas) + + self.alphas = 1.0 - self.betas + self.alphas_cumprod = torch.cumprod(self.alphas, dim=0) + + # At every step in ddim, we are looking into the previous alphas_cumprod + # For the final step, there is no previous alphas_cumprod because we are already at 0 + # `set_alpha_to_one` decides whether we set this parameter simply to one or + # whether we use the final alpha of the "non-previous" one. + self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0] + + # standard deviation of the initial noise distribution + self.init_noise_sigma = 1.0 + + # setable values + self.num_inference_steps = None + self.step_length = None + self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)) + + def to(self, device): + self.betas = self.betas.to(device) + self.alphas = self.alphas.to(device) + self.alphas_cumprod = self.alphas_cumprod.to(device) + self.final_alpha_cumprod = self.final_alpha_cumprod.to(device) + + def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`int`, *optional*): + The current timestep in the diffusion chain. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + return sample + + def _get_variance(self, alpha_prod_t, alpha_prod_t_prev): + beta_prod_t = 1 - alpha_prod_t + beta_prod_t_prev = 1 - alpha_prod_t_prev + + variance = (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev) + + return variance + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device] = None, num_train_timesteps = None): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`): + The number of diffusion steps used when generating samples with a pre-trained model. + """ + if num_inference_steps > self.config.num_train_timesteps: + raise ValueError( + f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:" + f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle" + f" maximal {self.config.num_train_timesteps} timesteps." + ) + + self.num_inference_steps = num_inference_steps + + # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891 + if self.config.timestep_spacing == "linspace": + timesteps = ( + np.linspace(0, self.config.num_train_timesteps - 1, num_inference_steps) + .round()[::-1] + .copy() + .astype(np.int64) + ) + elif self.config.timestep_spacing == "leading": + step_ratio = self.config.num_train_timesteps // self.num_inference_steps + # creates integer timesteps by multiplying by ratio + # casting to int to avoid issues when num_inference_step is power of 3 + timesteps = (np.arange(0, num_inference_steps) * step_ratio).round()[::-1].copy().astype(np.int64) + timesteps += self.config.steps_offset + elif self.config.timestep_spacing == "trailing": + step_ratio = self.config.num_train_timesteps / self.num_inference_steps + # creates integer timesteps by multiplying by ratio + # casting to int to avoid issues when num_inference_step is power of 3 + timesteps = np.round(np.arange(self.config.num_train_timesteps, 0, -step_ratio)).astype(np.int64) + timesteps -= 1 + else: + raise ValueError( + f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'leading' or 'trailing'." + ) + + self.timesteps = torch.from_numpy(timesteps).to(device) + + def set_step_length(self, step_length): + self.step_length = step_length + + def step( + self, + model_output: torch.FloatTensor, + timestep: int, + sample: torch.FloatTensor, + eta: float = 0.0, + use_clipped_model_output: bool = False, + generator=None, + variance_noise: Optional[torch.FloatTensor] = None, + return_dict: bool = True, + ) -> Union[DDIMSchedulerOutput, Tuple]: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + eta (`float`): + The weight of noise for added noise in diffusion step. + use_clipped_model_output (`bool`, defaults to `False`): + If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary + because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no + clipping has happened, "corrected" `model_output` would coincide with the one provided as input and + `use_clipped_model_output` has no effect. + generator (`torch.Generator`, *optional*): + A random number generator. + variance_noise (`torch.FloatTensor`): + Alternative to generating noise with `generator` by directly providing the noise for the variance + itself. Useful for methods such as [`CycleDiffusion`]. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.DDIMSchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None and self.step_length is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf + # Ideally, read DDIM paper in-detail understanding + + # Notation ( -> + # - pred_noise_t -> e_theta(x_t, t) + # - pred_original_sample -> f_theta(x_t, t) or x_0 + # - std_dev_t -> sigma_t + # - eta -> η + # - pred_sample_direction -> "direction pointing to x_t" + # - pred_prev_sample -> "x_t-1" + + # 1. get previous step value (=t-1) + if self.step_length is not None: + prev_timestep = timestep - self.step_length + else: + prev_timestep = timestep - self.config.num_train_timesteps // self.num_inference_steps + + # 2. compute alphas, betas + alpha_prod_t = self.alphas_cumprod[timestep] + if isinstance(prev_timestep, torch.Tensor): + alpha_prod_t = alpha_prod_t.view(-1, 1, 1, 1) + prev_timestep = prev_timestep.to(self.alphas_cumprod.device) + alpha_prod_t_prev = torch.where(prev_timestep >= 0, self.alphas_cumprod[prev_timestep], self.final_alpha_cumprod) + alpha_prod_t_prev = alpha_prod_t_prev.view(-1, 1, 1, 1) + else: + alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod + + # alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod + + beta_prod_t = 1 - alpha_prod_t + + # 3. compute predicted original sample from predicted noise also called + # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + if self.config.prediction_type == "epsilon": + pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5) + pred_epsilon = model_output + elif self.config.prediction_type == "sample": + pred_original_sample = model_output + pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + elif self.config.prediction_type == "v_prediction": + pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output + pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or" + " `v_prediction`" + ) + + # 4. Clip or threshold "predicted x_0" + if self.config.thresholding: + pred_original_sample = self._threshold_sample(pred_original_sample) + elif self.config.clip_sample: + pred_original_sample = pred_original_sample.clamp( + -self.config.clip_sample_range, self.config.clip_sample_range + ) + + # 5. compute variance: "sigma_t(η)" -> see formula (16) + # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1) + variance = self._get_variance(alpha_prod_t, alpha_prod_t_prev) + std_dev_t = eta * variance ** (0.5) + + if use_clipped_model_output: + # the pred_epsilon is always re-derived from the clipped x_0 in Glide + pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5) + + # 6. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + pred_sample_direction = (1 - alpha_prod_t_prev - std_dev_t**2) ** (0.5) * pred_epsilon + + # 7. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf + prev_sample = alpha_prod_t_prev ** (0.5) * pred_original_sample + pred_sample_direction + + if eta > 0: + if variance_noise is not None and generator is not None: + raise ValueError( + "Cannot pass both generator and variance_noise. Please make sure that either `generator` or" + " `variance_noise` stays `None`." + ) + + if variance_noise is None: + variance_noise = randn_tensor( + model_output.shape, generator=generator, device=model_output.device, dtype=model_output.dtype + ) + variance = std_dev_t * variance_noise + + prev_sample = prev_sample + variance + + if not return_dict: + return (prev_sample, pred_original_sample) + + return DDIMSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample) + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise + def add_noise( + self, + original_samples: torch.FloatTensor, + noise: torch.FloatTensor, + timesteps: torch.IntTensor, + ) -> torch.FloatTensor: + # Make sure alphas_cumprod and timestep have same device and dtype as original_samples + # Move the self.alphas_cumprod to device to avoid redundant CPU to GPU data movement + # for the subsequent add_noise calls + shape_len = len(original_samples.shape) + if shape_len == 5: + b, c, f, h, w = original_samples.shape + original_samples = rearrange(original_samples, "b c f h w -> (b f) c h w") + noise = rearrange(noise, "b c f h w -> (b f) c h w") + + self.alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device) + alphas_cumprod = self.alphas_cumprod.to(dtype=original_samples.dtype) + timesteps = timesteps.to(original_samples.device) + + sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 + sqrt_alpha_prod = sqrt_alpha_prod.flatten() + while len(sqrt_alpha_prod.shape) < len(original_samples.shape): + sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) + + sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() + while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape): + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) + + noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise + + if shape_len == 5: + noisy_samples = rearrange(noisy_samples, "(b f) c h w -> b c f h w", b=b) + return noisy_samples + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity + def get_velocity( + self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor + ) -> torch.FloatTensor: + # Make sure alphas_cumprod and timestep have same device and dtype as sample + self.alphas_cumprod = self.alphas_cumprod.to(device=sample.device) + alphas_cumprod = self.alphas_cumprod.to(dtype=sample.dtype) + timesteps = timesteps.to(sample.device) + + sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5 + sqrt_alpha_prod = sqrt_alpha_prod.flatten() + while len(sqrt_alpha_prod.shape) < len(sample.shape): + sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1) + + sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5 + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten() + while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape): + sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1) + + velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample + return velocity + + def __len__(self): + return self.config.num_train_timesteps diff --git a/src/scope/core/pipelines/personalive/modules/transformer_2d.py b/src/scope/core/pipelines/personalive/modules/transformer_2d.py new file mode 100644 index 00000000..4c0995bf --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/transformer_2d.py @@ -0,0 +1,395 @@ +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.embeddings import PixArtAlphaTextProjection as CaptionProjection +from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear +from diffusers.models.modeling_utils import ModelMixin +from diffusers.models.normalization import AdaLayerNormSingle +from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, is_torch_version +from torch import nn + +from .attention import BasicTransformerBlock + + +@dataclass +class Transformer2DModelOutput(BaseOutput): + """ + The output of [`Transformer2DModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete): + The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability + distributions for the unnoised latent pixels. + """ + + sample: torch.FloatTensor + ref_feature: torch.FloatTensor + + +class Transformer2DModel(ModelMixin, ConfigMixin): + """ + A 2D Transformer model for image-like data. + + Parameters: + num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention. + attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head. + in_channels (`int`, *optional*): + The number of channels in the input and output (specify if the input is **continuous**). + num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use. + sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**). + This is fixed during training since it is used to learn a number of position embeddings. + num_vector_embeds (`int`, *optional*): + The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**). + Includes the class for the masked latent pixel. + activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward. + num_embeds_ada_norm ( `int`, *optional*): + The number of diffusion steps used during training. Pass if at least one of the norm_layers is + `AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are + added to the hidden states. + + During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`. + attention_bias (`bool`, *optional*): + Configure if the `TransformerBlocks` attention should contain a bias parameter. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + out_channels: Optional[int] = None, + num_layers: int = 1, + dropout: float = 0.0, + norm_num_groups: int = 32, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + sample_size: Optional[int] = None, + num_vector_embeds: Optional[int] = None, + patch_size: Optional[int] = None, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + double_self_attention: bool = False, + upcast_attention: bool = False, + norm_type: str = "layer_norm", + norm_elementwise_affine: bool = True, + norm_eps: float = 1e-5, + attention_type: str = "default", + caption_channels: int = None, + ): + super().__init__() + self.use_linear_projection = use_linear_projection + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + + conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv + linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear + + # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)` + # Define whether input is continuous or discrete depending on configuration + self.is_input_continuous = (in_channels is not None) and (patch_size is None) + self.is_input_vectorized = num_vector_embeds is not None + self.is_input_patches = in_channels is not None and patch_size is not None + + if norm_type == "layer_norm" and num_embeds_ada_norm is not None: + deprecation_message = ( + f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or" + " incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config." + " Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect" + " results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it" + " would be very nice if you could open a Pull request for the `transformer/config.json` file" + ) + deprecate( + "norm_type!=num_embeds_ada_norm", + "1.0.0", + deprecation_message, + standard_warn=False, + ) + norm_type = "ada_norm" + + if self.is_input_continuous and self.is_input_vectorized: + raise ValueError( + f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make" + " sure that either `in_channels` or `num_vector_embeds` is None." + ) + elif self.is_input_vectorized and self.is_input_patches: + raise ValueError( + f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make" + " sure that either `num_vector_embeds` or `num_patches` is None." + ) + elif ( + not self.is_input_continuous + and not self.is_input_vectorized + and not self.is_input_patches + ): + raise ValueError( + f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:" + f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None." + ) + + # 2. Define input layers + self.in_channels = in_channels + + self.norm = torch.nn.GroupNorm( + num_groups=norm_num_groups, + num_channels=in_channels, + eps=1e-6, + affine=True, + ) + if use_linear_projection: + self.proj_in = linear_cls(in_channels, inner_dim) + else: + self.proj_in = conv_cls( + in_channels, inner_dim, kernel_size=1, stride=1, padding=0 + ) + + # 3. Define transformers blocks + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + double_self_attention=double_self_attention, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + attention_type=attention_type, + ) + for d in range(num_layers) + ] + ) + + # 4. Define output layers + self.out_channels = in_channels if out_channels is None else out_channels + # TODO: should use out_channels for continuous projections + if use_linear_projection: + self.proj_out = linear_cls(inner_dim, in_channels) + else: + self.proj_out = conv_cls( + inner_dim, in_channels, kernel_size=1, stride=1, padding=0 + ) + + # 5. PixArt-Alpha blocks. + self.adaln_single = None + self.use_additional_conditions = False + if norm_type == "ada_norm_single": + self.use_additional_conditions = self.config.sample_size == 128 + # TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use + # additional conditions until we find better name + self.adaln_single = AdaLayerNormSingle( + inner_dim, use_additional_conditions=self.use_additional_conditions + ) + + self.caption_projection = None + if caption_channels is not None: + self.caption_projection = CaptionProjection( + in_features=caption_channels, hidden_size=inner_dim + ) + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + timestep: Optional[torch.LongTensor] = None, + added_cond_kwargs: Dict[str, torch.Tensor] = None, + class_labels: Optional[torch.LongTensor] = None, + cross_attention_kwargs: Dict[str, Any] = None, + attention_mask: Optional[torch.Tensor] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + ): + """ + The [`Transformer2DModel`] forward method. + + Args: + hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous): + Input `hidden_states`. + encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*): + Conditional embeddings for cross attention layer. If not given, cross-attention defaults to + self-attention. + timestep ( `torch.LongTensor`, *optional*): + Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`. + class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*): + Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in + `AdaLayerZeroNorm`. + cross_attention_kwargs ( `Dict[str, Any]`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + attention_mask ( `torch.Tensor`, *optional*): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + encoder_attention_mask ( `torch.Tensor`, *optional*): + Cross-attention mask applied to `encoder_hidden_states`. Two formats supported: + + * Mask `(batch, sequence_length)` True = keep, False = discard. + * Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard. + + If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format + above. This bias will be added to the cross-attention scores. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + + Returns: + If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a + `tuple` where the first element is the sample tensor. + """ + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension. + # we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward. + # we can tell by counting dims; if ndim == 2: it's a mask rather than a bias. + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None and attention_mask.ndim == 2: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2: + encoder_attention_mask = ( + 1 - encoder_attention_mask.to(hidden_states.dtype) + ) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # Retrieve lora scale. + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + + # 1. Input + batch, _, height, width = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + if not self.use_linear_projection: + hidden_states = ( + self.proj_in(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_in(hidden_states) + ) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * width, inner_dim + ) + else: + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * width, inner_dim + ) + hidden_states = ( + self.proj_in(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_in(hidden_states) + ) + + # 2. Blocks + if self.caption_projection is not None: + batch_size = hidden_states.shape[0] + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view( + batch_size, -1, hidden_states.shape[-1] + ) + + ref_feature = hidden_states.reshape(batch, height, width, inner_dim) + for block in self.transformer_blocks: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = ( + {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + hidden_states, + attention_mask, + encoder_hidden_states, + encoder_attention_mask, + timestep, + cross_attention_kwargs, + class_labels, + **ckpt_kwargs, + ) + else: + hidden_states = block( + hidden_states, + attention_mask=attention_mask, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + timestep=timestep, + cross_attention_kwargs=cross_attention_kwargs, + class_labels=class_labels, + ) + + # 3. Output + if self.is_input_continuous: + if not self.use_linear_projection: + hidden_states = ( + hidden_states.reshape(batch, height, width, inner_dim) + .permute(0, 3, 1, 2) + .contiguous() + ) + hidden_states = ( + self.proj_out(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_out(hidden_states) + ) + else: + hidden_states = ( + self.proj_out(hidden_states, scale=lora_scale) + if not USE_PEFT_BACKEND + else self.proj_out(hidden_states) + ) + hidden_states = ( + hidden_states.reshape(batch, height, width, inner_dim) + .permute(0, 3, 1, 2) + .contiguous() + ) + + output = hidden_states + residual + if not return_dict: + return (output, ref_feature) + + return Transformer2DModelOutput(sample=output, ref_feature=ref_feature) diff --git a/src/scope/core/pipelines/personalive/modules/transformer_3d.py b/src/scope/core/pipelines/personalive/modules/transformer_3d.py new file mode 100644 index 00000000..d10a3a48 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/transformer_3d.py @@ -0,0 +1,182 @@ +from dataclasses import dataclass +from typing import Optional, Dict + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models import ModelMixin +from diffusers.utils import BaseOutput +from diffusers.utils.import_utils import is_xformers_available +from einops import rearrange, repeat +from torch import nn + +from .attention import TemporalBasicTransformerBlock + + +@dataclass +class Transformer3DModelOutput(BaseOutput): + sample: torch.FloatTensor + + +if is_xformers_available(): + import xformers + import xformers.ops +else: + xformers = None + + +class Transformer3DModel(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 16, + attention_head_dim: int = 88, + in_channels: Optional[int] = None, + num_layers: int = 1, + dropout: float = 0.0, + norm_num_groups: int = 32, + cross_attention_dim: Optional[int] = None, + attention_bias: bool = False, + activation_fn: str = "geglu", + num_embeds_ada_norm: Optional[int] = None, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + ): + super().__init__() + self.use_linear_projection = use_linear_projection + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + inner_dim = num_attention_heads * attention_head_dim + + # Define input layers + self.in_channels = in_channels + + self.norm = torch.nn.GroupNorm( + num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True + ) + if use_linear_projection: + self.proj_in = nn.Linear(in_channels, inner_dim) + else: + self.proj_in = nn.Conv2d( + in_channels, inner_dim, kernel_size=1, stride=1, padding=0 + ) + + # Define transformers blocks + self.transformer_blocks = nn.ModuleList( + [ + TemporalBasicTransformerBlock( + inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + cross_attention_dim=cross_attention_dim, + activation_fn=activation_fn, + num_embeds_ada_norm=num_embeds_ada_norm, + attention_bias=attention_bias, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + ) + for d in range(num_layers) + ] + ) + + # 4. Define output layers + if use_linear_projection: + self.proj_out = nn.Linear(in_channels, inner_dim) + else: + self.proj_out = nn.Conv2d( + inner_dim, in_channels, kernel_size=1, stride=1, padding=0 + ) + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def forward( + self, + hidden_states, + encoder_hidden_states=None, + timestep=None, + return_dict: bool = True, + reference=None, + **block_kwargs, + ): + # Input + assert ( + hidden_states.dim() == 5 + ), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}." + video_length = hidden_states.shape[2] + hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w") + if encoder_hidden_states.shape[0] != hidden_states.shape[0]: + encoder_hidden_states = repeat( + encoder_hidden_states, "b n c -> (b f) n c", f=video_length + ) + + batch, channel, height, weight = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + if not self.use_linear_projection: + hidden_states = self.proj_in(hidden_states) + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * weight, inner_dim + ) + else: + inner_dim = hidden_states.shape[1] + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape( + batch, height * weight, inner_dim + ) + hidden_states = self.proj_in(hidden_states) + + # Blocks + for i, block in enumerate(self.transformer_blocks): + if reference is None: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + video_length=video_length, + **block_kwargs + ) + else: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + video_length=video_length, + reference=reference, + **block_kwargs + ) + + # Output + if not self.use_linear_projection: + hidden_states = ( + hidden_states.reshape(batch, height, weight, inner_dim) + .permute(0, 3, 1, 2) + .contiguous() + ) + hidden_states = self.proj_out(hidden_states) + else: + hidden_states = self.proj_out(hidden_states) + hidden_states = ( + hidden_states.reshape(batch, height, weight, inner_dim) + .permute(0, 3, 1, 2) + .contiguous() + ) + + output = hidden_states + residual + + output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length) + if not return_dict: + return (output,) + + return Transformer3DModelOutput(sample=output) diff --git a/src/scope/core/pipelines/personalive/modules/unet_2d_blocks.py b/src/scope/core/pipelines/personalive/modules/unet_2d_blocks.py new file mode 100644 index 00000000..99c7ed02 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/unet_2d_blocks.py @@ -0,0 +1,1073 @@ +from typing import Any, Dict, Optional, Tuple, Union + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import Attention +from diffusers.models import DualTransformer2DModel +from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D +from diffusers.utils import is_torch_version, logging +from diffusers.utils.torch_utils import apply_freeu +from torch import nn + +from .transformer_2d import Transformer2DModel + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def get_down_block( + down_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + temb_channels: int, + add_downsample: bool, + resnet_eps: float, + resnet_act_fn: str, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + downsample_padding: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + downsample_type: Optional[str] = None, + dropout: float = 0.0, +): + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + down_block_type = ( + down_block_type[7:] + if down_block_type.startswith("UNetRes") + else down_block_type + ) + if down_block_type == "DownBlock2D": + return DownBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif down_block_type == "CrossAttnDownBlock2D": + if cross_attention_dim is None: + raise ValueError( + "cross_attention_dim must be specified for CrossAttnDownBlock2D" + ) + return CrossAttnDownBlock2D( + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + dropout=dropout, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block( + up_block_type: str, + num_layers: int, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + add_upsample: bool, + resnet_eps: float, + resnet_act_fn: str, + resolution_idx: Optional[int] = None, + transformer_layers_per_block: int = 1, + num_attention_heads: Optional[int] = None, + resnet_groups: Optional[int] = None, + cross_attention_dim: Optional[int] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + attention_type: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: float = 1.0, + cross_attention_norm: Optional[str] = None, + attention_head_dim: Optional[int] = None, + upsample_type: Optional[str] = None, + dropout: float = 0.0, +) -> nn.Module: + # If attn head dim is not defined, we default it to the number of heads + if attention_head_dim is None: + logger.warn( + f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}." + ) + attention_head_dim = num_attention_heads + + up_block_type = ( + up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type + ) + if up_block_type == "UpBlock2D": + return UpBlock2D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + ) + elif up_block_type == "CrossAttnUpBlock2D": + if cross_attention_dim is None: + raise ValueError( + "cross_attention_dim must be specified for CrossAttnUpBlock2D" + ) + return CrossAttnUpBlock2D( + num_layers=num_layers, + transformer_layers_per_block=transformer_layers_per_block, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + resolution_idx=resolution_idx, + dropout=dropout, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + ) + + raise ValueError(f"{up_block_type} does not exist.") + + +class AutoencoderTinyBlock(nn.Module): + """ + Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU + blocks. + + Args: + in_channels (`int`): The number of input channels. + out_channels (`int`): The number of output channels. + act_fn (`str`): + ` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`. + + Returns: + `torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to + `out_channels`. + """ + + def __init__(self, in_channels: int, out_channels: int, act_fn: str): + super().__init__() + act_fn = get_activation(act_fn) + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), + act_fn, + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + act_fn, + nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), + ) + self.skip = ( + nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + if in_channels != out_channels + else nn.Identity() + ) + self.fuse = nn.ReLU() + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + return self.fuse(self.conv(x) + self.skip(x)) + + +class UNetMidBlock2D(nn.Module): + """ + A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks. + + Args: + in_channels (`int`): The number of input channels. + temb_channels (`int`): The number of temporal embedding channels. + dropout (`float`, *optional*, defaults to 0.0): The dropout rate. + num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. + resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. + resnet_time_scale_shift (`str`, *optional*, defaults to `default`): + The type of normalization to apply to the time embeddings. This can help to improve the performance of the + model on tasks with long-range temporal dependencies. + resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks. + resnet_groups (`int`, *optional*, defaults to 32): + The number of groups to use in the group normalization layers of the resnet blocks. + attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks. + resnet_pre_norm (`bool`, *optional*, defaults to `True`): + Whether to use pre-normalization for the resnet blocks. + add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks. + attention_head_dim (`int`, *optional*, defaults to 1): + Dimension of a single attention head. The number of attention heads is determined based on this value and + the number of input channels. + output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor. + + Returns: + `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size, + in_channels, height, width)`. + + """ + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + attn_groups: Optional[int] = None, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + ): + super().__init__() + resnet_groups = ( + resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + ) + self.add_attention = add_attention + + if attn_groups is None: + attn_groups = ( + resnet_groups if resnet_time_scale_shift == "default" else None + ) + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=attn_groups, + spatial_norm_dim=temb_channels + if resnet_time_scale_shift == "spatial" + else None, + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward( + self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None + ) -> torch.FloatTensor: + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = attn(hidden_states, temb=temb) + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class UNetMidBlock2DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + output_scale_factor: float = 1.0, + cross_attention_dim: int = 1280, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + upcast_attention: bool = False, + attention_type: str = "default", + ): + super().__init__() + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + resnet_groups = ( + resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + ) + + # support for variable transformer layers per block + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * num_layers + + # there is always at least one resnet + resnets = [ + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ] + attentions = [] + + for i in range(num_layers): + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=transformer_layers_per_block[i], + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + in_channels // num_attention_heads, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = ( + {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + ) + hidden_states, ref_feature = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + else: + hidden_states, ref_feature = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + + return hidden_states + + +class CrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + output_scale_factor: float = 1.0, + downsample_padding: int = 1, + add_downsample: bool = True, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + attention_type: str = "default", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * num_layers + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block[i], + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + additional_residuals: Optional[torch.FloatTensor] = None, + ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + output_states = () + + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + + blocks = list(zip(self.resnets, self.attentions)) + + for i, (resnet, attn) in enumerate(blocks): + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = ( + {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states, ref_feature = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states, ref_feature = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + + # apply additional residuals to the output of the last pair of resnet and attention blocks + if i == len(blocks) - 1 and additional_residuals is not None: + hidden_states = hidden_states + additional_residuals + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=lora_scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class DownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_padding: int = 1, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock2D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample2D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + scale: float = 1.0, + ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: + output_states = () + + for resnet in self.resnets: + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + output_states = output_states + (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, scale=scale) + + output_states = output_states + (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + transformer_layers_per_block: Union[int, Tuple[int]] = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + num_attention_heads: int = 1, + cross_attention_dim: int = 1280, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + only_cross_attention: bool = False, + upcast_attention: bool = False, + attention_type: str = "default", + ): + super().__init__() + resnets = [] + attentions = [] + + self.has_cross_attention = True + self.num_attention_heads = num_attention_heads + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * num_layers + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + if not dual_cross_attention: + attentions.append( + Transformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=transformer_layers_per_block[i], + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + ) + else: + attentions.append( + DualTransformer2DModel( + num_attention_heads, + out_channels // num_attention_heads, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + ) + ) + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [Upsample2D(out_channels, use_conv=True, out_channels=out_channels)] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + encoder_hidden_states: Optional[torch.FloatTensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + upsample_size: Optional[int] = None, + attention_mask: Optional[torch.FloatTensor] = None, + encoder_attention_mask: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + + for resnet, attn in zip(self.resnets, self.attentions): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + ckpt_kwargs: Dict[str, Any] = ( + {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + **ckpt_kwargs, + ) + hidden_states, ref_feature = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + else: + hidden_states = resnet(hidden_states, temb, scale=lora_scale) + hidden_states, ref_feature = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + return_dict=False, + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler( + hidden_states, upsample_size, scale=lora_scale + ) + + return hidden_states + + +class UpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + resolution_idx: Optional[int] = None, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock2D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [Upsample2D(out_channels, use_conv=True, out_channels=out_channels)] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + self.resolution_idx = resolution_idx + + def forward( + self, + hidden_states: torch.FloatTensor, + res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], + temb: Optional[torch.FloatTensor] = None, + upsample_size: Optional[int] = None, + scale: float = 1.0, + ) -> torch.FloatTensor: + is_freeu_enabled = ( + getattr(self, "s1", None) + and getattr(self, "s2", None) + and getattr(self, "b1", None) + and getattr(self, "b2", None) + ) + + for resnet in self.resnets: + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + + # FreeU: Only operate on the first two stages + if is_freeu_enabled: + hidden_states, res_hidden_states = apply_freeu( + self.resolution_idx, + hidden_states, + res_hidden_states, + s1=self.s1, + s2=self.s2, + b1=self.b1, + b2=self.b2, + ) + + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), + hidden_states, + temb, + use_reentrant=False, + ) + else: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = resnet(hidden_states, temb, scale=scale) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size, scale=scale) + + return hidden_states diff --git a/src/scope/core/pipelines/personalive/modules/unet_2d_condition.py b/src/scope/core/pipelines/personalive/modules/unet_2d_condition.py new file mode 100644 index 00000000..38969ba3 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/unet_2d_condition.py @@ -0,0 +1,1311 @@ +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import UNet2DConditionLoadersMixin +from diffusers.models.activations import get_activation +from diffusers.models.attention_processor import ( + ADDED_KV_ATTENTION_PROCESSORS, + CROSS_ATTENTION_PROCESSORS, + AttentionProcessor, + AttnAddedKVProcessor, + AttnProcessor, +) +from diffusers.models.embeddings import ( + GaussianFourierProjection, + ImageHintTimeEmbedding, + ImageProjection, + ImageTimeEmbedding, + TextImageProjection, + TextImageTimeEmbedding, + TextTimeEmbedding, + TimestepEmbedding, + Timesteps, +) +from diffusers.models.embeddings import GLIGENTextBoundingboxProjection as PositionNet +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import ( + USE_PEFT_BACKEND, + BaseOutput, + deprecate, + logging, + scale_lora_layers, + unscale_lora_layers, +) + +from .unet_2d_blocks import ( + UNetMidBlock2D, + UNetMidBlock2DCrossAttn, + get_down_block, + get_up_block, +) + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNet2DConditionOutput(BaseOutput): + """ + The output of [`UNet2DConditionModel`]. + + Args: + sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model. + """ + + sample: torch.FloatTensor = None + ref_features: Tuple[torch.FloatTensor] = None + + +class UNet2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin): + r""" + A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample + shaped output. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented + for all models (such as downloading or saving). + + Parameters: + sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`): + Height and width of input/output sample. + in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample. + out_channels (`int`, *optional*, defaults to 4): Number of channels in the output. + center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample. + flip_sin_to_cos (`bool`, *optional*, defaults to `False`): + Whether to flip the sin to cos in the time embedding. + freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding. + down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`): + The tuple of downsample blocks to use. + mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`): + Block type for middle of UNet, it can be one of `UNetMidBlock2DCrossAttn`, `UNetMidBlock2D`, or + `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped. + up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`): + The tuple of upsample blocks to use. + only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`): + Whether to include self-attention in the basic transformer blocks, see + [`~models.attention.BasicTransformerBlock`]. + block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`): + The tuple of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block. + downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution. + mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block. + dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use. + act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use. + norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization. + If `None`, normalization and activation layers is skipped in post-processing. + norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization. + cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280): + The dimension of the cross attention features. + transformer_layers_per_block (`int`, `Tuple[int]`, or `Tuple[Tuple]` , *optional*, defaults to 1): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + reverse_transformer_layers_per_block : (`Tuple[Tuple]`, *optional*, defaults to None): + The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`], in the upsampling + blocks of the U-Net. Only relevant if `transformer_layers_per_block` is of type `Tuple[Tuple]` and for + [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`], + [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`]. + encoder_hid_dim (`int`, *optional*, defaults to None): + If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim` + dimension to `cross_attention_dim`. + encoder_hid_dim_type (`str`, *optional*, defaults to `None`): + If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text + embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`. + attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads. + num_attention_heads (`int`, *optional*): + The number of attention heads. If not defined, defaults to `attention_head_dim` + resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config + for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`. + class_embed_type (`str`, *optional*, defaults to `None`): + The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`, + `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`. + addition_embed_type (`str`, *optional*, defaults to `None`): + Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or + "text". "text" will use the `TextTimeEmbedding` layer. + addition_time_embed_dim: (`int`, *optional*, defaults to `None`): + Dimension for the timestep embeddings. + num_class_embeds (`int`, *optional*, defaults to `None`): + Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing + class conditioning with `class_embed_type` equal to `None`. + time_embedding_type (`str`, *optional*, defaults to `positional`): + The type of position embedding to use for timesteps. Choose from `positional` or `fourier`. + time_embedding_dim (`int`, *optional*, defaults to `None`): + An optional override for the dimension of the projected time embedding. + time_embedding_act_fn (`str`, *optional*, defaults to `None`): + Optional activation function to use only once on the time embeddings before they are passed to the rest of + the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`. + timestep_post_act (`str`, *optional*, defaults to `None`): + The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`. + time_cond_proj_dim (`int`, *optional*, defaults to `None`): + The dimension of `cond_proj` layer in the timestep embedding. + conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer. conv_out_kernel (`int`, + *optional*, default to `3`): The kernel size of `conv_out` layer. projection_class_embeddings_input_dim (`int`, + *optional*): The dimension of the `class_labels` input when + `class_embed_type="projection"`. Required when `class_embed_type="projection"`. + class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time + embeddings with the class embeddings. + mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`): + Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If + `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the + `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False` + otherwise. + """ + + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "CrossAttnDownBlock2D", + "DownBlock2D", + ), + mid_block_type: Optional[str] = "UNetMidBlock2DCrossAttn", + up_block_types: Tuple[str] = ( + "UpBlock2D", + "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", + "CrossAttnUpBlock2D", + ), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: Union[int, Tuple[int]] = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + dropout: float = 0.0, + act_fn: str = "silu", + norm_num_groups: Optional[int] = 32, + norm_eps: float = 1e-5, + cross_attention_dim: Union[int, Tuple[int]] = 1280, + transformer_layers_per_block: Union[int, Tuple[int], Tuple[Tuple]] = 1, + reverse_transformer_layers_per_block: Optional[Tuple[Tuple[int]]] = None, + encoder_hid_dim: Optional[int] = None, + encoder_hid_dim_type: Optional[str] = None, + attention_head_dim: Union[int, Tuple[int]] = 8, + num_attention_heads: Optional[Union[int, Tuple[int]]] = None, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + addition_embed_type: Optional[str] = None, + addition_time_embed_dim: Optional[int] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + resnet_skip_time_act: bool = False, + resnet_out_scale_factor: int = 1.0, + time_embedding_type: str = "positional", + time_embedding_dim: Optional[int] = None, + time_embedding_act_fn: Optional[str] = None, + timestep_post_act: Optional[str] = None, + time_cond_proj_dim: Optional[int] = None, + conv_in_kernel: int = 3, + conv_out_kernel: int = 3, + projection_class_embeddings_input_dim: Optional[int] = None, + attention_type: str = "default", + class_embeddings_concat: bool = False, + mid_block_only_cross_attention: Optional[bool] = None, + cross_attention_norm: Optional[str] = None, + addition_embed_type_num_heads=64, + ): + super().__init__() + + self.sample_size = sample_size + + if num_attention_heads is not None: + raise ValueError( + "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19." + ) + + # If `num_attention_heads` is not defined (which is the case for most models) + # it will default to `attention_head_dim`. This looks weird upon first reading it and it is. + # The reason for this behavior is to correct for incorrectly named variables that were introduced + # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131 + # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking + # which is why we correct for the naming here. + num_attention_heads = num_attention_heads or attention_head_dim + + # Check inputs + if len(down_block_types) != len(up_block_types): + raise ValueError( + f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}." + ) + + if len(block_out_channels) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(only_cross_attention, bool) and len( + only_cross_attention + ) != len(down_block_types): + raise ValueError( + f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len( + down_block_types + ): + raise ValueError( + f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len( + down_block_types + ): + raise ValueError( + f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}." + ) + + if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len( + down_block_types + ): + raise ValueError( + f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}." + ) + + if not isinstance(layers_per_block, int) and len(layers_per_block) != len( + down_block_types + ): + raise ValueError( + f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}." + ) + if ( + isinstance(transformer_layers_per_block, list) + and reverse_transformer_layers_per_block is None + ): + for layer_number_per_block in transformer_layers_per_block: + if isinstance(layer_number_per_block, list): + raise ValueError( + "Must provide 'reverse_transformer_layers_per_block` if using asymmetrical UNet." + ) + + # input + conv_in_padding = (conv_in_kernel - 1) // 2 + self.conv_in = nn.Conv2d( + in_channels, + block_out_channels[0], + kernel_size=conv_in_kernel, + padding=conv_in_padding, + ) + + # time + if time_embedding_type == "fourier": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 2 + if time_embed_dim % 2 != 0: + raise ValueError( + f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}." + ) + self.time_proj = GaussianFourierProjection( + time_embed_dim // 2, + set_W_to_weight=False, + log=False, + flip_sin_to_cos=flip_sin_to_cos, + ) + timestep_input_dim = time_embed_dim + elif time_embedding_type == "positional": + time_embed_dim = time_embedding_dim or block_out_channels[0] * 4 + + self.time_proj = Timesteps( + block_out_channels[0], flip_sin_to_cos, freq_shift + ) + timestep_input_dim = block_out_channels[0] + else: + raise ValueError( + f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`." + ) + + self.time_embedding = TimestepEmbedding( + timestep_input_dim, + time_embed_dim, + act_fn=act_fn, + post_act_fn=timestep_post_act, + cond_proj_dim=time_cond_proj_dim, + ) + + if encoder_hid_dim_type is None and encoder_hid_dim is not None: + encoder_hid_dim_type = "text_proj" + self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type) + logger.info( + "encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined." + ) + + if encoder_hid_dim is None and encoder_hid_dim_type is not None: + raise ValueError( + f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}." + ) + + if encoder_hid_dim_type == "text_proj": + self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim) + elif encoder_hid_dim_type == "text_image_proj": + # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)` + self.encoder_hid_proj = TextImageProjection( + text_embed_dim=encoder_hid_dim, + image_embed_dim=cross_attention_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type == "image_proj": + # Kandinsky 2.2 + self.encoder_hid_proj = ImageProjection( + image_embed_dim=encoder_hid_dim, + cross_attention_dim=cross_attention_dim, + ) + elif encoder_hid_dim_type is not None: + raise ValueError( + f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'." + ) + else: + self.encoder_hid_proj = None + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding( + timestep_input_dim, time_embed_dim, act_fn=act_fn + ) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + elif class_embed_type == "projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set" + ) + # The projection `class_embed_type` is the same as the timestep `class_embed_type` except + # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings + # 2. it projects from an arbitrary input dimension. + # + # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations. + # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings. + # As a result, `TimestepEmbedding` can be passed arbitrary vectors. + self.class_embedding = TimestepEmbedding( + projection_class_embeddings_input_dim, time_embed_dim + ) + elif class_embed_type == "simple_projection": + if projection_class_embeddings_input_dim is None: + raise ValueError( + "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set" + ) + self.class_embedding = nn.Linear( + projection_class_embeddings_input_dim, time_embed_dim + ) + else: + self.class_embedding = None + + if addition_embed_type == "text": + if encoder_hid_dim is not None: + text_time_embedding_from_dim = encoder_hid_dim + else: + text_time_embedding_from_dim = cross_attention_dim + + self.add_embedding = TextTimeEmbedding( + text_time_embedding_from_dim, + time_embed_dim, + num_heads=addition_embed_type_num_heads, + ) + elif addition_embed_type == "text_image": + # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much + # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use + # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)` + self.add_embedding = TextImageTimeEmbedding( + text_embed_dim=cross_attention_dim, + image_embed_dim=cross_attention_dim, + time_embed_dim=time_embed_dim, + ) + elif addition_embed_type == "text_time": + self.add_time_proj = Timesteps( + addition_time_embed_dim, flip_sin_to_cos, freq_shift + ) + self.add_embedding = TimestepEmbedding( + projection_class_embeddings_input_dim, time_embed_dim + ) + elif addition_embed_type == "image": + # Kandinsky 2.2 + self.add_embedding = ImageTimeEmbedding( + image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim + ) + elif addition_embed_type == "image_hint": + # Kandinsky 2.2 ControlNet + self.add_embedding = ImageHintTimeEmbedding( + image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim + ) + elif addition_embed_type is not None: + raise ValueError( + f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'." + ) + + if time_embedding_act_fn is None: + self.time_embed_act = None + else: + self.time_embed_act = get_activation(time_embedding_act_fn) + + self.down_blocks = nn.ModuleList([]) + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = only_cross_attention + + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if mid_block_only_cross_attention is None: + mid_block_only_cross_attention = False + + if isinstance(num_attention_heads, int): + num_attention_heads = (num_attention_heads,) * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + if isinstance(cross_attention_dim, int): + cross_attention_dim = (cross_attention_dim,) * len(down_block_types) + + if isinstance(layers_per_block, int): + layers_per_block = [layers_per_block] * len(down_block_types) + + if isinstance(transformer_layers_per_block, int): + transformer_layers_per_block = [transformer_layers_per_block] * len( + down_block_types + ) + + if class_embeddings_concat: + # The time embeddings are concatenated with the class embeddings. The dimension of the + # time embeddings passed to the down, middle, and up blocks is twice the dimension of the + # regular time embeddings + blocks_time_embed_dim = time_embed_dim * 2 + else: + blocks_time_embed_dim = time_embed_dim + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block[i], + transformer_layers_per_block=transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + temb_channels=blocks_time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim[i], + num_attention_heads=num_attention_heads[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] + if attention_head_dim[i] is not None + else output_channel, + dropout=dropout, + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock2DCrossAttn": + self.mid_block = UNetMidBlock2DCrossAttn( + transformer_layers_per_block=transformer_layers_per_block[-1], + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + dropout=dropout, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim[-1], + num_attention_heads=num_attention_heads[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + attention_type=attention_type, + ) + elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn": + raise NotImplementedError(f"Unsupport mid_block_type: {mid_block_type}") + elif mid_block_type == "UNetMidBlock2D": + self.mid_block = UNetMidBlock2D( + in_channels=block_out_channels[-1], + temb_channels=blocks_time_embed_dim, + dropout=dropout, + num_layers=0, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_groups=norm_num_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + add_attention=False, + ) + elif mid_block_type is None: + self.mid_block = None + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the images + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_num_attention_heads = list(reversed(num_attention_heads)) + reversed_layers_per_block = list(reversed(layers_per_block)) + reversed_cross_attention_dim = list(reversed(cross_attention_dim)) + reversed_transformer_layers_per_block = ( + list(reversed(transformer_layers_per_block)) + if reverse_transformer_layers_per_block is None + else reverse_transformer_layers_per_block + ) + only_cross_attention = list(reversed(only_cross_attention)) + + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[ + min(i + 1, len(block_out_channels) - 1) + ] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=reversed_layers_per_block[i] + 1, + transformer_layers_per_block=reversed_transformer_layers_per_block[i], + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=blocks_time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resolution_idx=i, + resnet_groups=norm_num_groups, + cross_attention_dim=reversed_cross_attention_dim[i], + num_attention_heads=reversed_num_attention_heads[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + attention_type=attention_type, + resnet_skip_time_act=resnet_skip_time_act, + resnet_out_scale_factor=resnet_out_scale_factor, + cross_attention_norm=cross_attention_norm, + attention_head_dim=attention_head_dim[i] + if attention_head_dim[i] is not None + else output_channel, + dropout=dropout, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_num_groups is not None: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=norm_eps, + ) + + self.conv_act = get_activation(act_fn) + + else: + self.conv_norm_out = None + self.conv_act = None + self.conv_norm_out = None + + conv_out_padding = (conv_out_kernel - 1) // 2 + # self.conv_out = nn.Conv2d( + # block_out_channels[0], + # out_channels, + # kernel_size=conv_out_kernel, + # padding=conv_out_padding, + # ) + + if attention_type in ["gated", "gated-text-image"]: + positive_len = 768 + if isinstance(cross_attention_dim, int): + positive_len = cross_attention_dim + elif isinstance(cross_attention_dim, tuple) or isinstance( + cross_attention_dim, list + ): + positive_len = cross_attention_dim[0] + + feature_type = "text-only" if attention_type == "gated" else "text-image" + self.position_net = PositionNet( + positive_len=positive_len, + out_dim=cross_attention_dim, + feature_type=feature_type, + ) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "get_processor"): + processors[f"{name}.processor"] = module.get_processor( + return_deprecated_lora=True + ) + + for sub_name, child in module.named_children(): + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor( + self, + processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]], + _remove_lora=False, + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor, _remove_lora=_remove_lora) + else: + module.set_processor( + processor.pop(f"{name}.processor"), _remove_lora=_remove_lora + ) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def set_default_attn_processor(self): + """ + Disables custom attention processors and sets the default attention implementation. + """ + if all( + proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS + for proc in self.attn_processors.values() + ): + processor = AttnAddedKVProcessor() + elif all( + proc.__class__ in CROSS_ATTENTION_PROCESSORS + for proc in self.attn_processors.values() + ): + processor = AttnProcessor() + else: + raise ValueError( + f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}" + ) + + self.set_attn_processor(processor, _remove_lora=True) + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module splits the input tensor in slices to compute attention in + several steps. This is useful for saving some memory in exchange for a small decrease in speed. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, input to the attention heads is halved, so attention is computed in two steps. If + `"max"`, maximum amount of memory is saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_sliceable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_sliceable_dims(module) + + num_sliceable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_sliceable_layers * [1] + + slice_size = ( + num_sliceable_layers * [slice_size] + if not isinstance(slice_size, list) + else slice_size + ) + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice( + module: torch.nn.Module, slice_size: List[int] + ): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + def enable_freeu(self, s1, s2, b1, b2): + r"""Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497. + + The suffixes after the scaling factors represent the stage blocks where they are being applied. + + Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that + are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL. + + Args: + s1 (`float`): + Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to + mitigate the "oversmoothing effect" in the enhanced denoising process. + s2 (`float`): + Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to + mitigate the "oversmoothing effect" in the enhanced denoising process. + b1 (`float`): Scaling factor for stage 1 to amplify the contributions of backbone features. + b2 (`float`): Scaling factor for stage 2 to amplify the contributions of backbone features. + """ + for i, upsample_block in enumerate(self.up_blocks): + setattr(upsample_block, "s1", s1) + setattr(upsample_block, "s2", s2) + setattr(upsample_block, "b1", b1) + setattr(upsample_block, "b2", b2) + + def disable_freeu(self): + """Disables the FreeU mechanism.""" + freeu_keys = {"s1", "s2", "b1", "b2"} + for i, upsample_block in enumerate(self.up_blocks): + for k in freeu_keys: + if ( + hasattr(upsample_block, k) + or getattr(upsample_block, k, None) is not None + ): + setattr(upsample_block, k, None) + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + class_labels: Optional[torch.Tensor] = None, + timestep_cond: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + added_cond_kwargs: Optional[Dict[str, torch.Tensor]] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + down_intrablock_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + encoder_attention_mask: Optional[torch.Tensor] = None, + return_dict: bool = True, + only_return_middle_fea: bool = False, + ) -> Union[UNet2DConditionOutput, Tuple]: + r""" + The [`UNet2DConditionModel`] forward method. + + Args: + sample (`torch.FloatTensor`): + The noisy input tensor with the following shape `(batch, channel, height, width)`. + timestep (`torch.FloatTensor` or `float` or `int`): The number of timesteps to denoise an input. + encoder_hidden_states (`torch.FloatTensor`): + The encoder hidden states with shape `(batch, sequence_length, feature_dim)`. + class_labels (`torch.Tensor`, *optional*, defaults to `None`): + Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings. + timestep_cond: (`torch.Tensor`, *optional*, defaults to `None`): + Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed + through the `self.time_embedding` layer to obtain the timestep embeddings. + attention_mask (`torch.Tensor`, *optional*, defaults to `None`): + An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask + is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large + negative values to the attention scores corresponding to "discard" tokens. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under + `self.processor` in + [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals: (`tuple` of `torch.Tensor`, *optional*): + A tuple of tensors that if specified are added to the residuals of down unet blocks. + mid_block_additional_residual: (`torch.Tensor`, *optional*): + A tensor that if specified is added to the residual of the middle unet block. + encoder_attention_mask (`torch.Tensor`): + A cross-attention mask of shape `(batch, sequence_length)` is applied to `encoder_hidden_states`. If + `True` the mask is kept, otherwise if `False` it is discarded. Mask will be converted into a bias, + which adds large negative values to the attention scores corresponding to "discard" tokens. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain + tuple. + cross_attention_kwargs (`dict`, *optional*): + A kwargs dictionary that if specified is passed along to the [`AttnProcessor`]. + added_cond_kwargs: (`dict`, *optional*): + A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that + are passed along to the UNet blocks. + down_block_additional_residuals (`tuple` of `torch.Tensor`, *optional*): + additional residuals to be added to UNet long skip connections from down blocks to up blocks for + example from ControlNet side model(s) + mid_block_additional_residual (`torch.Tensor`, *optional*): + additional residual to be added to UNet mid block output, for example from ControlNet side model + down_intrablock_additional_residuals (`tuple` of `torch.Tensor`, *optional*): + additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s) + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + If `return_dict` is True, an [`~models.unet_2d_condition.UNet2DConditionOutput`] is returned, otherwise + a `tuple` is returned where the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layers). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + for dim in sample.shape[-2:]: + if dim % default_overall_up_factor != 0: + # Forward upsample size to force interpolation output size. + forward_upsample_size = True + break + + # ensure attention_mask is a bias, and give it a singleton query_tokens dimension + # expects mask of shape: + # [batch, key_tokens] + # adds singleton query_tokens dimension: + # [batch, 1, key_tokens] + # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes: + # [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn) + # [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn) + if attention_mask is not None: + # assume that mask is expressed as: + # (1 = keep, 0 = discard) + # convert mask into a bias that can be added to attention scores: + # (keep = +0, discard = -10000.0) + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # convert encoder_attention_mask to a bias the same way we do for attention_mask + if encoder_attention_mask is not None: + encoder_attention_mask = ( + 1 - encoder_attention_mask.to(sample.dtype) + ) * -10000.0 + encoder_attention_mask = encoder_attention_mask.unsqueeze(1) + + # 0. center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # 1. time + timesteps = timestep + if not torch.is_tensor(timesteps): + # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + timesteps = timesteps.expand(sample.shape[0]) + + t_emb = self.time_proj(timesteps) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=sample.dtype) + + emb = self.time_embedding(t_emb, timestep_cond) + aug_emb = None + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError( + "class_labels should be provided when num_class_embeds > 0" + ) + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + # `Timesteps` does not contain any weights and will always return f32 tensors + # there might be better ways to encapsulate this. + class_labels = class_labels.to(dtype=sample.dtype) + + class_emb = self.class_embedding(class_labels).to(dtype=sample.dtype) + + if self.config.class_embeddings_concat: + emb = torch.cat([emb, class_emb], dim=-1) + else: + emb = emb + class_emb + + if self.config.addition_embed_type == "text": + aug_emb = self.add_embedding(encoder_hidden_states) + elif self.config.addition_embed_type == "text_image": + # Kandinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + + image_embs = added_cond_kwargs.get("image_embeds") + text_embs = added_cond_kwargs.get("text_embeds", encoder_hidden_states) + aug_emb = self.add_embedding(text_embs, image_embs) + elif self.config.addition_embed_type == "text_time": + # SDXL - style + if "text_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `text_embeds` to be passed in `added_cond_kwargs`" + ) + text_embeds = added_cond_kwargs.get("text_embeds") + if "time_ids" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'text_time' which requires the keyword argument `time_ids` to be passed in `added_cond_kwargs`" + ) + time_ids = added_cond_kwargs.get("time_ids") + time_embeds = self.add_time_proj(time_ids.flatten()) + time_embeds = time_embeds.reshape((text_embeds.shape[0], -1)) + add_embeds = torch.concat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(emb.dtype) + aug_emb = self.add_embedding(add_embeds) + elif self.config.addition_embed_type == "image": + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image' which requires the keyword argument `image_embeds` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + aug_emb = self.add_embedding(image_embs) + elif self.config.addition_embed_type == "image_hint": + # Kandinsky 2.2 - style + if ( + "image_embeds" not in added_cond_kwargs + or "hint" not in added_cond_kwargs + ): + raise ValueError( + f"{self.__class__} has the config param `addition_embed_type` set to 'image_hint' which requires the keyword arguments `image_embeds` and `hint` to be passed in `added_cond_kwargs`" + ) + image_embs = added_cond_kwargs.get("image_embeds") + hint = added_cond_kwargs.get("hint") + aug_emb, hint = self.add_embedding(image_embs, hint) + sample = torch.cat([sample, hint], dim=1) + + emb = emb + aug_emb if aug_emb is not None else emb + + if self.time_embed_act is not None: + emb = self.time_embed_act(emb) + + if ( + self.encoder_hid_proj is not None + and self.config.encoder_hid_dim_type == "text_proj" + ): + encoder_hidden_states = self.encoder_hid_proj(encoder_hidden_states) + elif ( + self.encoder_hid_proj is not None + and self.config.encoder_hid_dim_type == "text_image_proj" + ): + # Kadinsky 2.1 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'text_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj( + encoder_hidden_states, image_embeds + ) + elif ( + self.encoder_hid_proj is not None + and self.config.encoder_hid_dim_type == "image_proj" + ): + # Kandinsky 2.2 - style + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + encoder_hidden_states = self.encoder_hid_proj(image_embeds) + elif ( + self.encoder_hid_proj is not None + and self.config.encoder_hid_dim_type == "ip_image_proj" + ): + if "image_embeds" not in added_cond_kwargs: + raise ValueError( + f"{self.__class__} has the config param `encoder_hid_dim_type` set to 'ip_image_proj' which requires the keyword argument `image_embeds` to be passed in `added_conditions`" + ) + image_embeds = added_cond_kwargs.get("image_embeds") + image_embeds = self.encoder_hid_proj(image_embeds).to( + encoder_hidden_states.dtype + ) + encoder_hidden_states = torch.cat( + [encoder_hidden_states, image_embeds], dim=1 + ) + + # 2. pre-process + sample = self.conv_in(sample) + + # 2.5 GLIGEN position net + if ( + cross_attention_kwargs is not None + and cross_attention_kwargs.get("gligen", None) is not None + ): + cross_attention_kwargs = cross_attention_kwargs.copy() + gligen_args = cross_attention_kwargs.pop("gligen") + cross_attention_kwargs["gligen"] = { + "objs": self.position_net(**gligen_args) + } + + # 3. down + lora_scale = ( + cross_attention_kwargs.get("scale", 1.0) + if cross_attention_kwargs is not None + else 1.0 + ) + if USE_PEFT_BACKEND: + # weight the lora layers by setting `lora_scale` for each PEFT layer + scale_lora_layers(self, lora_scale) + + is_controlnet = ( + mid_block_additional_residual is not None + and down_block_additional_residuals is not None + ) + # using new arg down_intrablock_additional_residuals for T2I-Adapters, to distinguish from controlnets + is_adapter = down_intrablock_additional_residuals is not None + # maintain backward compatibility for legacy usage, where + # T2I-Adapter and ControlNet both use down_block_additional_residuals arg + # but can only use one or the other + if ( + not is_adapter + and mid_block_additional_residual is None + and down_block_additional_residuals is not None + ): + deprecate( + "T2I should not use down_block_additional_residuals", + "1.3.0", + "Passing intrablock residual connections with `down_block_additional_residuals` is deprecated \ + and will be removed in diffusers 1.3.0. `down_block_additional_residuals` should only be used \ + for ControlNet. Please make sure use `down_intrablock_additional_residuals` instead. ", + standard_warn=False, + ) + down_intrablock_additional_residuals = down_block_additional_residuals + is_adapter = True + + down_block_res_samples = (sample,) + tot_referece_features = () + for downsample_block in self.down_blocks: + if ( + hasattr(downsample_block, "has_cross_attention") + and downsample_block.has_cross_attention + ): + # For t2i-adapter CrossAttnDownBlock2D + additional_residuals = {} + if is_adapter and len(down_intrablock_additional_residuals) > 0: + additional_residuals[ + "additional_residuals" + ] = down_intrablock_additional_residuals.pop(0) + + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + **additional_residuals, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, temb=emb, scale=lora_scale + ) + if is_adapter and len(down_intrablock_additional_residuals) > 0: + sample += down_intrablock_additional_residuals.pop(0) + + down_block_res_samples += res_samples + + if only_return_middle_fea: + return sample + + if is_controlnet: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = ( + down_block_res_sample + down_block_additional_residual + ) + new_down_block_res_samples = new_down_block_res_samples + ( + down_block_res_sample, + ) + + down_block_res_samples = new_down_block_res_samples + + # 4. mid + if self.mid_block is not None: + if ( + hasattr(self.mid_block, "has_cross_attention") + and self.mid_block.has_cross_attention + ): + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + cross_attention_kwargs=cross_attention_kwargs, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = self.mid_block(sample, emb) + + # To support T2I-Adapter-XL + if ( + is_adapter + and len(down_intrablock_additional_residuals) > 0 + and sample.shape == down_intrablock_additional_residuals[0].shape + ): + sample += down_intrablock_additional_residuals.pop(0) + + if is_controlnet: + sample = sample + mid_block_additional_residual + + # 5. up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[ + : -len(upsample_block.resnets) + ] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if ( + hasattr(upsample_block, "has_cross_attention") + and upsample_block.has_cross_attention + ): + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + cross_attention_kwargs=cross_attention_kwargs, + upsample_size=upsample_size, + attention_mask=attention_mask, + encoder_attention_mask=encoder_attention_mask, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + scale=lora_scale, + ) + + # 6. post-process + # if self.conv_norm_out: + # sample = self.conv_norm_out(sample) + # sample = self.conv_act(sample) + # sample = self.conv_out(sample) + + if USE_PEFT_BACKEND: + # remove `lora_scale` from each PEFT layer + unscale_lora_layers(self, lora_scale) + + if not return_dict: + return (sample,) + + return UNet2DConditionOutput(sample=sample) diff --git a/src/scope/core/pipelines/personalive/modules/unet_3d.py b/src/scope/core/pipelines/personalive/modules/unet_3d.py new file mode 100644 index 00000000..fa9df12f --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/unet_3d.py @@ -0,0 +1,717 @@ +from collections import OrderedDict +from dataclasses import dataclass +import pdb +from os import PathLike +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint +import torch.nn.functional as F +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention_processor import AttentionProcessor +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, BaseOutput, logging +from safetensors.torch import load_file + +from .resnet import InflatedConv3d, InflatedGroupNorm +from .unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block +from .motion_module import get_motion_module + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class UNet3DConditionOutput(BaseOutput): + sample: torch.FloatTensor + + +class UNet3DConditionModel(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "DownBlock3D", + ), + mid_block_type: str = "UNetMidBlock3DCrossAttn", + up_block_types: Tuple[str] = ( + "UpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + ), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-5, + cross_attention_dim: int = 1280, + attention_head_dim: Union[int, Tuple[int]] = 8, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + use_inflated_groupnorm=False, + # Additional + use_motion_module=False, + use_temporal_module=False, + motion_module_resolutions=(1, 2, 4, 8), + motion_module_mid_block=False, + motion_module_decoder_only=False, + motion_module_type=None, + temporal_module_type=None, + motion_module_kwargs={}, + temporal_module_kwargs={}, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + ): + super().__init__() + + self.sample_size = sample_size + time_embed_dim = block_out_channels[0] * 4 + + # input + self.conv_in = InflatedConv3d( + in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1) + ) + + # time + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + + self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + else: + self.class_embedding = None + + self.down_blocks = nn.ModuleList([]) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + res = 2**i + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + temb_channels=time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module + and (res in motion_module_resolutions) + and (not motion_module_decoder_only), + use_temporal_module=use_temporal_module + and (res in motion_module_resolutions) + and (not motion_module_decoder_only), + motion_module_type=motion_module_type, + temporal_module_type=temporal_module_type, + motion_module_kwargs=motion_module_kwargs, + temporal_module_kwargs=temporal_module_kwargs + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock3DCrossAttn": + self.mid_block = UNetMidBlock3DCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module and motion_module_mid_block, + use_temporal_module=use_temporal_module and motion_module_mid_block, + motion_module_type=motion_module_type, + temporal_module_type=temporal_module_type, + motion_module_kwargs=motion_module_kwargs, + temporal_module_kwargs=temporal_module_kwargs, + ) + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the videos + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_attention_head_dim = list(reversed(attention_head_dim)) + only_cross_attention = list(reversed(only_cross_attention)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + res = 2 ** (3 - i) + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[ + min(i + 1, len(block_out_channels) - 1) + ] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=reversed_attention_head_dim[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module + and (res in motion_module_resolutions), + use_temporal_module=use_temporal_module + and (res in motion_module_resolutions), + motion_module_type=motion_module_type, + temporal_module_type=temporal_module_type, + motion_module_kwargs=motion_module_kwargs, + temporal_module_kwargs=temporal_module_kwargs, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if use_inflated_groupnorm: + self.conv_norm_out = InflatedGroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=norm_eps, + ) + else: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=norm_eps, + ) + self.conv_act = nn.SiLU() + self.conv_out = InflatedConv3d( + block_out_channels[0], out_channels, kernel_size=3, padding=1 + ) + + @property + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors + def attn_processors(self) -> Dict[str, AttentionProcessor]: + r""" + Returns: + `dict` of attention processors: A dictionary containing all attention processors used in the model with + indexed by its weight name. + """ + # set recursively + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "set_processor"): + processors[f"{name}.processor"] = module.processor + + for sub_name, child in module.named_children(): + if "temporal_transformer" not in sub_name: + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + if "temporal_transformer" not in name: + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attention_slice(self, slice_size): + r""" + Enable sliced attention computation. + + When this option is enabled, the attention module will split the input tensor in slices, to compute attention + in several steps. This is useful to save some memory in exchange for a small speed decrease. + + Args: + slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`): + When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If + `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is + provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim` + must be a multiple of `slice_size`. + """ + sliceable_head_dims = [] + + def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module): + if hasattr(module, "set_attention_slice"): + sliceable_head_dims.append(module.sliceable_head_dim) + + for child in module.children(): + fn_recursive_retrieve_slicable_dims(child) + + # retrieve number of attention layers + for module in self.children(): + fn_recursive_retrieve_slicable_dims(module) + + num_slicable_layers = len(sliceable_head_dims) + + if slice_size == "auto": + # half the attention head size is usually a good trade-off between + # speed and memory + slice_size = [dim // 2 for dim in sliceable_head_dims] + elif slice_size == "max": + # make smallest slice possible + slice_size = num_slicable_layers * [1] + + slice_size = ( + num_slicable_layers * [slice_size] + if not isinstance(slice_size, list) + else slice_size + ) + + if len(slice_size) != len(sliceable_head_dims): + raise ValueError( + f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different" + f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}." + ) + + for i in range(len(slice_size)): + size = slice_size[i] + dim = sliceable_head_dims[i] + if size is not None and size > dim: + raise ValueError(f"size {size} has to be smaller or equal to {dim}.") + + # Recursively walk through all the children. + # Any children which exposes the set_attention_slice method + # gets the message + def fn_recursive_set_attention_slice( + module: torch.nn.Module, slice_size: List[int] + ): + if hasattr(module, "set_attention_slice"): + module.set_attention_slice(slice_size.pop()) + + for child in module.children(): + fn_recursive_set_attention_slice(child, slice_size) + + reversed_slice_size = list(reversed(slice_size)) + for module in self.children(): + fn_recursive_set_attention_slice(module, reversed_slice_size) + + def set_use_cross_frame_attention(self, value): + + def fn_recursive_set_use_cf_att(module: torch.nn.Module, value): + if hasattr(module, "set_use_cross_frame_attention"): + module.set_use_cross_frame_attention(value) + + for child in module.children(): + fn_recursive_set_use_cf_att(child, value) + + for module in self.children(): + fn_recursive_set_use_cf_att(module, value) + + def _set_gradient_checkpointing(self, module, value=False): + if hasattr(module, "gradient_checkpointing"): + module.gradient_checkpointing = value + + # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor + def set_attn_processor( + self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] + ): + r""" + Sets the attention processor to use to compute attention. + + Parameters: + processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`): + The instantiated processor class or a dictionary of processor classes that will be set as the processor + for **all** `Attention` layers. + + If `processor` is a dict, the key needs to define the path to the corresponding cross attention + processor. This is strongly recommended when setting trainable attention processors. + + """ + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + class_labels: Optional[torch.Tensor] = None, + pose_cond_fea = None, + attention_mask: Optional[torch.Tensor] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + return_dict: bool = True, + skip_mm: bool = False, + ) -> Union[UNet3DConditionOutput, Tuple]: + r""" + Args: + sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor + timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps + encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple. + + Returns: + [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`: + [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When + returning a tuple, the first element is the sample tensor. + """ + # By default samples have to be AT least a multiple of the overall upsampling factor. + # The overall upsampling factor is equal to 2 ** (# num of upsampling layears). + # However, the upsampling interpolation output size can be forced to fit any upsampling size + # on the fly if necessary. + default_overall_up_factor = 2**self.num_upsamplers + + # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor` + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + logger.info("Forward upsample size to force interpolation output size.") + forward_upsample_size = True + + # prepare attention_mask + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + # center input if necessary + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # time + timesteps = timestep + if not torch.is_tensor(timesteps): + # This would be a good case for the `match` statement (Python 3.10+) + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + # broadcast to batch dimension in a way that's compatible with ONNX/Core ML + # try: + # timesteps = timesteps.expand(sample.shape[0]) + # except Exception: + # timesteps = timesteps.expand(sample.shape[0] * sample.shape[2]) + + t_emb = self.time_proj(timesteps) + + # timesteps does not contain any weights and will always return f32 tensors + # but time_embedding might actually be running in fp16. so we need to cast here. + # there might be better ways to encapsulate this. + t_emb = t_emb.to(dtype=self.dtype) + emb = self.time_embedding(t_emb) + if self.class_embedding is not None: + if class_labels is None: + raise ValueError( + "class_labels should be provided when num_class_embeds > 0" + ) + + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + + class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) + emb = emb + class_emb + + # pre-process + sample = self.conv_in(sample) + if pose_cond_fea is not None: + sample = sample + pose_cond_fea + + # down + down_block_res_samples = (sample,) + block_count = 1 + for downsample_block in self.down_blocks: + if ( + hasattr(downsample_block, "has_cross_attention") + and downsample_block.has_cross_attention + ): + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + skip_mm=skip_mm, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + skip_mm=skip_mm, + ) + # if pose_cond_fea is not None: + # sample = sample + pose_cond_fea[block_count] + # block_count += 1 + down_block_res_samples += res_samples + + if down_block_additional_residuals is not None: + new_down_block_res_samples = () + + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = ( + down_block_res_sample + down_block_additional_residual + ) + new_down_block_res_samples += (down_block_res_sample,) + + down_block_res_samples = new_down_block_res_samples + + # mid + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + skip_mm=skip_mm, + ) + + if mid_block_additional_residual is not None: + sample = sample + mid_block_additional_residual + + # up + for i, upsample_block in enumerate(self.up_blocks): + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[ + : -len(upsample_block.resnets) + ] + + # if we have not reached the final block and need to forward the + # upsample size, we do it here + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if ( + hasattr(upsample_block, "has_cross_attention") + and upsample_block.has_cross_attention + ): + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + upsample_size=upsample_size, + attention_mask=attention_mask, + skip_mm=skip_mm, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + encoder_hidden_states=encoder_hidden_states, + skip_mm=skip_mm, + ) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + if not return_dict: + return (sample,) + + return UNet3DConditionOutput(sample=sample) + + @classmethod + def from_pretrained_2d( + cls, + pretrained_model_path: PathLike, + motion_module_path: PathLike, + subfolder=None, + unet_additional_kwargs=None, + mm_zero_proj_out=False, + ): + pretrained_model_path = Path(pretrained_model_path) + motion_module_path = Path(motion_module_path) + if subfolder is not None: + pretrained_model_path = pretrained_model_path.joinpath(subfolder) + logger.info( + f"loaded temporal unet's pretrained weights from {pretrained_model_path} ..." + ) + + config_file = pretrained_model_path / "config.json" + if not (config_file.exists() and config_file.is_file()): + raise RuntimeError(f"{config_file} does not exist or is not a file") + + unet_config = cls.load_config(config_file) + unet_config["_class_name"] = cls.__name__ + unet_config["down_block_types"] = [ + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "DownBlock3D", + ] + unet_config["up_block_types"] = [ + "UpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + ] + unet_config["mid_block_type"] = "UNetMidBlock3DCrossAttn" + + model = cls.from_config(unet_config, **unet_additional_kwargs) + # load the vanilla weights + if pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME).exists(): + logger.debug( + f"loading safeTensors weights from {pretrained_model_path} ..." + ) + state_dict = load_file( + pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME), device="cpu" + ) + + elif pretrained_model_path.joinpath(WEIGHTS_NAME).exists(): + logger.debug(f"loading weights from {pretrained_model_path} ...") + state_dict = torch.load( + pretrained_model_path.joinpath(WEIGHTS_NAME), + map_location="cpu", + weights_only=True, + ) + else: + raise FileNotFoundError(f"no weights file found in {pretrained_model_path}") + + # load the motion module weights + if motion_module_path.exists() and motion_module_path.is_file(): + if motion_module_path.suffix.lower() in [".pth", ".pt", ".ckpt"]: + logger.info(f"Load motion module params from {motion_module_path}") + motion_state_dict = torch.load( + motion_module_path, map_location="cpu", weights_only=True + ) + elif motion_module_path.suffix.lower() == ".safetensors": + motion_state_dict = load_file(motion_module_path, device="cpu") + else: + raise RuntimeError( + f"unknown file format for motion module weights: {motion_module_path.suffix}" + ) + + motion_state_dict = { + k.replace('motion_modules.', 'temporal_modules.'): v for k, v in motion_state_dict.items() if not "pos_encoder" in k + } + + if mm_zero_proj_out: + logger.info(f"Zero initialize proj_out layers in motion module...") + new_motion_state_dict = OrderedDict() + for k in motion_state_dict: + if "proj_out" in k: + continue + new_motion_state_dict[k] = motion_state_dict[k] + motion_state_dict = new_motion_state_dict + + # merge the state dicts + state_dict.update(motion_state_dict) + + # load the weights into the model + m, u = model.load_state_dict(state_dict, strict=False) + logger.debug(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};") + + params = [ + p.numel() if "temporal_modules" in n else 0 + for n, p in model.named_parameters() + ] + mm_params = [ + p.numel() if "motion_modules" in n else 0 + for n, p in model.named_parameters() + ] + logger.info( + f"Loaded {sum(mm_params) / 1e6}M-parameter motion module, Loaded {sum(params) / 1e6}M-parameter temporal module" + ) + + return model diff --git a/src/scope/core/pipelines/personalive/modules/unet_3d_blocks.py b/src/scope/core/pipelines/personalive/modules/unet_3d_blocks.py new file mode 100644 index 00000000..26de3be2 --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/unet_3d_blocks.py @@ -0,0 +1,1119 @@ +import pdb +from typing import Dict, Optional +import torch +from torch import nn + +from .motion_module import get_motion_module +from .resnet import Downsample3D, ResnetBlock3D, Upsample3D +from .transformer_3d import Transformer3DModel + + +def get_down_block( + down_block_type, + num_layers, + in_channels, + out_channels, + temb_channels, + add_downsample, + resnet_eps, + resnet_act_fn, + attn_num_head_channels, + resnet_groups=None, + cross_attention_dim=None, + downsample_padding=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift="default", + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + use_inflated_groupnorm=None, + use_motion_module=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, +): + down_block_type = ( + down_block_type[7:] + if down_block_type.startswith("UNetRes") + else down_block_type + ) + if down_block_type == "DownBlock3D": + return DownBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + resnet_time_scale_shift=resnet_time_scale_shift, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + use_temporal_module=use_temporal_module, + temporal_module_type=temporal_module_type, + temporal_module_kwargs=temporal_module_kwargs, + ) + elif down_block_type == "CrossAttnDownBlock3D": + if cross_attention_dim is None: + raise ValueError( + "cross_attention_dim must be specified for CrossAttnDownBlock3D" + ) + return CrossAttnDownBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + add_downsample=add_downsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + downsample_padding=downsample_padding, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + use_temporal_module=use_temporal_module, + temporal_module_type=temporal_module_type, + temporal_module_kwargs=temporal_module_kwargs, + ) + raise ValueError(f"{down_block_type} does not exist.") + + +def get_up_block( + up_block_type, + num_layers, + in_channels, + out_channels, + prev_output_channel, + temb_channels, + add_upsample, + resnet_eps, + resnet_act_fn, + attn_num_head_channels, + resnet_groups=None, + cross_attention_dim=None, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + resnet_time_scale_shift="default", + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + use_inflated_groupnorm=None, + use_motion_module=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, +): + up_block_type = ( + up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type + ) + if up_block_type == "UpBlock3D": + return UpBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_time_scale_shift=resnet_time_scale_shift, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + use_temporal_module=use_temporal_module, + temporal_module_type=temporal_module_type, + temporal_module_kwargs=temporal_module_kwargs, + ) + elif up_block_type == "CrossAttnUpBlock3D": + if cross_attention_dim is None: + raise ValueError( + "cross_attention_dim must be specified for CrossAttnUpBlock3D" + ) + return CrossAttnUpBlock3D( + num_layers=num_layers, + in_channels=in_channels, + out_channels=out_channels, + prev_output_channel=prev_output_channel, + temb_channels=temb_channels, + add_upsample=add_upsample, + resnet_eps=resnet_eps, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attn_num_head_channels, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + use_temporal_module=use_temporal_module, + temporal_module_type=temporal_module_type, + temporal_module_kwargs=temporal_module_kwargs, + ) + raise ValueError(f"{up_block_type} does not exist.") + + +class UNetMidBlock3DCrossAttn(nn.Module): + + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + output_scale_factor=1.0, + cross_attention_dim=1280, + dual_cross_attention=False, + use_linear_projection=False, + upcast_attention=False, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + use_inflated_groupnorm=None, + use_motion_module=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, + **transformer_kwargs, + ): + super().__init__() + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + resnet_groups = ( + resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + ) + + # there is always at least one resnet + resnets = [ + ResnetBlock3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_inflated_groupnorm=use_inflated_groupnorm, + ) + ] + attentions = [] + motion_modules = [] + + for _ in range(num_layers): + if dual_cross_attention: + raise NotImplementedError + attentions.append( + Transformer3DModel( + attn_num_head_channels, + in_channels // attn_num_head_channels, + in_channels=in_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + **transformer_kwargs + ) + ) + motion_modules.append( + get_motion_module( + in_channels=in_channels, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + ) + if use_motion_module + else None + ) + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_inflated_groupnorm=use_inflated_groupnorm, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + self.temporal_modules = nn.ModuleList( + [ + ( + get_motion_module( + in_channels=in_channels, + motion_module_type=temporal_module_type, + motion_module_kwargs=temporal_module_kwargs, + ) + if use_temporal_module + else None + ) + for _ in range(num_layers) + ] + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + skip_mm=False, + mid_reference=None, + ): + if isinstance(encoder_hidden_states, list): + encoder_hidden_states, motion_hidden_states = encoder_hidden_states + else: + motion_hidden_states = encoder_hidden_states + + hidden_states = self.resnets[0](hidden_states, temb) + for attn, resnet, motion_module, temporal_module in zip( + self.attentions, self.resnets[1:], self.motion_modules, self.temporal_modules + ): + if self.training and self.gradient_checkpointing: + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + )[0] + if (motion_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states, + temb, + motion_hidden_states, + ) + if (temporal_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(temporal_module), + hidden_states.requires_grad_(), + temb, + None, + ) + # hidden_states = ( + # temporal_module(hidden_states, temb, encoder_hidden_states=None) + # if (temporal_module is not None) and not skip_mm + # else hidden_states + # ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + else: + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + reference=mid_reference, + ).sample + hidden_states = ( + motion_module( + hidden_states, temb, encoder_hidden_states=motion_hidden_states + ) + if (motion_module is not None) and not skip_mm + else hidden_states + ) + hidden_states = ( + temporal_module(hidden_states, temb, encoder_hidden_states=None, debug=True) + if (temporal_module is not None) and not skip_mm + else hidden_states + ) + + hidden_states = resnet(hidden_states, temb) + + return hidden_states + + +class CrossAttnDownBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + downsample_padding=1, + add_downsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + use_inflated_groupnorm=None, + use_motion_module=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, + **transformer_kwargs, + ): + super().__init__() + resnets = [] + attentions = [] + motion_modules = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_inflated_groupnorm=use_inflated_groupnorm, + ) + ) + if dual_cross_attention: + raise NotImplementedError + attentions.append( + Transformer3DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + **transformer_kwargs, + ) + ) + motion_modules.append( + get_motion_module( + in_channels=out_channels, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + ) + if use_motion_module + else None + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + self.temporal_modules = nn.ModuleList( + [ + ( + get_motion_module( + in_channels=out_channels, + motion_module_type=temporal_module_type, + motion_module_kwargs=temporal_module_kwargs, + ) + if use_temporal_module + else None + ) + for _ in range(num_layers) + ] + ) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + temb=None, + encoder_hidden_states=None, + attention_mask=None, + skip_mm=False, + down_reference=None, + ): + if isinstance(encoder_hidden_states, list): + encoder_hidden_states, motion_hidden_states = encoder_hidden_states + else: + motion_hidden_states = encoder_hidden_states + + output_states = () + + for i, (resnet, attn, motion_module, temporal_module) in enumerate( + zip(self.resnets, self.attentions, self.motion_modules, self.temporal_modules) + ): + reference = None + if not(down_reference is None): + reference = down_reference[i] + # self.gradient_checkpointing = False + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + )[0] + + # add motion module + if (motion_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states, + temb, + motion_hidden_states, + ) + if (temporal_module is not None) and not skip_mm: + # hidden_states = torch.utils.checkpoint.checkpoint( + # create_custom_forward(temporal_module), + # hidden_states.requires_grad_(), + # temb, + # None, + # ) + hidden_states = ( + temporal_module(hidden_states, temb, encoder_hidden_states=None) + if (temporal_module is not None) and not skip_mm + else hidden_states + ) + + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + reference=reference, + ).sample + + # add motion module + hidden_states = ( + motion_module( + hidden_states, temb, encoder_hidden_states=motion_hidden_states + ) + if (motion_module is not None) and not skip_mm + else hidden_states + ) + hidden_states = ( + temporal_module(hidden_states, temb, encoder_hidden_states=None, debug=True) + if (temporal_module is not None) and not skip_mm + else hidden_states + ) + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class DownBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_downsample=True, + downsample_padding=1, + use_inflated_groupnorm=None, + use_motion_module=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, + ): + super().__init__() + resnets = [] + motion_modules = [] + + # use_motion_module = False + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_inflated_groupnorm=use_inflated_groupnorm, + ) + ) + motion_modules.append( + get_motion_module( + in_channels=out_channels, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + ) + if use_motion_module + else None + ) + + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + self.temporal_modules = nn.ModuleList( + [ + ( + get_motion_module( + in_channels=out_channels, + motion_module_type=temporal_module_type, + motion_module_kwargs=temporal_module_kwargs, + ) + if use_temporal_module + else None + ) + for _ in range(num_layers) + ] + ) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + Downsample3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + ) + ] + ) + else: + self.downsamplers = None + + self.gradient_checkpointing = False + + def forward(self, hidden_states, temb=None, encoder_hidden_states=None, skip_mm=False): + output_states = () + if isinstance(encoder_hidden_states, list): + encoder_hidden_states, motion_hidden_states = encoder_hidden_states + else: + motion_hidden_states = encoder_hidden_states + for resnet, motion_module, temporal_module in zip( + self.resnets, self.motion_modules, self.temporal_modules + ): + # print(f"DownBlock3D {self.gradient_checkpointing = }") + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + if (motion_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states, + temb, + motion_hidden_states, + ) + + if (temporal_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(temporal_module), + hidden_states.requires_grad_(), + temb, + None, + ) + else: + hidden_states = resnet(hidden_states, temb) + + # add motion module + hidden_states = ( + motion_module( + hidden_states, temb, encoder_hidden_states=motion_hidden_states + ) + if (motion_module is not None) and not skip_mm + else hidden_states + ) + hidden_states = ( + temporal_module( + hidden_states, temb, encoder_hidden_states=None, debug=True + ) + if (temporal_module is not None) and not skip_mm + else hidden_states + ) + + output_states += (hidden_states,) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states) + + output_states += (hidden_states,) + + return hidden_states, output_states + + +class CrossAttnUpBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + out_channels: int, + prev_output_channel: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + attn_num_head_channels=1, + cross_attention_dim=1280, + output_scale_factor=1.0, + add_upsample=True, + dual_cross_attention=False, + use_linear_projection=False, + only_cross_attention=False, + upcast_attention=False, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + use_motion_module=None, + use_inflated_groupnorm=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, + **transformer_kwargs, + ): + super().__init__() + resnets = [] + attentions = [] + motion_modules = [] + + self.has_cross_attention = True + self.attn_num_head_channels = attn_num_head_channels + + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock3D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_inflated_groupnorm=use_inflated_groupnorm, + ) + ) + if dual_cross_attention: + raise NotImplementedError + attentions.append( + Transformer3DModel( + attn_num_head_channels, + out_channels // attn_num_head_channels, + in_channels=out_channels, + num_layers=1, + cross_attention_dim=cross_attention_dim, + norm_num_groups=resnet_groups, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention, + upcast_attention=upcast_attention, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + **transformer_kwargs, + ) + ) + motion_modules.append( + get_motion_module( + in_channels=out_channels, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + ) + if use_motion_module + else None + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + self.temporal_modules = nn.ModuleList( + [ + ( + get_motion_module( + in_channels=out_channels, + motion_module_type=temporal_module_type, + motion_module_kwargs=temporal_module_kwargs, + ) + if use_temporal_module + else None + ) + for _ in range(num_layers) + ] + ) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [Upsample3D(out_channels, use_conv=True, out_channels=out_channels)] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + res_hidden_states_tuple, + temb=None, + encoder_hidden_states=None, + upsample_size=None, + attention_mask=None, + skip_mm=False, + up_reference=None, + ): + if isinstance(encoder_hidden_states, list): + encoder_hidden_states, motion_hidden_states = encoder_hidden_states + else: + motion_hidden_states = encoder_hidden_states + for i, (resnet, attn, motion_module, temporal_module) in enumerate( + zip(self.resnets, self.attentions, self.motion_modules, self.temporal_modules) + ): + reference = None + if not(up_reference is None): + reference = up_reference[i] + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module, return_dict=None): + def custom_forward(*inputs): + if return_dict is not None: + return module(*inputs, return_dict=return_dict) + else: + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + # hidden_states = attn( + # hidden_states, + # encoder_hidden_states=encoder_hidden_states, + # ).sample + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(attn, return_dict=False), + hidden_states, + encoder_hidden_states, + )[0] + if (motion_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states, + temb, + motion_hidden_states, + ) + if (temporal_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(temporal_module), + hidden_states.requires_grad_(), + temb, + None, + ) + + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = attn( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + reference=reference + ).sample + + # add motion module + hidden_states = ( + motion_module( + hidden_states, temb, encoder_hidden_states=motion_hidden_states + ) + if (motion_module is not None) and not skip_mm + else hidden_states + ) + + # add temporal_module + hidden_states = ( + temporal_module(hidden_states, temb, encoder_hidden_states=None, debug=True) + if (temporal_module is not None) and not skip_mm + else hidden_states + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states + + +class UpBlock3D(nn.Module): + + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor=1.0, + add_upsample=True, + use_inflated_groupnorm=None, + use_motion_module=None, + motion_module_type=None, + motion_module_kwargs=None, + use_temporal_module=None, + temporal_module_type=None, + temporal_module_kwargs=None, + ): + super().__init__() + resnets = [] + motion_modules = [] + + # use_motion_module = False + for i in range(num_layers): + res_skip_channels = in_channels if (i == num_layers - 1) else out_channels + resnet_in_channels = prev_output_channel if i == 0 else out_channels + + resnets.append( + ResnetBlock3D( + in_channels=resnet_in_channels + res_skip_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + use_inflated_groupnorm=use_inflated_groupnorm, + ) + ) + motion_modules.append( + get_motion_module( + in_channels=out_channels, + motion_module_type=motion_module_type, + motion_module_kwargs=motion_module_kwargs, + ) + if use_motion_module + else None + ) + + self.resnets = nn.ModuleList(resnets) + self.motion_modules = nn.ModuleList(motion_modules) + self.temporal_modules = nn.ModuleList( + [ + ( + get_motion_module( + in_channels=out_channels, + motion_module_type=temporal_module_type, + motion_module_kwargs=temporal_module_kwargs, + ) + if use_temporal_module + else None + ) + for _ in range(num_layers) + ] + ) + + if add_upsample: + self.upsamplers = nn.ModuleList( + [Upsample3D(out_channels, use_conv=True, out_channels=out_channels)] + ) + else: + self.upsamplers = None + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + res_hidden_states_tuple, + temb=None, + upsample_size=None, + encoder_hidden_states=None, + skip_mm=False, + ): + if isinstance(encoder_hidden_states, list): + encoder_hidden_states, motion_hidden_states = encoder_hidden_states + else: + motion_hidden_states = encoder_hidden_states + for resnet, motion_module, temporal_module in zip(self.resnets, self.motion_modules, self.temporal_modules): + # pop res hidden states + res_hidden_states = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) + + # print(f"UpBlock3D {self.gradient_checkpointing = }") + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(resnet), hidden_states, temb + ) + if (motion_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(motion_module), + hidden_states, + temb, + motion_hidden_states, + ) + if (temporal_module is not None) and not skip_mm: + hidden_states = torch.utils.checkpoint.checkpoint( + create_custom_forward(temporal_module), + hidden_states.requires_grad_(), + temb, + None, + ) + + else: + hidden_states = resnet(hidden_states, temb) + hidden_states = ( + motion_module( + hidden_states, temb, encoder_hidden_states=motion_hidden_states + ) + if (motion_module is not None) and not skip_mm + else hidden_states + ) + hidden_states = ( + temporal_module(hidden_states, temb, encoder_hidden_states=None, debug=True) + if (temporal_module is not None) and not skip_mm + else hidden_states + ) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, upsample_size) + + return hidden_states diff --git a/src/scope/core/pipelines/personalive/modules/util.py b/src/scope/core/pipelines/personalive/modules/util.py new file mode 100644 index 00000000..7f9ccbfe --- /dev/null +++ b/src/scope/core/pipelines/personalive/modules/util.py @@ -0,0 +1,313 @@ +import importlib +import os +import os.path as osp +import shutil +import sys +from pathlib import Path + +import av +import numpy as np +import torch +import torchvision +from einops import rearrange +from PIL import Image +import cv2 + + +def save_checkpoint(model, save_dir, prefix, ckpt_num, logger, total_limit=None): + save_path = osp.join(save_dir, f"{prefix}-{ckpt_num}.pth") + + if total_limit is not None: + checkpoints = os.listdir(save_dir) + checkpoints = [d for d in checkpoints if d.startswith(prefix)] + checkpoints = sorted( + checkpoints, key=lambda x: int(x.split("-")[1].split(".")[0]) + ) + + if len(checkpoints) >= total_limit: + num_to_remove = len(checkpoints) - total_limit + 1 + removing_checkpoints = checkpoints[0:num_to_remove] + logger.info( + f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints" + ) + logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}") + + for removing_checkpoint in removing_checkpoints: + removing_checkpoint = os.path.join(save_dir, removing_checkpoint) + os.remove(removing_checkpoint) + + state_dict = model.state_dict() + torch.save(state_dict, save_path) + + +def create_code_snapshot(root, dst_path, extensions=(".py", ".h", ".cpp", ".cu", ".cc", ".cuh", ".json", ".sh", ".bat", ".yaml"), exclude=()): + """Creates tarball with the source code""" + import tarfile + from pathlib import Path + with tarfile.open(str(dst_path), "w:gz") as tar: + for path in Path(root).rglob("*"): + if '.git' in path.parts: + continue + exclude_flag = False + if len(exclude) > 0: + for k in exclude: + if k in path.parts: + exclude_flag = True + if exclude_flag: + continue + if path.suffix.lower() in extensions: + try: + tar.add(path.as_posix(), arcname=path.relative_to( + root).as_posix(), recursive=True) + except: + print(path) + assert False, 'Error occur in create_code_snapshot' + +def seed_everything(seed): + import random + + import numpy as np + + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + np.random.seed(seed % (2**32)) + random.seed(seed) + + +def import_filename(filename): + spec = importlib.util.spec_from_file_location("mymodule", filename) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def delete_additional_ckpt(base_path, num_keep): + dirs = [] + for d in os.listdir(base_path): + if d.startswith("checkpoint-"): + dirs.append(d) + num_tot = len(dirs) + if num_tot <= num_keep: + return + # ensure ckpt is sorted and delete the ealier! + del_dirs = sorted(dirs, key=lambda x: int(x.split("-")[-1]))[: num_tot - num_keep] + for d in del_dirs: + path_to_dir = osp.join(base_path, d) + if osp.exists(path_to_dir): + shutil.rmtree(path_to_dir) + + +def save_videos_from_pil(pil_images, path, fps=8, crf=None): + import av + + save_fmt = Path(path).suffix + os.makedirs(os.path.dirname(path), exist_ok=True) + width, height = pil_images[0].size + + if save_fmt == ".mp4": + if True: + codec = "libx264" + container = av.open(path, "w") + stream = container.add_stream(codec, rate=fps) + + stream.width = width + stream.height = height + if crf is not None: + stream.options = {'crf': str(crf)} + + for pil_image in pil_images: + # pil_image = Image.fromarray(image_arr).convert("RGB") + av_frame = av.VideoFrame.from_image(pil_image) + container.mux(stream.encode(av_frame)) + container.mux(stream.encode()) + container.close() + else: + + video_writer = cv2.VideoWriter( + path.replace('.mp4', '_cv.mp4'), cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height) + ) + for pil_image in pil_images: + img_np = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR) + video_writer.write(img_np) + video_writer.release() + + elif save_fmt == ".gif": + pil_images[0].save( + fp=path, + format="GIF", + append_images=pil_images[1:], + save_all=True, + duration=(1 / fps * 1000), + loop=0, + ) + else: + raise ValueError("Unsupported file type. Use .mp4 or .gif.") + + +def save_videos_grid(videos_, path: str, rescale=False, n_rows=6, fps=8, crf=None): + if not isinstance(videos_, list): videos_ = [videos_] + + outputs = [] + vid_len = videos_[0].shape[2] + for i in range(vid_len): + output = [] + for videos in videos_: + videos = rearrange(videos, "b c t h w -> t b c h w") + height, width = videos.shape[-2:] + + x = torchvision.utils.make_grid(videos[i], nrow=n_rows) # (c h w) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = (x * 255).numpy().astype(np.uint8) + output.append(x) + + output = Image.fromarray(np.concatenate(output, axis=0)) + outputs.append(output) + + os.makedirs(os.path.dirname(path), exist_ok=True) + save_videos_from_pil(outputs, path, fps, crf) + + +def save_videos_grid_ori(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8): + videos = rearrange(videos, "b c t h w -> t b c h w") + height, width = videos.shape[-2:] + outputs = [] + + for x in videos: + x = torchvision.utils.make_grid(x, nrow=n_rows) # (c h w) + x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c) + if rescale: + x = (x + 1.0) / 2.0 # -1,1 -> 0,1 + x = (x * 255).numpy().astype(np.uint8) + x = Image.fromarray(x) + + outputs.append(x) + + os.makedirs(os.path.dirname(path), exist_ok=True) + + save_videos_from_pil(outputs, path, fps) + + +def read_frames(video_path): + container = av.open(video_path) + + video_stream = next(s for s in container.streams if s.type == "video") + frames = [] + for packet in container.demux(video_stream): + for frame in packet.decode(): + image = Image.frombytes( + "RGB", + (frame.width, frame.height), + frame.to_rgb().to_ndarray(), + ) + frames.append(image) + + return frames + + +def get_fps(video_path): + container = av.open(video_path) + video_stream = next(s for s in container.streams if s.type == "video") + fps = video_stream.average_rate + container.close() + return fps + +def draw_keypoints(keypoints, height=512, width=512, device="cuda"): + colors = torch.tensor([ + [255, 0, 0], + [255, 255, 0], + [0, 255, 0], + [0, 255, 255], + [0, 0, 255], + [255, 0, 255], + [255, 0, 85], + ], device=device, dtype=torch.float32) + + selected = torch.tensor([1, 2, 3, 4, 12, 15, 20], device=device) + B = keypoints.shape[0] + + # [B, len(selected), 2] + pts = keypoints[:, selected] * 0.5 + 0.5 + pts[..., 0] *= width + pts[..., 1] *= height + pts = pts.long() + + canvas = torch.zeros((B, 3, height, width), device=device) + radius = 4 + + for i, color in enumerate(colors): + x = pts[:, i, 0] + y = pts[:, i, 1] + mask = ( + (x[:, None, None] - torch.arange(width, device=device)) ** 2 + + (y[:, None, None] - torch.arange(height, device=device)[:, None]) ** 2 + ) <= radius**2 + canvas[:, 0] += color[0] / 255.0 * mask + canvas[:, 1] += color[1] / 255.0 * mask + canvas[:, 2] += color[2] / 255.0 * mask + + return canvas.clamp(0, 1) + +def get_boxes(keypoints, height=512, width=512): + selected = torch.tensor([1, 2, 3, 4, 12, 15, 20]) + + # [B, len(selected), 2] + pts = keypoints[:, selected] * 0.5 + 0.5 + pts[..., 0] *= width + pts[..., 1] *= height + pts = pts.long() + + cx = pts[..., 0].float().mean(dim=1) # [B] + cy = pts[..., 1].float().mean(dim=1) # [B] + + min_y = pts[..., 1].float().min(dim=1)[0] # [B] + + side = (cy - min_y) * 2.0 + side = side * 1.7 + + x1 = (cx - side / 2 * 0.95).clamp(0, width - 1).long() + y1 = (cy - side / 2 * 0.95).clamp(0, height - 1).long() + x2 = (cx + side / 2 * 1.05).clamp(0, width - 1).long() + y2 = (cy + side / 2 * 1.05).clamp(0, height - 1).long() + + boxes = torch.stack([x1, y1, x2, y2], dim=1) # [B, 4] + + return boxes + +def crop_face(image_pil, face_mesh): + image = np.array(image_pil) + h, w = image.shape[:2] + results = face_mesh.process(image) + face_landmarks = results.multi_face_landmarks[0] + coords = [(int(l.x * w), int(l.y * h)) for l in face_landmarks.landmark] + xs, ys = zip(*coords) + x1, y1 = min(xs), min(ys) + x2, y2 = max(xs), max(ys) + face_box = (x1, y1, x2, y2) + + left, top, right, bot = scale_bb(face_box, scale=1.1, size=image.shape[:2]) + + face_patch = image[int(top) : int(bot), int(left) : int(right)] + + return face_patch + +def scale_bb(bbox, scale, size): + left, top, right, bot = bbox + width = right - left + height = bot - top + length = max(width, height) * scale + center_X = (left + right) * 0.5 + center_Y = (top + bot) * 0.5 + left, top, right, bot = [ + center_X - length / 2, + center_Y - length / 2, + center_X + length / 2, + center_Y + length / 2, + ] + left = max(0, left) + top = max(0, top) + right = min(size[1] - 1, right) + bot = min(size[0] - 1, bot) + return np.array([left, top, right, bot]) \ No newline at end of file diff --git a/src/scope/core/pipelines/personalive/pipeline.py b/src/scope/core/pipelines/personalive/pipeline.py new file mode 100644 index 00000000..1a3da79c --- /dev/null +++ b/src/scope/core/pipelines/personalive/pipeline.py @@ -0,0 +1,656 @@ +"""PersonaLive pipeline for real-time portrait animation. + +Animates a reference portrait image using driving video frames. +Based on: https://github.com/GVCLab/PersonaLive +""" + +import logging +import time +from collections import deque +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import torch +import torch.nn.functional as F +from diffusers import AutoencoderKL +from diffusers.image_processor import VaeImageProcessor +from einops import rearrange +from omegaconf import OmegaConf +from PIL import Image +from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection + +from ..interface import Pipeline, Requirements +from ..process import preprocess_chunk, postprocess_chunk +from ..schema import PersonaLiveConfig +from ..utils import load_model_config +from .components import FaceDetector +from .modules import ( + DDIMScheduler, + MotEncoder, + MotionExtractor, + PoseGuider, + ReferenceAttentionControl, + UNet2DConditionModel, + UNet3DConditionModel, + draw_keypoints, + get_boxes, +) + +if TYPE_CHECKING: + from ..schema import BasePipelineConfig + +logger = logging.getLogger(__name__) + + +class PersonaLivePipeline(Pipeline): + """PersonaLive portrait animation pipeline. + + This pipeline animates a reference portrait image using driving video frames. + It requires: + 1. A reference image to be set once via `fuse_reference()` + 2. Continuous driving video frames via `__call__()` + + The output is animated video frames of the reference portrait following + the motion from the driving video. + """ + + @classmethod + def get_config_class(cls) -> type["BasePipelineConfig"]: + return PersonaLiveConfig + + def __init__( + self, + config, + device: torch.device | None = None, + dtype: torch.dtype = torch.float16, + ): + """Initialize PersonaLive pipeline. + + Args: + config: Pipeline configuration (OmegaConf or dict-like). + device: Target device for models. + dtype: Data type for model weights. + """ + if device is None: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + self.device = device + self.dtype = dtype + self.numpy_dtype = np.float16 if dtype == torch.float16 else np.float32 + + # Load model config + model_config = load_model_config(config, __file__) + + # Get paths from config + model_dir = Path(getattr(config, "model_dir", ".")) + personalive_dir = model_dir / "PersonaLive" / "pretrained_weights" + + # Model paths + pretrained_base_path = personalive_dir / "sd-image-variations-diffusers" + vae_path = personalive_dir / "sd-vae-ft-mse" + personalive_weights = personalive_dir / "personalive" + image_encoder_path = pretrained_base_path / "image_encoder" + + # Pipeline parameters + self.temporal_window_size = getattr(model_config, "temporal_window_size", 4) + self.temporal_adaptive_step = getattr(model_config, "temporal_adaptive_step", 4) + self.num_inference_steps = getattr(model_config, "num_inference_steps", 4) + self.batch_size = getattr(model_config, "batch_size", 1) + + # Resolution + self.height = getattr(config, "height", 512) + self.width = getattr(config, "width", 512) + self.ref_height = getattr(model_config, "reference_image_height", 512) + self.ref_width = getattr(model_config, "reference_image_width", 512) + + # Load inference config for UNet kwargs + unet_additional_kwargs = getattr(model_config, "unet_additional_kwargs", {}) + if isinstance(unet_additional_kwargs, OmegaConf): + unet_additional_kwargs = OmegaConf.to_container(unet_additional_kwargs) + + # Scheduler config + sched_kwargs = getattr(model_config, "noise_scheduler_kwargs", {}) + if isinstance(sched_kwargs, OmegaConf): + sched_kwargs = OmegaConf.to_container(sched_kwargs) + + logger.info("Loading PersonaLive models...") + start = time.time() + + # Load pose guider + self.pose_guider = PoseGuider().to(device=device, dtype=dtype) + pose_guider_path = personalive_weights / "pose_guider.pth" + if pose_guider_path.exists(): + state_dict = torch.load(pose_guider_path, map_location="cpu") + self.pose_guider.load_state_dict(state_dict) + del state_dict + logger.info(f"Loaded pose_guider in {time.time() - start:.2f}s") + + # Load motion encoder + start = time.time() + self.motion_encoder = MotEncoder().to(dtype=dtype, device=device).eval() + motion_encoder_path = personalive_weights / "motion_encoder.pth" + if motion_encoder_path.exists(): + state_dict = torch.load(motion_encoder_path, map_location="cpu") + self.motion_encoder.load_state_dict(state_dict) + del state_dict + logger.info(f"Loaded motion_encoder in {time.time() - start:.2f}s") + + # Load pose encoder (motion extractor) + start = time.time() + self.pose_encoder = MotionExtractor(num_kp=21).to(device=device, dtype=dtype).eval() + pose_encoder_path = personalive_weights / "motion_extractor.pth" + if pose_encoder_path.exists(): + state_dict = torch.load(pose_encoder_path, map_location="cpu") + self.pose_encoder.load_state_dict(state_dict, strict=False) + del state_dict + logger.info(f"Loaded pose_encoder in {time.time() - start:.2f}s") + + # Load denoising UNet (3D) + start = time.time() + self.denoising_unet = UNet3DConditionModel.from_pretrained_2d( + str(pretrained_base_path), + "", + subfolder="unet", + unet_additional_kwargs=unet_additional_kwargs, + ).to(dtype=dtype, device=device) + + denoising_unet_path = personalive_weights / "denoising_unet.pth" + if denoising_unet_path.exists(): + state_dict = torch.load(denoising_unet_path, map_location="cpu") + self.denoising_unet.load_state_dict(state_dict, strict=False) + del state_dict + + temporal_module_path = personalive_weights / "temporal_module.pth" + if temporal_module_path.exists(): + state_dict = torch.load(temporal_module_path, map_location="cpu") + self.denoising_unet.load_state_dict(state_dict, strict=False) + del state_dict + logger.info(f"Loaded denoising_unet in {time.time() - start:.2f}s") + + # Load reference UNet (2D) + start = time.time() + self.reference_unet = UNet2DConditionModel.from_pretrained( + str(pretrained_base_path), + subfolder="unet", + ).to(dtype=dtype, device=device) + + reference_unet_path = personalive_weights / "reference_unet.pth" + if reference_unet_path.exists(): + state_dict = torch.load(reference_unet_path, map_location="cpu") + self.reference_unet.load_state_dict(state_dict) + del state_dict + logger.info(f"Loaded reference_unet in {time.time() - start:.2f}s") + + # Setup reference attention control + self.reference_control_writer = ReferenceAttentionControl( + self.reference_unet, + do_classifier_free_guidance=False, + mode="write", + batch_size=self.batch_size, + fusion_blocks="full", + ) + self.reference_control_reader = ReferenceAttentionControl( + self.denoising_unet, + do_classifier_free_guidance=False, + mode="read", + batch_size=self.batch_size, + fusion_blocks="full", + cache_kv=True, + ) + + # Load VAE + start = time.time() + self.vae = AutoencoderKL.from_pretrained(str(vae_path)).to( + device=device, dtype=dtype + ) + logger.info(f"Loaded VAE in {time.time() - start:.2f}s") + + # Load image encoder + start = time.time() + self.image_encoder = CLIPVisionModelWithProjection.from_pretrained( + str(image_encoder_path), + ).to(device=device, dtype=dtype) + logger.info(f"Loaded image_encoder in {time.time() - start:.2f}s") + + # Setup scheduler + self.scheduler = DDIMScheduler(**sched_kwargs) + self.timesteps = torch.tensor([999, 666, 333, 0], device=device).long() + self.scheduler.set_step_length(333) + + # Setup generator + seed = getattr(config, "seed", 42) + self.generator = torch.Generator(device) + self.generator.manual_seed(seed) + + # Image processors + self.vae_scale_factor = 8 + self.ref_image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True + ) + self.clip_image_processor = CLIPImageProcessor() + self.cond_image_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=True + ) + + # Face detector + self.face_detector = FaceDetector() + + # State variables + self.first_frame = True + self.motion_bank = None + self.count = 0 + self.num_khf = 0 # Number of history keyframes + + # Temporal buffers + self.latents_pile = deque([]) + self.pose_pile = deque([]) + self.motion_pile = deque([]) + + # Reference state (set via fuse_reference) + self.reference_fused = False + self.encoder_hidden_states = None + self.ref_image_tensor = None + self.ref_image_latents = None + self.ref_cond_tensor = None + self.kps_ref = None + self.kps_frame1 = None + + # Enable memory efficient attention + self._enable_xformers() + + torch.cuda.empty_cache() + logger.info("PersonaLive pipeline initialized") + + def _enable_xformers(self): + """Enable xformers memory efficient attention if available.""" + try: + self.reference_unet.enable_xformers_memory_efficient_attention() + self.denoising_unet.enable_xformers_memory_efficient_attention() + logger.info("Enabled xformers memory efficient attention") + except Exception as e: + logger.warning(f"Could not enable xformers: {e}") + + def prepare(self, **kwargs) -> Requirements | None: + """Return input requirements. + + PersonaLive requires: + - Reference image to be fused first via fuse_reference() + - Driving video frames (temporal_window_size frames per call) + + Returns Requirements only when video mode is signaled (video key in kwargs), + following the same pattern as other video pipelines. + """ + # Only return requirements when video mode is signaled + # This is indicated by FrameProcessor setting video=True in prepare_params + if kwargs.get("video") is not None: + return Requirements(input_size=self.temporal_window_size) + return None + + def _fast_resize(self, images: torch.Tensor, height: int, width: int) -> torch.Tensor: + """Fast bilinear resize of image tensor.""" + return F.interpolate( + images, + size=(height, width), + mode="bilinear", + align_corners=False, + ) + + def _interpolate_tensors(self, a: torch.Tensor, b: torch.Tensor, num: int) -> torch.Tensor: + """Linear interpolation between tensors a and b.""" + if a.shape != b.shape: + raise ValueError(f"Shape mismatch: a.shape={a.shape}, b.shape={b.shape}") + + B, _, *rest = a.shape + alphas = torch.linspace(0, 1, num, device=a.device, dtype=a.dtype) + view_shape = (1, num) + (1,) * len(rest) + alphas = alphas.view(view_shape) + + result = (1 - alphas) * a + alphas * b + return result + + def _calculate_dis(self, A: torch.Tensor, B: torch.Tensor, threshold: float = 10.0): + """Calculate distance between motion features for keyframe selection.""" + A_flat = A.view(A.size(1), -1).clone() + B_flat = B.view(B.size(1), -1).clone() + + dist = torch.cdist(B_flat.to(torch.float32), A_flat.to(torch.float32), p=2) + min_dist, min_idx = dist.min(dim=1) + + idx_to_add = torch.nonzero(min_dist[:1] > threshold, as_tuple=False).squeeze(1).tolist() + + if len(idx_to_add) > 0: + B_to_add = B[:, idx_to_add] + A_new = torch.cat([A, B_to_add], dim=1) + else: + A_new = A + + return idx_to_add, A_new, min_idx + + def _crop_face_tensor(self, image_tensor: torch.Tensor, boxes) -> torch.Tensor: + """Crop face from tensor using bounding box.""" + left, top, right, bot = boxes + left, top, right, bottom = map(int, (left, top, right, bot)) + + face_patch = image_tensor[:, top:bottom, left:right] + face_patch = F.interpolate( + face_patch.unsqueeze(0), + size=(224, 224), + mode="bilinear", + align_corners=False, + ) + return face_patch + + @torch.no_grad() + def fuse_reference(self, ref_image: Image.Image): + """Fuse reference image into the pipeline. + + This must be called before processing driving frames. + + Args: + ref_image: PIL Image of the reference portrait. + """ + logger.info("Fusing reference image...") + + # Process for CLIP - resize to 224x224 first (matching official implementation) + clip_image = self.clip_image_processor.preprocess( + ref_image.resize((224, 224)), return_tensors="pt" + ).pixel_values + + # Process for VAE + ref_image_tensor = self.ref_image_processor.preprocess( + ref_image, height=self.ref_height, width=self.ref_width + ) + + # Get CLIP embeddings + clip_image_embeds = self.image_encoder( + clip_image.to(self.image_encoder.device, dtype=self.image_encoder.dtype) + ).image_embeds + self.encoder_hidden_states = clip_image_embeds.unsqueeze(1) + + # Encode reference image + ref_image_tensor = ref_image_tensor.to(dtype=self.vae.dtype, device=self.vae.device) + self.ref_image_tensor = ref_image_tensor.squeeze(0) + ref_image_latents = self.vae.encode(ref_image_tensor).latent_dist.mean + ref_image_latents = ref_image_latents * 0.18215 + + # Run reference UNet to cache features + self.reference_unet( + ref_image_latents.to(self.reference_unet.device), + torch.zeros((self.batch_size,), dtype=self.dtype, device=self.reference_unet.device), + encoder_hidden_states=self.encoder_hidden_states, + return_dict=False, + ) + self.reference_control_reader.update(self.reference_control_writer) + self.encoder_hidden_states = self.encoder_hidden_states.to(self.device) + + # Prepare conditioning tensor for pose encoder + ref_cond_tensor = self.cond_image_processor.preprocess( + ref_image, height=256, width=256 + ).to(device=self.device, dtype=self.pose_encoder.dtype) + self.ref_cond_tensor = ref_cond_tensor / 2 + 0.5 + self.ref_image_latents = ref_image_latents + + # Reset state - clear ALL piles first to handle re-fusing + self.first_frame = True + self.motion_bank = None + self.count = 0 + self.num_khf = 0 + self.latents_pile.clear() + self.pose_pile.clear() + self.motion_pile.clear() + + # Initialize latent piles with padding + padding_num = (self.temporal_adaptive_step - 1) * self.temporal_window_size + init_latents = ref_image_latents.unsqueeze(2).repeat(1, 1, padding_num, 1, 1) + noise = torch.randn_like(init_latents) + init_timesteps = reversed(self.timesteps).repeat_interleave(self.temporal_window_size, dim=0) + noisy_latents_first = self.scheduler.add_noise(init_latents, noise, init_timesteps[:padding_num]) + + for i in range(self.temporal_adaptive_step - 1): + l = i * self.temporal_window_size + r = (i + 1) * self.temporal_window_size + self.latents_pile.append(noisy_latents_first[:, :, l:r]) + + self.reference_fused = True + logger.info("Reference image fused successfully") + + @torch.no_grad() + def __call__( + self, + video: torch.Tensor | list[torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + """Process driving video frames. + + Args: + video: Driving video frames. Can be: + - Tensor of shape (B, C, T, H, W) in [0, 1] range + - List of frame tensors in THWC format and [0, 255] range + - None (not supported) + **kwargs: Additional parameters (ignored). + + Returns: + Animated output frames in THWC format and [0, 1] range. + """ + if not self.reference_fused: + raise RuntimeError( + "Reference image must be fused before processing. " + "Call fuse_reference() first." + ) + + # Convert input to expected format + if video is None: + # No video frames yet - return empty tensor + # This can happen during streaming startup + logger.debug("No video frames provided, returning empty output") + return torch.empty(0, self.height, self.width, 3, device="cpu") + + # Use standard preprocess_chunk for list input (from frame_processor) + if isinstance(video, list): + # preprocess_chunk expects list of (1, H, W, C) tensors + # Returns (B, C, T, H, W) tensor in [-1, 1] range + frames = preprocess_chunk( + video, + self.device, + self.dtype, + height=self.height, + width=self.width, + ) + else: + # Already a tensor - ensure correct device/dtype + frames = video.to(device=self.device, dtype=self.dtype) + + # Process input frames + output = self._process_frames(frames) + + return output + + @torch.no_grad() + def _process_frames(self, images: torch.Tensor) -> torch.Tensor: + """Process driving frames through the pipeline. + + Args: + images: Input tensor of shape (B, C, T, H, W) in [-1, 1] or [0, 1] range. + + Returns: + Output frames tensor of shape (T, H, W, C) in [0, 1] range. + """ + batch_size = self.batch_size + device = self.device + temporal_window_size = self.temporal_window_size + temporal_adaptive_step = self.temporal_adaptive_step + + # Reshape from (B, C, T, H, W) to (T, C, H, W) for processing + if images.dim() == 5: + images = images.squeeze(0).permute(1, 0, 2, 3) # (T, C, H, W) + + # PersonaLive expects exactly temporal_window_size frames per call + # If we receive more, process only the required number + num_frames = images.shape[0] + if num_frames != temporal_window_size: + logger.warning( + f"Expected {temporal_window_size} frames but received {num_frames}. " + f"Processing only the last {temporal_window_size} frames." + ) + images = images[-temporal_window_size:] + + # Store original images for face cropping (they should be in [0, 1] after preprocess_chunk) + # But ensure they're in [0, 1] range for consistency + if images.min() < 0: + images = images / 2 + 0.5 + images = images.clamp(0, 1) + + # Resize for pose encoder (matching official: resize first, then normalize) + # Official does: fast_resize(images, 256, 256) then / 2 + 0.5 + # This suggests images are in [-1, 1] before resize in official + # But since we have [0, 1], we'll resize directly and ensure [0, 1] range + tgt_cond_tensor = self._fast_resize(images, 256, 256) + # Ensure [0, 1] range (official normalizes after resize, but we already have [0, 1]) + tgt_cond_tensor = tgt_cond_tensor.clamp(0, 1) + + # Get keypoints + if self.first_frame: + mot_bbox_param, kps_ref, kps_frame1, kps_dri = self.pose_encoder.interpolate_kps_online( + self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + ) + self.kps_ref = kps_ref + self.kps_frame1 = kps_frame1 + else: + mot_bbox_param, kps_dri = self.pose_encoder.get_kps( + self.kps_ref, self.kps_frame1, tgt_cond_tensor + ) + + # Draw keypoints and get bounding boxes + keypoints = draw_keypoints(mot_bbox_param, device=device) + boxes = get_boxes(kps_dri) + keypoints = rearrange(keypoints.unsqueeze(2), 'f c b h w -> b c f h w') + keypoints = keypoints.to(device=device, dtype=self.pose_guider.dtype) + + # Process motion features + if self.first_frame: + ref_box = get_boxes(mot_bbox_param[:1]) + ref_face = self._crop_face_tensor(self.ref_image_tensor, ref_box[0]) + motion_face = [ref_face] + for i, frame in enumerate(images): + motion_face.append(self._crop_face_tensor(frame, boxes[i])) + pose_cond_tensor = torch.cat(motion_face, dim=0).transpose(0, 1) + pose_cond_tensor = pose_cond_tensor.unsqueeze(0) + motion_hidden_states = self.motion_encoder(pose_cond_tensor) + ref_motion = motion_hidden_states[:, :1] + dri_motion = motion_hidden_states[:, 1:] + + init_motion_hidden_states = self._interpolate_tensors( + ref_motion, dri_motion[:, :1], num=12 + 1 + )[:, :-1] + for i in range(temporal_adaptive_step - 1): + l = i * temporal_window_size + r = (i + 1) * temporal_window_size + self.motion_pile.append(init_motion_hidden_states[:, l:r]) + self.motion_pile.append(dri_motion) + + self.motion_bank = ref_motion + else: + motion_face = [] + for i, frame in enumerate(images): + motion_face.append(self._crop_face_tensor(frame, boxes[i])) + pose_cond_tensor = torch.cat(motion_face, dim=0).transpose(0, 1) + pose_cond_tensor = pose_cond_tensor.unsqueeze(0) + motion_hidden_states = self.motion_encoder(pose_cond_tensor) + self.motion_pile.append(motion_hidden_states) + + # Encode pose features + pose_fea = self.pose_guider(keypoints) + if self.first_frame: + for i in range(temporal_adaptive_step): + l = i * temporal_window_size + r = (i + 1) * temporal_window_size + self.pose_pile.append(pose_fea[:, :, l:r]) + self.first_frame = False + else: + self.pose_pile.append(pose_fea) + + # Prepare noisy latents for new frames + latents = self.ref_image_latents.unsqueeze(2).repeat(1, 1, temporal_window_size, 1, 1) + noise = torch.randn_like(latents) + latents = self.scheduler.add_noise(latents, noise, self.timesteps[:1]) + self.latents_pile.append(latents) + + # Combine piles + jump = 1 + motion_hidden_state = torch.cat(list(self.motion_pile), dim=1) + pose_cond_fea = torch.cat(list(self.pose_pile), dim=2) + + # Check for keyframe addition + idx_to_add = [] + if self.count > 8: + idx_to_add, self.motion_bank, idx_his = self._calculate_dis( + self.motion_bank, motion_hidden_state, threshold=17.0 + ) + + # Denoise + latents_model_input = torch.cat(list(self.latents_pile), dim=2) + for j in range(jump): + timesteps = reversed(self.timesteps[j::jump]).repeat_interleave(temporal_window_size, dim=0) + timesteps = torch.stack([timesteps] * batch_size) + timesteps = rearrange(timesteps, 'b f -> (b f)') + + noise_pred = self.denoising_unet( + latents_model_input, + timesteps, + encoder_hidden_states=[ + self.encoder_hidden_states, + motion_hidden_state, + ], + pose_cond_fea=pose_cond_fea, + return_dict=False, + )[0] + + clip_length = noise_pred.shape[2] + mid_noise_pred = rearrange(noise_pred, 'b c f h w -> (b f) c h w') + mid_latents = rearrange(latents_model_input, 'b c f h w -> (b f) c h w') + latents_model_input, pred_original_sample = self.scheduler.step( + mid_noise_pred, timesteps, mid_latents, + generator=self.generator, return_dict=False + ) + latents_model_input = rearrange(latents_model_input, '(b f) c h w -> b c f h w', f=clip_length) + pred_original_sample = rearrange(pred_original_sample, '(b f) c h w -> b c f h w', f=clip_length) + latents_model_input = torch.cat([ + pred_original_sample[:, :, :temporal_window_size], + latents_model_input[:, :, temporal_window_size:] + ], dim=2) + latents_model_input = latents_model_input.to(dtype=self.dtype) + + # History keyframe mechanism + if len(idx_to_add) > 0 and self.num_khf < 3: + self.reference_control_writer.clear() + self.reference_unet( + pred_original_sample[:, :, 0].to(self.reference_unet.dtype), + torch.zeros((batch_size,), dtype=self.dtype, device=self.reference_unet.device), + encoder_hidden_states=self.encoder_hidden_states, + return_dict=False, + ) + self.reference_control_reader.update_hkf(self.reference_control_writer) + logger.debug("Added history keyframe") + self.num_khf += 1 + + # Update latent piles + for i in range(len(self.latents_pile)): + self.latents_pile[i] = latents_model_input[ + :, :, i * temporal_adaptive_step:(i + 1) * temporal_adaptive_step, :, : + ] + + # Pop oldest latents and decode + self.pose_pile.popleft() + self.motion_pile.popleft() + latents = self.latents_pile.popleft() + latents = 1 / 0.18215 * latents + latents = rearrange(latents, "b c f h w -> (b f) c h w") + video = self.vae.decode(latents).sample + video = rearrange(video, "b c h w -> b h w c") + video = (video / 2 + 0.5).clamp(0, 1) + + self.count += 1 + + # Return in THWC format [0, 1] + return video.cpu() diff --git a/src/scope/core/pipelines/registry.py b/src/scope/core/pipelines/registry.py index 493804be..d1f45c3e 100644 --- a/src/scope/core/pipelines/registry.py +++ b/src/scope/core/pipelines/registry.py @@ -71,6 +71,7 @@ def _register_pipelines(): from .krea_realtime_video.pipeline import KreaRealtimeVideoPipeline from .longlive.pipeline import LongLivePipeline from .passthrough.pipeline import PassthroughPipeline + from .personalive.pipeline import PersonaLivePipeline from .reward_forcing.pipeline import RewardForcingPipeline from .streamdiffusionv2.pipeline import StreamDiffusionV2Pipeline @@ -81,6 +82,7 @@ def _register_pipelines(): StreamDiffusionV2Pipeline, PassthroughPipeline, RewardForcingPipeline, + PersonaLivePipeline, ]: config_class = pipeline_class.get_config_class() PipelineRegistry.register(config_class.pipeline_id, pipeline_class) diff --git a/src/scope/core/pipelines/schema.py b/src/scope/core/pipelines/schema.py index db11474e..d9aa0651 100644 --- a/src/scope/core/pipelines/schema.py +++ b/src/scope/core/pipelines/schema.py @@ -418,6 +418,49 @@ class PassthroughConfig(BasePipelineConfig): ) +class PersonaLiveConfig(BasePipelineConfig): + """Configuration for PersonaLive portrait animation pipeline. + + PersonaLive animates a reference portrait image using driving video frames. + It requires a reference image to be set once, then processes driving video frames. + """ + + pipeline_id: ClassVar[str] = "personalive" + pipeline_name: ClassVar[str] = "PersonaLive" + pipeline_description: ClassVar[str] = ( + "Real-time portrait animation from reference image and driving video" + ) + + # Mode support - video only (requires both reference image + driving video) + supported_modes: ClassVar[list[InputMode]] = ["video"] + default_mode: ClassVar[InputMode] = "video" + + # PersonaLive defaults + height: int = Field(default=512, ge=1, description="Output height in pixels") + width: int = Field(default=512, ge=1, description="Output width in pixels") + input_size: int | None = Field( + default=4, + description="Number of driving video frames per chunk", + ) + + # PersonaLive-specific parameters + temporal_window_size: int = Field( + default=4, + ge=1, + description="Temporal window size for processing", + ) + temporal_adaptive_step: int = Field( + default=4, + ge=1, + description="Temporal adaptive step for denoising", + ) + num_inference_steps: int = Field( + default=4, + ge=1, + description="Number of denoising steps (typically 4 for real-time)", + ) + + # Registry of pipeline config classes PIPELINE_CONFIGS: dict[str, type[BasePipelineConfig]] = { "longlive": LongLiveConfig, @@ -425,6 +468,7 @@ class PassthroughConfig(BasePipelineConfig): "krea-realtime-video": KreaRealtimeVideoConfig, "reward-forcing": RewardForcingConfig, "passthrough": PassthroughConfig, + "personalive": PersonaLiveConfig, } diff --git a/src/scope/server/app.py b/src/scope/server/app.py index 0f854727..35a24e76 100644 --- a/src/scope/server/app.py +++ b/src/scope/server/app.py @@ -37,6 +37,7 @@ IceCandidateRequest, IceServerConfig, IceServersResponse, + PersonaLiveReferenceResponse, PipelineLoadRequest, PipelineSchemasResponse, PipelineStatusResponse, @@ -311,6 +312,70 @@ async def get_pipeline_status( raise HTTPException(status_code=500, detail=str(e)) from e +@app.post("/api/v1/personalive/reference") +async def set_personalive_reference( + request: Request, + pipeline_manager: PipelineManager = Depends(get_pipeline_manager), +): + """Set the reference image for PersonaLive pipeline. + + This endpoint accepts a reference portrait image that will be animated + based on driving video frames. The image should be a face-centered portrait. + + The request body should be the raw image bytes with appropriate content-type + header (image/jpeg, image/png, etc.). + + This must be called after loading the PersonaLive pipeline and before + starting video streaming. + """ + from io import BytesIO + from PIL import Image + + try: + # Check if PersonaLive pipeline is loaded + status_info = await pipeline_manager.get_status_info_async() + if status_info.get("pipeline_id") != "personalive": + raise HTTPException( + status_code=400, + detail="PersonaLive pipeline must be loaded first" + ) + if status_info.get("status") != "loaded": + raise HTTPException( + status_code=400, + detail=f"Pipeline not ready. Current status: {status_info.get('status')}" + ) + + # Get image from request body + body = await request.body() + if not body: + raise HTTPException(status_code=400, detail="No image data provided") + + # Load image + try: + image = Image.open(BytesIO(body)).convert("RGB") + except Exception as e: + raise HTTPException(status_code=400, detail=f"Invalid image data: {e}") from e + + # Get pipeline and fuse reference + pipeline = pipeline_manager.get_pipeline() + if pipeline is None: + raise HTTPException(status_code=500, detail="Pipeline not available") + + # Fuse reference image + pipeline.fuse_reference(image) + + return PersonaLiveReferenceResponse( + success=True, + message="Reference image set successfully. Ready to process driving video." + ) + + except HTTPException: + raise + except Exception as e: + logger.error(f"Error setting PersonaLive reference: {e}") + raise HTTPException(status_code=500, detail=str(e)) from e + + @app.get("/api/v1/pipelines/schemas", response_model=PipelineSchemasResponse) async def get_pipeline_schemas(): """Get configuration schemas and defaults for all available pipelines. diff --git a/src/scope/server/download_models.py b/src/scope/server/download_models.py index 354b8a1d..15df3dc8 100644 --- a/src/scope/server/download_models.py +++ b/src/scope/server/download_models.py @@ -219,6 +219,73 @@ def download_reward_forcing_pipeline() -> None: ) +def download_personalive_pipeline() -> None: + """ + Download models for the PersonaLive portrait animation pipeline. + + Downloads from three sources: + 1. lambdalabs/sd-image-variations-diffusers - Base UNet + CLIP Image Encoder + (SD fine-tuned to accept CLIP image embeddings instead of text) + 2. stabilityai/sd-vae-ft-mse - VAE (improved fine-tuned version) + 3. huaichang/PersonaLive - PersonaLive-specific weights: + - denoising_unet.pth, reference_unet.pth + - motion_encoder.pth, motion_extractor.pth + - pose_guider.pth, temporal_module.pth + """ + models_root = ensure_models_dir() + personalive_root = models_root / "PersonaLive" / "pretrained_weights" + personalive_root.mkdir(parents=True, exist_ok=True) + + # 1. Download sd-image-variations-diffusers (base model with CLIP image encoder) + # This provides: image_encoder/, unet/, model_index.json + sd_variations_dst = personalive_root / "sd-image-variations-diffusers" + if not sd_variations_dst.exists(): + print("Downloading sd-image-variations-diffusers from lambdalabs...") + snapshot_download( + repo_id="lambdalabs/sd-image-variations-diffusers", + local_dir=str(sd_variations_dst), + local_dir_use_symlinks=False, + revision="v2.0", # Use v2.0 as recommended + ignore_patterns=["*.md", "*.txt"], + ) + print(f"[OK] Downloaded sd-image-variations-diffusers to: {sd_variations_dst}") + else: + print(f"[SKIP] sd-image-variations-diffusers already exists at: {sd_variations_dst}") + + # 2. Download sd-vae-ft-mse (improved VAE) + # This provides: config.json, diffusion_pytorch_model.safetensors + vae_dst = personalive_root / "sd-vae-ft-mse" + if not vae_dst.exists(): + print("Downloading sd-vae-ft-mse from stabilityai...") + snapshot_download( + repo_id="stabilityai/sd-vae-ft-mse", + local_dir=str(vae_dst), + local_dir_use_symlinks=False, + ignore_patterns=["*.md", "*.txt"], + ) + print(f"[OK] Downloaded sd-vae-ft-mse to: {vae_dst}") + else: + print(f"[SKIP] sd-vae-ft-mse already exists at: {vae_dst}") + + # 3. Download PersonaLive-specific weights + # This provides: denoising_unet.pth, reference_unet.pth, motion_encoder.pth, etc. + personalive_weights_dst = personalive_root / "personalive" + if not personalive_weights_dst.exists(): + print("Downloading PersonaLive weights from huaichang/PersonaLive...") + snapshot_download( + repo_id="huaichang/PersonaLive", + local_dir=str(personalive_weights_dst), + local_dir_use_symlinks=False, + # Only download the .pth weight files, not ONNX/TRT + allow_patterns=["*.pth"], + ) + print(f"[OK] Downloaded PersonaLive weights to: {personalive_weights_dst}") + else: + print(f"[SKIP] PersonaLive weights already exist at: {personalive_weights_dst}") + + print(f"\n[OK] PersonaLive pipeline models ready at: {personalive_root}") + + def download_models(pipeline_id: str | None = None) -> None: """ Download models. If pipeline_id is None, downloads all pipelines. @@ -226,7 +293,7 @@ def download_models(pipeline_id: str | None = None) -> None: Args: pipeline_id: Optional pipeline ID. Supports "streamdiffusionv2", "longlive", - "krea-realtime-video", or "reward-forcing". + "krea-realtime-video", "reward-forcing", or "personalive". If None, downloads all pipelines. """ if pipeline_id is None: @@ -235,6 +302,7 @@ def download_models(pipeline_id: str | None = None) -> None: download_longlive_pipeline() download_krea_realtime_video_pipeline() download_reward_forcing_pipeline() + download_personalive_pipeline() elif pipeline_id == "streamdiffusionv2": download_streamdiffusionv2_pipeline() elif pipeline_id == "longlive": @@ -243,9 +311,11 @@ def download_models(pipeline_id: str | None = None) -> None: download_krea_realtime_video_pipeline() elif pipeline_id == "reward-forcing": download_reward_forcing_pipeline() + elif pipeline_id == "personalive": + download_personalive_pipeline() else: raise ValueError( - f"Unknown pipeline: {pipeline_id}. Supported pipelines: streamdiffusionv2, longlive, krea-realtime-video, reward-forcing" + f"Unknown pipeline: {pipeline_id}. Supported pipelines: streamdiffusionv2, longlive, krea-realtime-video, reward-forcing, personalive" ) print("\nAll downloads complete.") @@ -266,6 +336,7 @@ def main(): python download_models.py --pipeline longlive python download_models.py --pipeline krea-realtime-video python download_models.py --pipeline reward-forcing + python download_models.py --pipeline personalive python download_models.py -p streamdiffusionv2 """, ) @@ -274,7 +345,7 @@ def main(): "-p", type=str, default=None, - help="Pipeline ID to download (e.g., 'streamdiffusionv2', 'longlive', 'krea-realtime-video', 'reward-forcing'). If not specified, downloads all pipelines.", + help="Pipeline ID to download (e.g., 'streamdiffusionv2', 'longlive', 'krea-realtime-video', 'reward-forcing', 'personalive'). If not specified, downloads all pipelines.", ) args = parser.parse_args() diff --git a/src/scope/server/models_config.py b/src/scope/server/models_config.py index c6ce6ace..cfc2b133 100644 --- a/src/scope/server/models_config.py +++ b/src/scope/server/models_config.py @@ -123,6 +123,19 @@ def get_required_model_files(pipeline_id: str | None = None) -> list[Path]: models_dir / "Reward-Forcing-T2V-1.3B" / "rewardforcing.pt", ] + # personalive pipeline + if pipeline_id == "personalive": + return [ + models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "denoising_unet.pth", + models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "reference_unet.pth", + models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "motion_encoder.pth", + models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "motion_extractor.pth", + models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "pose_guider.pth", + models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "temporal_module.pth", + models_dir / "PersonaLive" / "pretrained_weights" / "sd-image-variations-diffusers" / "model_index.json", + models_dir / "PersonaLive" / "pretrained_weights" / "sd-vae-ft-mse" / "config.json", + ] + # Default: nothing is required return [] diff --git a/src/scope/server/patch_xformers.py b/src/scope/server/patch_xformers.py new file mode 100644 index 00000000..89281ca6 --- /dev/null +++ b/src/scope/server/patch_xformers.py @@ -0,0 +1,85 @@ +"""Patch xformers flash.py to set FLASH_VER_LAST to 2.8.3.""" + +import importlib.util +import sys +from pathlib import Path + + +def find_xformers_flash_file() -> Path | None: + """Find the xformers flash.py file in the installed package.""" + try: + import xformers + + # Get the package path + spec = importlib.util.find_spec("xformers") + if spec is None or spec.origin is None: + return None + + # xformers package location + package_path = Path(spec.origin).parent + flash_file = package_path / "ops" / "fmha" / "flash.py" + + if flash_file.exists(): + return flash_file + + return None + except ImportError: + print("Error: xformers is not installed", file=sys.stderr) + return None + + +def patch_xformers_flash() -> bool: + """Patch xformers flash.py to set FLASH_VER_LAST to 2.8.3.""" + flash_file = find_xformers_flash_file() + + if flash_file is None: + print("Error: Could not find xformers flash.py file", file=sys.stderr) + return False + + try: + # Read the file + content = flash_file.read_text(encoding="utf-8") + + # Check if already patched + if 'FLASH_VER_LAST = parse_version("2.8.3")' in content: + print(f"✓ xformers flash.py already patched: {flash_file}") + return True + + # Find and replace the FLASH_VER_LAST line + lines = content.splitlines() + patched = False + + for i, line in enumerate(lines): + # Look for FLASH_VER_LAST assignment + if "FLASH_VER_LAST" in line and "parse_version" in line: + # Preserve original indentation + indent = len(line) - len(line.lstrip()) + # Replace with our desired version, preserving indentation + lines[i] = ( + " " * indent + + 'FLASH_VER_LAST = parse_version("2.8.3") # last supported, inclusive' + ) + patched = True + break + + if not patched: + print( + "Warning: Could not find FLASH_VER_LAST line to patch", + file=sys.stderr, + ) + return False + + # Write the patched content back + flash_file.write_text("\n".join(lines) + "\n", encoding="utf-8") + print(f"✓ Successfully patched xformers flash.py: {flash_file}") + return True + + except Exception as e: + print(f"Error patching xformers flash.py: {e}", file=sys.stderr) + return False + + +def main() -> None: + """Main entry point for the patch command.""" + success = patch_xformers_flash() + sys.exit(0 if success else 1) diff --git a/src/scope/server/pipeline_manager.py b/src/scope/server/pipeline_manager.py index 884bd43b..fdb5b954 100644 --- a/src/scope/server/pipeline_manager.py +++ b/src/scope/server/pipeline_manager.py @@ -491,6 +491,38 @@ def _load_pipeline_implementation( logger.info("RewardForcing pipeline initialized") return pipeline + elif pipeline_id == "personalive": + from scope.core.pipelines import PersonaLivePipeline + + from .models_config import get_models_dir + + config = OmegaConf.create( + { + "model_dir": str(get_models_dir()), + } + ) + + # Apply load parameters (resolution, seed) to config + height = 512 + width = 512 + seed = 42 + if load_params: + height = load_params.get("height", 512) + width = load_params.get("width", 512) + seed = load_params.get("seed", 42) + + config["height"] = height + config["width"] = width + config["seed"] = seed + + pipeline = PersonaLivePipeline( + config, + device=torch.device("cuda"), + dtype=torch.float16, # PersonaLive uses fp16 + ) + logger.info("PersonaLive pipeline initialized") + return pipeline + else: raise ValueError(f"Invalid pipeline ID: {pipeline_id}") diff --git a/src/scope/server/schema.py b/src/scope/server/schema.py index 1916dab7..7052358d 100644 --- a/src/scope/server/schema.py +++ b/src/scope/server/schema.py @@ -8,6 +8,7 @@ from scope.core.pipelines.schema import ( KreaRealtimeVideoConfig, LongLiveConfig, + PersonaLiveConfig, StreamDiffusionV2Config, ) from scope.core.pipelines.utils import Quantization @@ -360,6 +361,40 @@ class KreaRealtimeVideoLoadParams(LoRAEnabledLoadParams): ) +class PersonaLiveLoadParams(PipelineLoadParams): + """Load parameters for PersonaLive portrait animation pipeline. + + Defaults are derived from PersonaLiveConfig to ensure consistency. + Note: A reference image must be uploaded via /api/v1/personalive/reference + before the pipeline can process driving video frames. + """ + + height: int = Field( + default=PersonaLiveConfig.model_fields["height"].default, + description="Target video height", + ge=256, + le=1024, + ) + width: int = Field( + default=PersonaLiveConfig.model_fields["width"].default, + description="Target video width", + ge=256, + le=1024, + ) + seed: int = Field( + default=42, # Default seed from BasePipelineConfig + description="Random seed for generation", + ge=0, + ) + + +class PersonaLiveReferenceResponse(BaseModel): + """Response after setting PersonaLive reference image.""" + + success: bool = Field(..., description="Whether the reference image was set successfully") + message: str = Field(..., description="Status message") + + class PipelineLoadRequest(BaseModel): """Pipeline load request schema.""" @@ -371,6 +406,7 @@ class PipelineLoadRequest(BaseModel): | PassthroughLoadParams | LongLiveLoadParams | KreaRealtimeVideoLoadParams + | PersonaLiveLoadParams | None ) = Field(default=None, description="Pipeline-specific load parameters") diff --git a/uv.lock b/uv.lock index 82d4a89b..24068139 100644 --- a/uv.lock +++ b/uv.lock @@ -2,15 +2,35 @@ version = 1 revision = 3 requires-python = ">=3.10.12" resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version < '3.11' and sys_platform == 'win32'", - "python_full_version >= '3.12' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.11.*' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] + +[[package]] +name = "absl-py" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/10/2a/c93173ffa1b39c1d0395b7e842bbdc62e556ca9d8d3b5572926f3e4ca752/absl_py-2.3.1.tar.gz", hash = "sha256:a97820526f7fbfd2ec1bce83f3f25e3a14840dac0d8e02a0b71cd75db3f77fc9", size = 116588, upload-time = "2025-07-03T09:31:44.05Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8f/aa/ba0014cc4659328dc818a28827be78e6d97312ab0cb98105a770924dc11e/absl_py-2.3.1-py3-none-any.whl", hash = "sha256:eeecf07f0c2a93ace0772c92e596ace6d3d3996c042b2128459aaae2a76de11d", size = 135811, upload-time = "2025-07-03T09:31:42.253Z" }, ] [[package]] @@ -19,8 +39,7 @@ version = "1.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy" }, { name = "packaging" }, { name = "psutil" }, { name = "pyyaml" }, @@ -551,6 +570,179 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "contourpy" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/54/eb9bfc647b19f2009dd5c7f5ec51c4e6ca831725f1aea7a993034f483147/contourpy-1.3.2.tar.gz", hash = "sha256:b6945942715a034c671b7fc54f9588126b0b8bf23db2696e3ca8328f3ff0ab54", size = 13466130, upload-time = "2025-04-15T17:47:53.79Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/a3/da4153ec8fe25d263aa48c1a4cbde7f49b59af86f0b6f7862788c60da737/contourpy-1.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:ba38e3f9f330af820c4b27ceb4b9c7feee5fe0493ea53a8720f4792667465934", size = 268551, upload-time = "2025-04-15T17:34:46.581Z" }, + { url = "https://files.pythonhosted.org/packages/2f/6c/330de89ae1087eb622bfca0177d32a7ece50c3ef07b28002de4757d9d875/contourpy-1.3.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:dc41ba0714aa2968d1f8674ec97504a8f7e334f48eeacebcaa6256213acb0989", size = 253399, upload-time = "2025-04-15T17:34:51.427Z" }, + { url = "https://files.pythonhosted.org/packages/c1/bd/20c6726b1b7f81a8bee5271bed5c165f0a8e1f572578a9d27e2ccb763cb2/contourpy-1.3.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9be002b31c558d1ddf1b9b415b162c603405414bacd6932d031c5b5a8b757f0d", size = 312061, upload-time = "2025-04-15T17:34:55.961Z" }, + { url = "https://files.pythonhosted.org/packages/22/fc/a9665c88f8a2473f823cf1ec601de9e5375050f1958cbb356cdf06ef1ab6/contourpy-1.3.2-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8d2e74acbcba3bfdb6d9d8384cdc4f9260cae86ed9beee8bd5f54fee49a430b9", size = 351956, upload-time = "2025-04-15T17:35:00.992Z" }, + { url = "https://files.pythonhosted.org/packages/25/eb/9f0a0238f305ad8fb7ef42481020d6e20cf15e46be99a1fcf939546a177e/contourpy-1.3.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e259bced5549ac64410162adc973c5e2fb77f04df4a439d00b478e57a0e65512", size = 320872, upload-time = "2025-04-15T17:35:06.177Z" }, + { url = "https://files.pythonhosted.org/packages/32/5c/1ee32d1c7956923202f00cf8d2a14a62ed7517bdc0ee1e55301227fc273c/contourpy-1.3.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ad687a04bc802cbe8b9c399c07162a3c35e227e2daccf1668eb1f278cb698631", size = 325027, upload-time = "2025-04-15T17:35:11.244Z" }, + { url = "https://files.pythonhosted.org/packages/83/bf/9baed89785ba743ef329c2b07fd0611d12bfecbedbdd3eeecf929d8d3b52/contourpy-1.3.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cdd22595308f53ef2f891040ab2b93d79192513ffccbd7fe19be7aa773a5e09f", size = 1306641, upload-time = "2025-04-15T17:35:26.701Z" }, + { url = "https://files.pythonhosted.org/packages/d4/cc/74e5e83d1e35de2d28bd97033426b450bc4fd96e092a1f7a63dc7369b55d/contourpy-1.3.2-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b4f54d6a2defe9f257327b0f243612dd051cc43825587520b1bf74a31e2f6ef2", size = 1374075, upload-time = "2025-04-15T17:35:43.204Z" }, + { url = "https://files.pythonhosted.org/packages/0c/42/17f3b798fd5e033b46a16f8d9fcb39f1aba051307f5ebf441bad1ecf78f8/contourpy-1.3.2-cp310-cp310-win32.whl", hash = "sha256:f939a054192ddc596e031e50bb13b657ce318cf13d264f095ce9db7dc6ae81c0", size = 177534, upload-time = "2025-04-15T17:35:46.554Z" }, + { url = "https://files.pythonhosted.org/packages/54/ec/5162b8582f2c994721018d0c9ece9dc6ff769d298a8ac6b6a652c307e7df/contourpy-1.3.2-cp310-cp310-win_amd64.whl", hash = "sha256:c440093bbc8fc21c637c03bafcbef95ccd963bc6e0514ad887932c18ca2a759a", size = 221188, upload-time = "2025-04-15T17:35:50.064Z" }, + { url = "https://files.pythonhosted.org/packages/b3/b9/ede788a0b56fc5b071639d06c33cb893f68b1178938f3425debebe2dab78/contourpy-1.3.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6a37a2fb93d4df3fc4c0e363ea4d16f83195fc09c891bc8ce072b9d084853445", size = 269636, upload-time = "2025-04-15T17:35:54.473Z" }, + { url = "https://files.pythonhosted.org/packages/e6/75/3469f011d64b8bbfa04f709bfc23e1dd71be54d05b1b083be9f5b22750d1/contourpy-1.3.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b7cd50c38f500bbcc9b6a46643a40e0913673f869315d8e70de0438817cb7773", size = 254636, upload-time = "2025-04-15T17:35:58.283Z" }, + { url = "https://files.pythonhosted.org/packages/8d/2f/95adb8dae08ce0ebca4fd8e7ad653159565d9739128b2d5977806656fcd2/contourpy-1.3.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d6658ccc7251a4433eebd89ed2672c2ed96fba367fd25ca9512aa92a4b46c4f1", size = 313053, upload-time = "2025-04-15T17:36:03.235Z" }, + { url = "https://files.pythonhosted.org/packages/c3/a6/8ccf97a50f31adfa36917707fe39c9a0cbc24b3bbb58185577f119736cc9/contourpy-1.3.2-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:70771a461aaeb335df14deb6c97439973d253ae70660ca085eec25241137ef43", size = 352985, upload-time = "2025-04-15T17:36:08.275Z" }, + { url = "https://files.pythonhosted.org/packages/1d/b6/7925ab9b77386143f39d9c3243fdd101621b4532eb126743201160ffa7e6/contourpy-1.3.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:65a887a6e8c4cd0897507d814b14c54a8c2e2aa4ac9f7686292f9769fcf9a6ab", size = 323750, upload-time = "2025-04-15T17:36:13.29Z" }, + { url = "https://files.pythonhosted.org/packages/c2/f3/20c5d1ef4f4748e52d60771b8560cf00b69d5c6368b5c2e9311bcfa2a08b/contourpy-1.3.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3859783aefa2b8355697f16642695a5b9792e7a46ab86da1118a4a23a51a33d7", size = 326246, upload-time = "2025-04-15T17:36:18.329Z" }, + { url = "https://files.pythonhosted.org/packages/8c/e5/9dae809e7e0b2d9d70c52b3d24cba134dd3dad979eb3e5e71f5df22ed1f5/contourpy-1.3.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:eab0f6db315fa4d70f1d8ab514e527f0366ec021ff853d7ed6a2d33605cf4b83", size = 1308728, upload-time = "2025-04-15T17:36:33.878Z" }, + { url = "https://files.pythonhosted.org/packages/e2/4a/0058ba34aeea35c0b442ae61a4f4d4ca84d6df8f91309bc2d43bb8dd248f/contourpy-1.3.2-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d91a3ccc7fea94ca0acab82ceb77f396d50a1f67412efe4c526f5d20264e6ecd", size = 1375762, upload-time = "2025-04-15T17:36:51.295Z" }, + { url = "https://files.pythonhosted.org/packages/09/33/7174bdfc8b7767ef2c08ed81244762d93d5c579336fc0b51ca57b33d1b80/contourpy-1.3.2-cp311-cp311-win32.whl", hash = "sha256:1c48188778d4d2f3d48e4643fb15d8608b1d01e4b4d6b0548d9b336c28fc9b6f", size = 178196, upload-time = "2025-04-15T17:36:55.002Z" }, + { url = "https://files.pythonhosted.org/packages/5e/fe/4029038b4e1c4485cef18e480b0e2cd2d755448bb071eb9977caac80b77b/contourpy-1.3.2-cp311-cp311-win_amd64.whl", hash = "sha256:5ebac872ba09cb8f2131c46b8739a7ff71de28a24c869bcad554477eb089a878", size = 222017, upload-time = "2025-04-15T17:36:58.576Z" }, + { url = "https://files.pythonhosted.org/packages/34/f7/44785876384eff370c251d58fd65f6ad7f39adce4a093c934d4a67a7c6b6/contourpy-1.3.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:4caf2bcd2969402bf77edc4cb6034c7dd7c0803213b3523f111eb7460a51b8d2", size = 271580, upload-time = "2025-04-15T17:37:03.105Z" }, + { url = "https://files.pythonhosted.org/packages/93/3b/0004767622a9826ea3d95f0e9d98cd8729015768075d61f9fea8eeca42a8/contourpy-1.3.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:82199cb78276249796419fe36b7386bd8d2cc3f28b3bc19fe2454fe2e26c4c15", size = 255530, upload-time = "2025-04-15T17:37:07.026Z" }, + { url = "https://files.pythonhosted.org/packages/e7/bb/7bd49e1f4fa805772d9fd130e0d375554ebc771ed7172f48dfcd4ca61549/contourpy-1.3.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:106fab697af11456fcba3e352ad50effe493a90f893fca6c2ca5c033820cea92", size = 307688, upload-time = "2025-04-15T17:37:11.481Z" }, + { url = "https://files.pythonhosted.org/packages/fc/97/e1d5dbbfa170725ef78357a9a0edc996b09ae4af170927ba8ce977e60a5f/contourpy-1.3.2-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:d14f12932a8d620e307f715857107b1d1845cc44fdb5da2bc8e850f5ceba9f87", size = 347331, upload-time = "2025-04-15T17:37:18.212Z" }, + { url = "https://files.pythonhosted.org/packages/6f/66/e69e6e904f5ecf6901be3dd16e7e54d41b6ec6ae3405a535286d4418ffb4/contourpy-1.3.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:532fd26e715560721bb0d5fc7610fce279b3699b018600ab999d1be895b09415", size = 318963, upload-time = "2025-04-15T17:37:22.76Z" }, + { url = "https://files.pythonhosted.org/packages/a8/32/b8a1c8965e4f72482ff2d1ac2cd670ce0b542f203c8e1d34e7c3e6925da7/contourpy-1.3.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f26b383144cf2d2c29f01a1e8170f50dacf0eac02d64139dcd709a8ac4eb3cfe", size = 323681, upload-time = "2025-04-15T17:37:33.001Z" }, + { url = "https://files.pythonhosted.org/packages/30/c6/12a7e6811d08757c7162a541ca4c5c6a34c0f4e98ef2b338791093518e40/contourpy-1.3.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:c49f73e61f1f774650a55d221803b101d966ca0c5a2d6d5e4320ec3997489441", size = 1308674, upload-time = "2025-04-15T17:37:48.64Z" }, + { url = "https://files.pythonhosted.org/packages/2a/8a/bebe5a3f68b484d3a2b8ffaf84704b3e343ef1addea528132ef148e22b3b/contourpy-1.3.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:3d80b2c0300583228ac98d0a927a1ba6a2ba6b8a742463c564f1d419ee5b211e", size = 1380480, upload-time = "2025-04-15T17:38:06.7Z" }, + { url = "https://files.pythonhosted.org/packages/34/db/fcd325f19b5978fb509a7d55e06d99f5f856294c1991097534360b307cf1/contourpy-1.3.2-cp312-cp312-win32.whl", hash = "sha256:90df94c89a91b7362e1142cbee7568f86514412ab8a2c0d0fca72d7e91b62912", size = 178489, upload-time = "2025-04-15T17:38:10.338Z" }, + { url = "https://files.pythonhosted.org/packages/01/c8/fadd0b92ffa7b5eb5949bf340a63a4a496a6930a6c37a7ba0f12acb076d6/contourpy-1.3.2-cp312-cp312-win_amd64.whl", hash = "sha256:8c942a01d9163e2e5cfb05cb66110121b8d07ad438a17f9e766317bcb62abf73", size = 223042, upload-time = "2025-04-15T17:38:14.239Z" }, + { url = "https://files.pythonhosted.org/packages/2e/61/5673f7e364b31e4e7ef6f61a4b5121c5f170f941895912f773d95270f3a2/contourpy-1.3.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:de39db2604ae755316cb5967728f4bea92685884b1e767b7c24e983ef5f771cb", size = 271630, upload-time = "2025-04-15T17:38:19.142Z" }, + { url = "https://files.pythonhosted.org/packages/ff/66/a40badddd1223822c95798c55292844b7e871e50f6bfd9f158cb25e0bd39/contourpy-1.3.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3f9e896f447c5c8618f1edb2bafa9a4030f22a575ec418ad70611450720b5b08", size = 255670, upload-time = "2025-04-15T17:38:23.688Z" }, + { url = "https://files.pythonhosted.org/packages/1e/c7/cf9fdee8200805c9bc3b148f49cb9482a4e3ea2719e772602a425c9b09f8/contourpy-1.3.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:71e2bd4a1c4188f5c2b8d274da78faab884b59df20df63c34f74aa1813c4427c", size = 306694, upload-time = "2025-04-15T17:38:28.238Z" }, + { url = "https://files.pythonhosted.org/packages/dd/e7/ccb9bec80e1ba121efbffad7f38021021cda5be87532ec16fd96533bb2e0/contourpy-1.3.2-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:de425af81b6cea33101ae95ece1f696af39446db9682a0b56daaa48cfc29f38f", size = 345986, upload-time = "2025-04-15T17:38:33.502Z" }, + { url = "https://files.pythonhosted.org/packages/dc/49/ca13bb2da90391fa4219fdb23b078d6065ada886658ac7818e5441448b78/contourpy-1.3.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:977e98a0e0480d3fe292246417239d2d45435904afd6d7332d8455981c408b85", size = 318060, upload-time = "2025-04-15T17:38:38.672Z" }, + { url = "https://files.pythonhosted.org/packages/c8/65/5245ce8c548a8422236c13ffcdcdada6a2a812c361e9e0c70548bb40b661/contourpy-1.3.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:434f0adf84911c924519d2b08fc10491dd282b20bdd3fa8f60fd816ea0b48841", size = 322747, upload-time = "2025-04-15T17:38:43.712Z" }, + { url = "https://files.pythonhosted.org/packages/72/30/669b8eb48e0a01c660ead3752a25b44fdb2e5ebc13a55782f639170772f9/contourpy-1.3.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:c66c4906cdbc50e9cba65978823e6e00b45682eb09adbb78c9775b74eb222422", size = 1308895, upload-time = "2025-04-15T17:39:00.224Z" }, + { url = "https://files.pythonhosted.org/packages/05/5a/b569f4250decee6e8d54498be7bdf29021a4c256e77fe8138c8319ef8eb3/contourpy-1.3.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8b7fc0cd78ba2f4695fd0a6ad81a19e7e3ab825c31b577f384aa9d7817dc3bef", size = 1379098, upload-time = "2025-04-15T17:43:29.649Z" }, + { url = "https://files.pythonhosted.org/packages/19/ba/b227c3886d120e60e41b28740ac3617b2f2b971b9f601c835661194579f1/contourpy-1.3.2-cp313-cp313-win32.whl", hash = "sha256:15ce6ab60957ca74cff444fe66d9045c1fd3e92c8936894ebd1f3eef2fff075f", size = 178535, upload-time = "2025-04-15T17:44:44.532Z" }, + { url = "https://files.pythonhosted.org/packages/12/6e/2fed56cd47ca739b43e892707ae9a13790a486a3173be063681ca67d2262/contourpy-1.3.2-cp313-cp313-win_amd64.whl", hash = "sha256:e1578f7eafce927b168752ed7e22646dad6cd9bca673c60bff55889fa236ebf9", size = 223096, upload-time = "2025-04-15T17:44:48.194Z" }, + { url = "https://files.pythonhosted.org/packages/54/4c/e76fe2a03014a7c767d79ea35c86a747e9325537a8b7627e0e5b3ba266b4/contourpy-1.3.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0475b1f6604896bc7c53bb070e355e9321e1bc0d381735421a2d2068ec56531f", size = 285090, upload-time = "2025-04-15T17:43:34.084Z" }, + { url = "https://files.pythonhosted.org/packages/7b/e2/5aba47debd55d668e00baf9651b721e7733975dc9fc27264a62b0dd26eb8/contourpy-1.3.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:c85bb486e9be652314bb5b9e2e3b0d1b2e643d5eec4992c0fbe8ac71775da739", size = 268643, upload-time = "2025-04-15T17:43:38.626Z" }, + { url = "https://files.pythonhosted.org/packages/a1/37/cd45f1f051fe6230f751cc5cdd2728bb3a203f5619510ef11e732109593c/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:745b57db7758f3ffc05a10254edd3182a2a83402a89c00957a8e8a22f5582823", size = 310443, upload-time = "2025-04-15T17:43:44.522Z" }, + { url = "https://files.pythonhosted.org/packages/8b/a2/36ea6140c306c9ff6dd38e3bcec80b3b018474ef4d17eb68ceecd26675f4/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:970e9173dbd7eba9b4e01aab19215a48ee5dd3f43cef736eebde064a171f89a5", size = 349865, upload-time = "2025-04-15T17:43:49.545Z" }, + { url = "https://files.pythonhosted.org/packages/95/b7/2fc76bc539693180488f7b6cc518da7acbbb9e3b931fd9280504128bf956/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c6c4639a9c22230276b7bffb6a850dfc8258a2521305e1faefe804d006b2e532", size = 321162, upload-time = "2025-04-15T17:43:54.203Z" }, + { url = "https://files.pythonhosted.org/packages/f4/10/76d4f778458b0aa83f96e59d65ece72a060bacb20cfbee46cf6cd5ceba41/contourpy-1.3.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:cc829960f34ba36aad4302e78eabf3ef16a3a100863f0d4eeddf30e8a485a03b", size = 327355, upload-time = "2025-04-15T17:44:01.025Z" }, + { url = "https://files.pythonhosted.org/packages/43/a3/10cf483ea683f9f8ab096c24bad3cce20e0d1dd9a4baa0e2093c1c962d9d/contourpy-1.3.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:d32530b534e986374fc19eaa77fcb87e8a99e5431499949b828312bdcd20ac52", size = 1307935, upload-time = "2025-04-15T17:44:17.322Z" }, + { url = "https://files.pythonhosted.org/packages/78/73/69dd9a024444489e22d86108e7b913f3528f56cfc312b5c5727a44188471/contourpy-1.3.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:e298e7e70cf4eb179cc1077be1c725b5fd131ebc81181bf0c03525c8abc297fd", size = 1372168, upload-time = "2025-04-15T17:44:33.43Z" }, + { url = "https://files.pythonhosted.org/packages/0f/1b/96d586ccf1b1a9d2004dd519b25fbf104a11589abfd05484ff12199cca21/contourpy-1.3.2-cp313-cp313t-win32.whl", hash = "sha256:d0e589ae0d55204991450bb5c23f571c64fe43adaa53f93fc902a84c96f52fe1", size = 189550, upload-time = "2025-04-15T17:44:37.092Z" }, + { url = "https://files.pythonhosted.org/packages/b0/e6/6000d0094e8a5e32ad62591c8609e269febb6e4db83a1c75ff8868b42731/contourpy-1.3.2-cp313-cp313t-win_amd64.whl", hash = "sha256:78e9253c3de756b3f6a5174d024c4835acd59eb3f8e2ca13e775dbffe1558f69", size = 238214, upload-time = "2025-04-15T17:44:40.827Z" }, + { url = "https://files.pythonhosted.org/packages/33/05/b26e3c6ecc05f349ee0013f0bb850a761016d89cec528a98193a48c34033/contourpy-1.3.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:fd93cc7f3139b6dd7aab2f26a90dde0aa9fc264dbf70f6740d498a70b860b82c", size = 265681, upload-time = "2025-04-15T17:44:59.314Z" }, + { url = "https://files.pythonhosted.org/packages/2b/25/ac07d6ad12affa7d1ffed11b77417d0a6308170f44ff20fa1d5aa6333f03/contourpy-1.3.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:107ba8a6a7eec58bb475329e6d3b95deba9440667c4d62b9b6063942b61d7f16", size = 315101, upload-time = "2025-04-15T17:45:04.165Z" }, + { url = "https://files.pythonhosted.org/packages/8f/4d/5bb3192bbe9d3f27e3061a6a8e7733c9120e203cb8515767d30973f71030/contourpy-1.3.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:ded1706ed0c1049224531b81128efbd5084598f18d8a2d9efae833edbd2b40ad", size = 220599, upload-time = "2025-04-15T17:45:08.456Z" }, + { url = "https://files.pythonhosted.org/packages/ff/c0/91f1215d0d9f9f343e4773ba6c9b89e8c0cc7a64a6263f21139da639d848/contourpy-1.3.2-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:5f5964cdad279256c084b69c3f412b7801e15356b16efa9d78aa974041903da0", size = 266807, upload-time = "2025-04-15T17:45:15.535Z" }, + { url = "https://files.pythonhosted.org/packages/d4/79/6be7e90c955c0487e7712660d6cead01fa17bff98e0ea275737cc2bc8e71/contourpy-1.3.2-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:49b65a95d642d4efa8f64ba12558fcb83407e58a2dfba9d796d77b63ccfcaff5", size = 318729, upload-time = "2025-04-15T17:45:20.166Z" }, + { url = "https://files.pythonhosted.org/packages/87/68/7f46fb537958e87427d98a4074bcde4b67a70b04900cfc5ce29bc2f556c1/contourpy-1.3.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:8c5acb8dddb0752bf252e01a3035b21443158910ac16a3b0d20e7fed7d534ce5", size = 221791, upload-time = "2025-04-15T17:45:24.794Z" }, +] + +[[package]] +name = "contourpy" +version = "1.3.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/2e/c4390a31919d8a78b90e8ecf87cd4b4c4f05a5b48d05ec17db8e5404c6f4/contourpy-1.3.3-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:709a48ef9a690e1343202916450bc48b9e51c049b089c7f79a267b46cffcdaa1", size = 288773, upload-time = "2025-07-26T12:01:02.277Z" }, + { url = "https://files.pythonhosted.org/packages/0d/44/c4b0b6095fef4dc9c420e041799591e3b63e9619e3044f7f4f6c21c0ab24/contourpy-1.3.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:23416f38bfd74d5d28ab8429cc4d63fa67d5068bd711a85edb1c3fb0c3e2f381", size = 270149, upload-time = "2025-07-26T12:01:04.072Z" }, + { url = "https://files.pythonhosted.org/packages/30/2e/dd4ced42fefac8470661d7cb7e264808425e6c5d56d175291e93890cce09/contourpy-1.3.3-cp311-cp311-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:929ddf8c4c7f348e4c0a5a3a714b5c8542ffaa8c22954862a46ca1813b667ee7", size = 329222, upload-time = "2025-07-26T12:01:05.688Z" }, + { url = "https://files.pythonhosted.org/packages/f2/74/cc6ec2548e3d276c71389ea4802a774b7aa3558223b7bade3f25787fafc2/contourpy-1.3.3-cp311-cp311-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9e999574eddae35f1312c2b4b717b7885d4edd6cb46700e04f7f02db454e67c1", size = 377234, upload-time = "2025-07-26T12:01:07.054Z" }, + { url = "https://files.pythonhosted.org/packages/03/b3/64ef723029f917410f75c09da54254c5f9ea90ef89b143ccadb09df14c15/contourpy-1.3.3-cp311-cp311-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:0bf67e0e3f482cb69779dd3061b534eb35ac9b17f163d851e2a547d56dba0a3a", size = 380555, upload-time = "2025-07-26T12:01:08.801Z" }, + { url = "https://files.pythonhosted.org/packages/5f/4b/6157f24ca425b89fe2eb7e7be642375711ab671135be21e6faa100f7448c/contourpy-1.3.3-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:51e79c1f7470158e838808d4a996fa9bac72c498e93d8ebe5119bc1e6becb0db", size = 355238, upload-time = "2025-07-26T12:01:10.319Z" }, + { url = "https://files.pythonhosted.org/packages/98/56/f914f0dd678480708a04cfd2206e7c382533249bc5001eb9f58aa693e200/contourpy-1.3.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:598c3aaece21c503615fd59c92a3598b428b2f01bfb4b8ca9c4edeecc2438620", size = 1326218, upload-time = "2025-07-26T12:01:12.659Z" }, + { url = "https://files.pythonhosted.org/packages/fb/d7/4a972334a0c971acd5172389671113ae82aa7527073980c38d5868ff1161/contourpy-1.3.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:322ab1c99b008dad206d406bb61d014cf0174df491ae9d9d0fac6a6fda4f977f", size = 1392867, upload-time = "2025-07-26T12:01:15.533Z" }, + { url = "https://files.pythonhosted.org/packages/75/3e/f2cc6cd56dc8cff46b1a56232eabc6feea52720083ea71ab15523daab796/contourpy-1.3.3-cp311-cp311-win32.whl", hash = "sha256:fd907ae12cd483cd83e414b12941c632a969171bf90fc937d0c9f268a31cafff", size = 183677, upload-time = "2025-07-26T12:01:17.088Z" }, + { url = "https://files.pythonhosted.org/packages/98/4b/9bd370b004b5c9d8045c6c33cf65bae018b27aca550a3f657cdc99acdbd8/contourpy-1.3.3-cp311-cp311-win_amd64.whl", hash = "sha256:3519428f6be58431c56581f1694ba8e50626f2dd550af225f82fb5f5814d2a42", size = 225234, upload-time = "2025-07-26T12:01:18.256Z" }, + { url = "https://files.pythonhosted.org/packages/d9/b6/71771e02c2e004450c12b1120a5f488cad2e4d5b590b1af8bad060360fe4/contourpy-1.3.3-cp311-cp311-win_arm64.whl", hash = "sha256:15ff10bfada4bf92ec8b31c62bf7c1834c244019b4a33095a68000d7075df470", size = 193123, upload-time = "2025-07-26T12:01:19.848Z" }, + { url = "https://files.pythonhosted.org/packages/be/45/adfee365d9ea3d853550b2e735f9d66366701c65db7855cd07621732ccfc/contourpy-1.3.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b08a32ea2f8e42cf1d4be3169a98dd4be32bafe4f22b6c4cb4ba810fa9e5d2cb", size = 293419, upload-time = "2025-07-26T12:01:21.16Z" }, + { url = "https://files.pythonhosted.org/packages/53/3e/405b59cfa13021a56bba395a6b3aca8cec012b45bf177b0eaf7a202cde2c/contourpy-1.3.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:556dba8fb6f5d8742f2923fe9457dbdd51e1049c4a43fd3986a0b14a1d815fc6", size = 273979, upload-time = "2025-07-26T12:01:22.448Z" }, + { url = "https://files.pythonhosted.org/packages/d4/1c/a12359b9b2ca3a845e8f7f9ac08bdf776114eb931392fcad91743e2ea17b/contourpy-1.3.3-cp312-cp312-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92d9abc807cf7d0e047b95ca5d957cf4792fcd04e920ca70d48add15c1a90ea7", size = 332653, upload-time = "2025-07-26T12:01:24.155Z" }, + { url = "https://files.pythonhosted.org/packages/63/12/897aeebfb475b7748ea67b61e045accdfcf0d971f8a588b67108ed7f5512/contourpy-1.3.3-cp312-cp312-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:b2e8faa0ed68cb29af51edd8e24798bb661eac3bd9f65420c1887b6ca89987c8", size = 379536, upload-time = "2025-07-26T12:01:25.91Z" }, + { url = "https://files.pythonhosted.org/packages/43/8a/a8c584b82deb248930ce069e71576fc09bd7174bbd35183b7943fb1064fd/contourpy-1.3.3-cp312-cp312-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:626d60935cf668e70a5ce6ff184fd713e9683fb458898e4249b63be9e28286ea", size = 384397, upload-time = "2025-07-26T12:01:27.152Z" }, + { url = "https://files.pythonhosted.org/packages/cc/8f/ec6289987824b29529d0dfda0d74a07cec60e54b9c92f3c9da4c0ac732de/contourpy-1.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d00e655fcef08aba35ec9610536bfe90267d7ab5ba944f7032549c55a146da1", size = 362601, upload-time = "2025-07-26T12:01:28.808Z" }, + { url = "https://files.pythonhosted.org/packages/05/0a/a3fe3be3ee2dceb3e615ebb4df97ae6f3828aa915d3e10549ce016302bd1/contourpy-1.3.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:451e71b5a7d597379ef572de31eeb909a87246974d960049a9848c3bc6c41bf7", size = 1331288, upload-time = "2025-07-26T12:01:31.198Z" }, + { url = "https://files.pythonhosted.org/packages/33/1d/acad9bd4e97f13f3e2b18a3977fe1b4a37ecf3d38d815333980c6c72e963/contourpy-1.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:459c1f020cd59fcfe6650180678a9993932d80d44ccde1fa1868977438f0b411", size = 1403386, upload-time = "2025-07-26T12:01:33.947Z" }, + { url = "https://files.pythonhosted.org/packages/cf/8f/5847f44a7fddf859704217a99a23a4f6417b10e5ab1256a179264561540e/contourpy-1.3.3-cp312-cp312-win32.whl", hash = "sha256:023b44101dfe49d7d53932be418477dba359649246075c996866106da069af69", size = 185018, upload-time = "2025-07-26T12:01:35.64Z" }, + { url = "https://files.pythonhosted.org/packages/19/e8/6026ed58a64563186a9ee3f29f41261fd1828f527dd93d33b60feca63352/contourpy-1.3.3-cp312-cp312-win_amd64.whl", hash = "sha256:8153b8bfc11e1e4d75bcb0bff1db232f9e10b274e0929de9d608027e0d34ff8b", size = 226567, upload-time = "2025-07-26T12:01:36.804Z" }, + { url = "https://files.pythonhosted.org/packages/d1/e2/f05240d2c39a1ed228d8328a78b6f44cd695f7ef47beb3e684cf93604f86/contourpy-1.3.3-cp312-cp312-win_arm64.whl", hash = "sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc", size = 193655, upload-time = "2025-07-26T12:01:37.999Z" }, + { url = "https://files.pythonhosted.org/packages/68/35/0167aad910bbdb9599272bd96d01a9ec6852f36b9455cf2ca67bd4cc2d23/contourpy-1.3.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:177fb367556747a686509d6fef71d221a4b198a3905fe824430e5ea0fda54eb5", size = 293257, upload-time = "2025-07-26T12:01:39.367Z" }, + { url = "https://files.pythonhosted.org/packages/96/e4/7adcd9c8362745b2210728f209bfbcf7d91ba868a2c5f40d8b58f54c509b/contourpy-1.3.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d002b6f00d73d69333dac9d0b8d5e84d9724ff9ef044fd63c5986e62b7c9e1b1", size = 274034, upload-time = "2025-07-26T12:01:40.645Z" }, + { url = "https://files.pythonhosted.org/packages/73/23/90e31ceeed1de63058a02cb04b12f2de4b40e3bef5e082a7c18d9c8ae281/contourpy-1.3.3-cp313-cp313-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:348ac1f5d4f1d66d3322420f01d42e43122f43616e0f194fc1c9f5d830c5b286", size = 334672, upload-time = "2025-07-26T12:01:41.942Z" }, + { url = "https://files.pythonhosted.org/packages/ed/93/b43d8acbe67392e659e1d984700e79eb67e2acb2bd7f62012b583a7f1b55/contourpy-1.3.3-cp313-cp313-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:655456777ff65c2c548b7c454af9c6f33f16c8884f11083244b5819cc214f1b5", size = 381234, upload-time = "2025-07-26T12:01:43.499Z" }, + { url = "https://files.pythonhosted.org/packages/46/3b/bec82a3ea06f66711520f75a40c8fc0b113b2a75edb36aa633eb11c4f50f/contourpy-1.3.3-cp313-cp313-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:644a6853d15b2512d67881586bd03f462c7ab755db95f16f14d7e238f2852c67", size = 385169, upload-time = "2025-07-26T12:01:45.219Z" }, + { url = "https://files.pythonhosted.org/packages/4b/32/e0f13a1c5b0f8572d0ec6ae2f6c677b7991fafd95da523159c19eff0696a/contourpy-1.3.3-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4debd64f124ca62069f313a9cb86656ff087786016d76927ae2cf37846b006c9", size = 362859, upload-time = "2025-07-26T12:01:46.519Z" }, + { url = "https://files.pythonhosted.org/packages/33/71/e2a7945b7de4e58af42d708a219f3b2f4cff7386e6b6ab0a0fa0033c49a9/contourpy-1.3.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a15459b0f4615b00bbd1e91f1b9e19b7e63aea7483d03d804186f278c0af2659", size = 1332062, upload-time = "2025-07-26T12:01:48.964Z" }, + { url = "https://files.pythonhosted.org/packages/12/fc/4e87ac754220ccc0e807284f88e943d6d43b43843614f0a8afa469801db0/contourpy-1.3.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ca0fdcd73925568ca027e0b17ab07aad764be4706d0a925b89227e447d9737b7", size = 1403932, upload-time = "2025-07-26T12:01:51.979Z" }, + { url = "https://files.pythonhosted.org/packages/a6/2e/adc197a37443f934594112222ac1aa7dc9a98faf9c3842884df9a9d8751d/contourpy-1.3.3-cp313-cp313-win32.whl", hash = "sha256:b20c7c9a3bf701366556e1b1984ed2d0cedf999903c51311417cf5f591d8c78d", size = 185024, upload-time = "2025-07-26T12:01:53.245Z" }, + { url = "https://files.pythonhosted.org/packages/18/0b/0098c214843213759692cc638fce7de5c289200a830e5035d1791d7a2338/contourpy-1.3.3-cp313-cp313-win_amd64.whl", hash = "sha256:1cadd8b8969f060ba45ed7c1b714fe69185812ab43bd6b86a9123fe8f99c3263", size = 226578, upload-time = "2025-07-26T12:01:54.422Z" }, + { url = "https://files.pythonhosted.org/packages/8a/9a/2f6024a0c5995243cd63afdeb3651c984f0d2bc727fd98066d40e141ad73/contourpy-1.3.3-cp313-cp313-win_arm64.whl", hash = "sha256:fd914713266421b7536de2bfa8181aa8c699432b6763a0ea64195ebe28bff6a9", size = 193524, upload-time = "2025-07-26T12:01:55.73Z" }, + { url = "https://files.pythonhosted.org/packages/c0/b3/f8a1a86bd3298513f500e5b1f5fd92b69896449f6cab6a146a5d52715479/contourpy-1.3.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:88df9880d507169449d434c293467418b9f6cbe82edd19284aa0409e7fdb933d", size = 306730, upload-time = "2025-07-26T12:01:57.051Z" }, + { url = "https://files.pythonhosted.org/packages/3f/11/4780db94ae62fc0c2053909b65dc3246bd7cecfc4f8a20d957ad43aa4ad8/contourpy-1.3.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:d06bb1f751ba5d417047db62bca3c8fde202b8c11fb50742ab3ab962c81e8216", size = 287897, upload-time = "2025-07-26T12:01:58.663Z" }, + { url = "https://files.pythonhosted.org/packages/ae/15/e59f5f3ffdd6f3d4daa3e47114c53daabcb18574a26c21f03dc9e4e42ff0/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e4e6b05a45525357e382909a4c1600444e2a45b4795163d3b22669285591c1ae", size = 326751, upload-time = "2025-07-26T12:02:00.343Z" }, + { url = "https://files.pythonhosted.org/packages/0f/81/03b45cfad088e4770b1dcf72ea78d3802d04200009fb364d18a493857210/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ab3074b48c4e2cf1a960e6bbeb7f04566bf36b1861d5c9d4d8ac04b82e38ba20", size = 375486, upload-time = "2025-07-26T12:02:02.128Z" }, + { url = "https://files.pythonhosted.org/packages/0c/ba/49923366492ffbdd4486e970d421b289a670ae8cf539c1ea9a09822b371a/contourpy-1.3.3-cp313-cp313t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c3d53c796f8647d6deb1abe867daeb66dcc8a97e8455efa729516b997b8ed99", size = 388106, upload-time = "2025-07-26T12:02:03.615Z" }, + { url = "https://files.pythonhosted.org/packages/9f/52/5b00ea89525f8f143651f9f03a0df371d3cbd2fccd21ca9b768c7a6500c2/contourpy-1.3.3-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:50ed930df7289ff2a8d7afeb9603f8289e5704755c7e5c3bbd929c90c817164b", size = 352548, upload-time = "2025-07-26T12:02:05.165Z" }, + { url = "https://files.pythonhosted.org/packages/32/1d/a209ec1a3a3452d490f6b14dd92e72280c99ae3d1e73da74f8277d4ee08f/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:4feffb6537d64b84877da813a5c30f1422ea5739566abf0bd18065ac040e120a", size = 1322297, upload-time = "2025-07-26T12:02:07.379Z" }, + { url = "https://files.pythonhosted.org/packages/bc/9e/46f0e8ebdd884ca0e8877e46a3f4e633f6c9c8c4f3f6e72be3fe075994aa/contourpy-1.3.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:2b7e9480ffe2b0cd2e787e4df64270e3a0440d9db8dc823312e2c940c167df7e", size = 1391023, upload-time = "2025-07-26T12:02:10.171Z" }, + { url = "https://files.pythonhosted.org/packages/b9/70/f308384a3ae9cd2209e0849f33c913f658d3326900d0ff5d378d6a1422d2/contourpy-1.3.3-cp313-cp313t-win32.whl", hash = "sha256:283edd842a01e3dcd435b1c5116798d661378d83d36d337b8dde1d16a5fc9ba3", size = 196157, upload-time = "2025-07-26T12:02:11.488Z" }, + { url = "https://files.pythonhosted.org/packages/b2/dd/880f890a6663b84d9e34a6f88cded89d78f0091e0045a284427cb6b18521/contourpy-1.3.3-cp313-cp313t-win_amd64.whl", hash = "sha256:87acf5963fc2b34825e5b6b048f40e3635dd547f590b04d2ab317c2619ef7ae8", size = 240570, upload-time = "2025-07-26T12:02:12.754Z" }, + { url = "https://files.pythonhosted.org/packages/80/99/2adc7d8ffead633234817ef8e9a87115c8a11927a94478f6bb3d3f4d4f7d/contourpy-1.3.3-cp313-cp313t-win_arm64.whl", hash = "sha256:3c30273eb2a55024ff31ba7d052dde990d7d8e5450f4bbb6e913558b3d6c2301", size = 199713, upload-time = "2025-07-26T12:02:14.4Z" }, + { url = "https://files.pythonhosted.org/packages/72/8b/4546f3ab60f78c514ffb7d01a0bd743f90de36f0019d1be84d0a708a580a/contourpy-1.3.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:fde6c716d51c04b1c25d0b90364d0be954624a0ee9d60e23e850e8d48353d07a", size = 292189, upload-time = "2025-07-26T12:02:16.095Z" }, + { url = "https://files.pythonhosted.org/packages/fd/e1/3542a9cb596cadd76fcef413f19c79216e002623158befe6daa03dbfa88c/contourpy-1.3.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:cbedb772ed74ff5be440fa8eee9bd49f64f6e3fc09436d9c7d8f1c287b121d77", size = 273251, upload-time = "2025-07-26T12:02:17.524Z" }, + { url = "https://files.pythonhosted.org/packages/b1/71/f93e1e9471d189f79d0ce2497007731c1e6bf9ef6d1d61b911430c3db4e5/contourpy-1.3.3-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:22e9b1bd7a9b1d652cd77388465dc358dafcd2e217d35552424aa4f996f524f5", size = 335810, upload-time = "2025-07-26T12:02:18.9Z" }, + { url = "https://files.pythonhosted.org/packages/91/f9/e35f4c1c93f9275d4e38681a80506b5510e9327350c51f8d4a5a724d178c/contourpy-1.3.3-cp314-cp314-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a22738912262aa3e254e4f3cb079a95a67132fc5a063890e224393596902f5a4", size = 382871, upload-time = "2025-07-26T12:02:20.418Z" }, + { url = "https://files.pythonhosted.org/packages/b5/71/47b512f936f66a0a900d81c396a7e60d73419868fba959c61efed7a8ab46/contourpy-1.3.3-cp314-cp314-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:afe5a512f31ee6bd7d0dda52ec9864c984ca3d66664444f2d72e0dc4eb832e36", size = 386264, upload-time = "2025-07-26T12:02:21.916Z" }, + { url = "https://files.pythonhosted.org/packages/04/5f/9ff93450ba96b09c7c2b3f81c94de31c89f92292f1380261bd7195bea4ea/contourpy-1.3.3-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f64836de09927cba6f79dcd00fdd7d5329f3fccc633468507079c829ca4db4e3", size = 363819, upload-time = "2025-07-26T12:02:23.759Z" }, + { url = "https://files.pythonhosted.org/packages/3e/a6/0b185d4cc480ee494945cde102cb0149ae830b5fa17bf855b95f2e70ad13/contourpy-1.3.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:1fd43c3be4c8e5fd6e4f2baeae35ae18176cf2e5cced681cca908addf1cdd53b", size = 1333650, upload-time = "2025-07-26T12:02:26.181Z" }, + { url = "https://files.pythonhosted.org/packages/43/d7/afdc95580ca56f30fbcd3060250f66cedbde69b4547028863abd8aa3b47e/contourpy-1.3.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:6afc576f7b33cf00996e5c1102dc2a8f7cc89e39c0b55df93a0b78c1bd992b36", size = 1404833, upload-time = "2025-07-26T12:02:28.782Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e2/366af18a6d386f41132a48f033cbd2102e9b0cf6345d35ff0826cd984566/contourpy-1.3.3-cp314-cp314-win32.whl", hash = "sha256:66c8a43a4f7b8df8b71ee1840e4211a3c8d93b214b213f590e18a1beca458f7d", size = 189692, upload-time = "2025-07-26T12:02:30.128Z" }, + { url = "https://files.pythonhosted.org/packages/7d/c2/57f54b03d0f22d4044b8afb9ca0e184f8b1afd57b4f735c2fa70883dc601/contourpy-1.3.3-cp314-cp314-win_amd64.whl", hash = "sha256:cf9022ef053f2694e31d630feaacb21ea24224be1c3ad0520b13d844274614fd", size = 232424, upload-time = "2025-07-26T12:02:31.395Z" }, + { url = "https://files.pythonhosted.org/packages/18/79/a9416650df9b525737ab521aa181ccc42d56016d2123ddcb7b58e926a42c/contourpy-1.3.3-cp314-cp314-win_arm64.whl", hash = "sha256:95b181891b4c71de4bb404c6621e7e2390745f887f2a026b2d99e92c17892339", size = 198300, upload-time = "2025-07-26T12:02:32.956Z" }, + { url = "https://files.pythonhosted.org/packages/1f/42/38c159a7d0f2b7b9c04c64ab317042bb6952b713ba875c1681529a2932fe/contourpy-1.3.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:33c82d0138c0a062380332c861387650c82e4cf1747aaa6938b9b6516762e772", size = 306769, upload-time = "2025-07-26T12:02:34.2Z" }, + { url = "https://files.pythonhosted.org/packages/c3/6c/26a8205f24bca10974e77460de68d3d7c63e282e23782f1239f226fcae6f/contourpy-1.3.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:ea37e7b45949df430fe649e5de8351c423430046a2af20b1c1961cae3afcda77", size = 287892, upload-time = "2025-07-26T12:02:35.807Z" }, + { url = "https://files.pythonhosted.org/packages/66/06/8a475c8ab718ebfd7925661747dbb3c3ee9c82ac834ccb3570be49d129f4/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d304906ecc71672e9c89e87c4675dc5c2645e1f4269a5063b99b0bb29f232d13", size = 326748, upload-time = "2025-07-26T12:02:37.193Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a3/c5ca9f010a44c223f098fccd8b158bb1cb287378a31ac141f04730dc49be/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ca658cd1a680a5c9ea96dc61cdbae1e85c8f25849843aa799dfd3cb370ad4fbe", size = 375554, upload-time = "2025-07-26T12:02:38.894Z" }, + { url = "https://files.pythonhosted.org/packages/80/5b/68bd33ae63fac658a4145088c1e894405e07584a316738710b636c6d0333/contourpy-1.3.3-cp314-cp314t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:ab2fd90904c503739a75b7c8c5c01160130ba67944a7b77bbf36ef8054576e7f", size = 388118, upload-time = "2025-07-26T12:02:40.642Z" }, + { url = "https://files.pythonhosted.org/packages/40/52/4c285a6435940ae25d7410a6c36bda5145839bc3f0beb20c707cda18b9d2/contourpy-1.3.3-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b7301b89040075c30e5768810bc96a8e8d78085b47d8be6e4c3f5a0b4ed478a0", size = 352555, upload-time = "2025-07-26T12:02:42.25Z" }, + { url = "https://files.pythonhosted.org/packages/24/ee/3e81e1dd174f5c7fefe50e85d0892de05ca4e26ef1c9a59c2a57e43b865a/contourpy-1.3.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:2a2a8b627d5cc6b7c41a4beff6c5ad5eb848c88255fda4a8745f7e901b32d8e4", size = 1322295, upload-time = "2025-07-26T12:02:44.668Z" }, + { url = "https://files.pythonhosted.org/packages/3c/b2/6d913d4d04e14379de429057cd169e5e00f6c2af3bb13e1710bcbdb5da12/contourpy-1.3.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:fd6ec6be509c787f1caf6b247f0b1ca598bef13f4ddeaa126b7658215529ba0f", size = 1391027, upload-time = "2025-07-26T12:02:47.09Z" }, + { url = "https://files.pythonhosted.org/packages/93/8a/68a4ec5c55a2971213d29a9374913f7e9f18581945a7a31d1a39b5d2dfe5/contourpy-1.3.3-cp314-cp314t-win32.whl", hash = "sha256:e74a9a0f5e3fff48fb5a7f2fd2b9b70a3fe014a67522f79b7cca4c0c7e43c9ae", size = 202428, upload-time = "2025-07-26T12:02:48.691Z" }, + { url = "https://files.pythonhosted.org/packages/fa/96/fd9f641ffedc4fa3ace923af73b9d07e869496c9cc7a459103e6e978992f/contourpy-1.3.3-cp314-cp314t-win_amd64.whl", hash = "sha256:13b68d6a62db8eafaebb8039218921399baf6e47bf85006fd8529f2a08ef33fc", size = 250331, upload-time = "2025-07-26T12:02:50.137Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8c/469afb6465b853afff216f9528ffda78a915ff880ed58813ba4faf4ba0b6/contourpy-1.3.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b", size = 203831, upload-time = "2025-07-26T12:02:51.449Z" }, + { url = "https://files.pythonhosted.org/packages/a5/29/8dcfe16f0107943fa92388c23f6e05cff0ba58058c4c95b00280d4c75a14/contourpy-1.3.3-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:cd5dfcaeb10f7b7f9dc8941717c6c2ade08f587be2226222c12b25f0483ed497", size = 278809, upload-time = "2025-07-26T12:02:52.74Z" }, + { url = "https://files.pythonhosted.org/packages/85/a9/8b37ef4f7dafeb335daee3c8254645ef5725be4d9c6aa70b50ec46ef2f7e/contourpy-1.3.3-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:0c1fc238306b35f246d61a1d416a627348b5cf0648648a031e14bb8705fcdfe8", size = 261593, upload-time = "2025-07-26T12:02:54.037Z" }, + { url = "https://files.pythonhosted.org/packages/0a/59/ebfb8c677c75605cc27f7122c90313fd2f375ff3c8d19a1694bda74aaa63/contourpy-1.3.3-pp311-pypy311_pp73-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:70f9aad7de812d6541d29d2bbf8feb22ff7e1c299523db288004e3157ff4674e", size = 302202, upload-time = "2025-07-26T12:02:55.947Z" }, + { url = "https://files.pythonhosted.org/packages/3c/37/21972a15834d90bfbfb009b9d004779bd5a07a0ec0234e5ba8f64d5736f4/contourpy-1.3.3-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5ed3657edf08512fc3fe81b510e35c2012fbd3081d2e26160f27ca28affec989", size = 329207, upload-time = "2025-07-26T12:02:57.468Z" }, + { url = "https://files.pythonhosted.org/packages/0c/58/bd257695f39d05594ca4ad60df5bcb7e32247f9951fd09a9b8edb82d1daa/contourpy-1.3.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:3d1a3799d62d45c18bafd41c5fa05120b96a28079f2393af559b843d1a966a77", size = 225315, upload-time = "2025-07-26T12:02:58.801Z" }, +] + [[package]] name = "cryptography" version = "46.0.3" @@ -616,6 +808,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/c3/e90f4a4feae6410f914f8ebac129b9ae7a8c92eb60a638012dde42030a9d/cryptography-46.0.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:6b5063083824e5509fdba180721d55909ffacccc8adbec85268b48439423d78c", size = 3438528, upload-time = "2025-10-15T23:18:26.227Z" }, ] +[[package]] +name = "cycler" +version = "0.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, +] + [[package]] name = "daydream-scope" version = "0.1.0a7" @@ -623,6 +824,7 @@ source = { editable = "." } dependencies = [ { name = "accelerate" }, { name = "aiortc" }, + { name = "decord" }, { name = "diffusers" }, { name = "easydict" }, { name = "einops" }, @@ -634,12 +836,14 @@ dependencies = [ { name = "huggingface-hub" }, { name = "kernels" }, { name = "lmdb" }, + { name = "mediapipe" }, { name = "omegaconf" }, { name = "peft" }, { name = "pyopengl", marker = "sys_platform == 'win32'" }, { name = "safetensors" }, { name = "sageattention", version = "2.2.0", source = { url = "https://github.com/daydreamlive/SageAttention/releases/download/v2.2.0-linux/sageattention-2.2.0-cp310-cp310-linux_x86_64.whl" }, marker = "sys_platform == 'linux'" }, { name = "sageattention", version = "2.2.0+cu128torch2.8.0.post3", source = { url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl" }, marker = "sys_platform == 'win32'" }, + { name = "scikit-image" }, { name = "spoutgl", marker = "sys_platform == 'win32'" }, { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -652,6 +856,7 @@ dependencies = [ { name = "triton-windows", marker = "sys_platform == 'win32'" }, { name = "twilio" }, { name = "uvicorn" }, + { name = "xformers" }, ] [package.dev-dependencies] @@ -669,6 +874,7 @@ dev = [ requires-dist = [ { name = "accelerate", specifier = ">=1.1.1" }, { name = "aiortc", specifier = ">=1.13.0" }, + { name = "decord", specifier = ">=0.6.0" }, { name = "diffusers", specifier = ">=0.31.0" }, { name = "easydict", specifier = ">=1.13" }, { name = "einops", specifier = ">=0.8.1" }, @@ -680,12 +886,14 @@ requires-dist = [ { name = "huggingface-hub", specifier = ">=0.25.0" }, { name = "kernels", specifier = ">=0.10.4" }, { name = "lmdb", specifier = ">=1.7.3" }, + { name = "mediapipe", specifier = ">=0.10.11" }, { name = "omegaconf", specifier = ">=2.3.0" }, { name = "peft", specifier = ">=0.17.1" }, { name = "pyopengl", marker = "sys_platform == 'win32'", specifier = ">=3.1.10" }, { name = "safetensors", specifier = ">=0.6.2" }, { name = "sageattention", marker = "sys_platform == 'linux'", url = "https://github.com/daydreamlive/SageAttention/releases/download/v2.2.0-linux/sageattention-2.2.0-cp310-cp310-linux_x86_64.whl" }, { name = "sageattention", marker = "sys_platform == 'win32'", url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl" }, + { name = "scikit-image", specifier = ">=0.22.0" }, { name = "spoutgl", marker = "sys_platform == 'win32'", specifier = ">=0.1.1" }, { name = "torch", marker = "sys_platform != 'linux' and sys_platform != 'win32'", specifier = "==2.8.0" }, { name = "torch", marker = "sys_platform == 'linux' or sys_platform == 'win32'", specifier = "==2.8.0", index = "https://download.pytorch.org/whl/cu128" }, @@ -698,6 +906,7 @@ requires-dist = [ { name = "triton-windows", marker = "sys_platform == 'win32'", specifier = "==3.4.0.post21" }, { name = "twilio", specifier = ">=9.8.0" }, { name = "uvicorn", specifier = ">=0.35.0" }, + { name = "xformers", specifier = ">=0.0.32.post2" }, ] [package.metadata.requires-dev] @@ -711,6 +920,18 @@ dev = [ { name = "twine", specifier = ">=5.0.0" }, ] +[[package]] +name = "decord" +version = "0.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" }, + { url = "https://files.pythonhosted.org/packages/6c/be/e15b5b866da452e62635a7b27513f31cb581fa2ea9cc9b768b535d62a955/decord-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02665d7c4f1193a330205a791bc128f7e108eb6ae5b67144437a02f700943bad", size = 24733380, upload-time = "2021-06-14T21:30:57.766Z" }, +] + [[package]] name = "diffusers" version = "0.35.2" @@ -719,8 +940,7 @@ dependencies = [ { name = "filelock" }, { name = "huggingface-hub" }, { name = "importlib-metadata" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy" }, { name = "pillow" }, { name = "regex" }, { name = "requests" }, @@ -817,9 +1037,14 @@ name = "flash-attn" version = "2.8.3" source = { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.8cxx11abiFALSE-cp310-cp310-linux_x86_64.whl" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", ] dependencies = [ { name = "einops", marker = "sys_platform == 'linux'" }, @@ -840,7 +1065,8 @@ name = "flash-attn" version = "2.8.3" source = { url = "https://github.com/kingbri1/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu128torch2.8.0cxx11abiFALSE-cp310-cp310-win_amd64.whl" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version < '3.11' and sys_platform == 'win32'", ] @@ -858,6 +1084,72 @@ requires-dist = [ { name = "torch" }, ] +[[package]] +name = "flatbuffers" +version = "25.9.23" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9d/1f/3ee70b0a55137442038f2a33469cc5fddd7e0ad2abf83d7497c18a2b6923/flatbuffers-25.9.23.tar.gz", hash = "sha256:676f9fa62750bb50cf531b42a0a2a118ad8f7f797a511eda12881c016f093b12", size = 22067, upload-time = "2025-09-24T05:25:30.106Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ee/1b/00a78aa2e8fbd63f9af08c9c19e6deb3d5d66b4dda677a0f61654680ee89/flatbuffers-25.9.23-py2.py3-none-any.whl", hash = "sha256:255538574d6cb6d0a79a17ec8bc0d30985913b87513a01cce8bcdb6b4c44d0e2", size = 30869, upload-time = "2025-09-24T05:25:28.912Z" }, +] + +[[package]] +name = "fonttools" +version = "4.61.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ec/ca/cf17b88a8df95691275a3d77dc0a5ad9907f328ae53acbe6795da1b2f5ed/fonttools-4.61.1.tar.gz", hash = "sha256:6675329885c44657f826ef01d9e4fb33b9158e9d93c537d84ad8399539bc6f69", size = 3565756, upload-time = "2025-12-12T17:31:24.246Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/94/8a28707adb00bed1bf22dac16ccafe60faf2ade353dcb32c3617ee917307/fonttools-4.61.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:7c7db70d57e5e1089a274cbb2b1fd635c9a24de809a231b154965d415d6c6d24", size = 2854799, upload-time = "2025-12-12T17:29:27.5Z" }, + { url = "https://files.pythonhosted.org/packages/94/93/c2e682faaa5ee92034818d8f8a8145ae73eb83619600495dcf8503fa7771/fonttools-4.61.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:5fe9fd43882620017add5eabb781ebfbc6998ee49b35bd7f8f79af1f9f99a958", size = 2403032, upload-time = "2025-12-12T17:29:30.115Z" }, + { url = "https://files.pythonhosted.org/packages/f1/62/1748f7e7e1ee41aa52279fd2e3a6d0733dc42a673b16932bad8e5d0c8b28/fonttools-4.61.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d8db08051fc9e7d8bc622f2112511b8107d8f27cd89e2f64ec45e9825e8288da", size = 4897863, upload-time = "2025-12-12T17:29:32.535Z" }, + { url = "https://files.pythonhosted.org/packages/69/69/4ca02ee367d2c98edcaeb83fc278d20972502ee071214ad9d8ca85e06080/fonttools-4.61.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a76d4cb80f41ba94a6691264be76435e5f72f2cb3cab0b092a6212855f71c2f6", size = 4859076, upload-time = "2025-12-12T17:29:34.907Z" }, + { url = "https://files.pythonhosted.org/packages/8c/f5/660f9e3cefa078861a7f099107c6d203b568a6227eef163dd173bfc56bdc/fonttools-4.61.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:a13fc8aeb24bad755eea8f7f9d409438eb94e82cf86b08fe77a03fbc8f6a96b1", size = 4875623, upload-time = "2025-12-12T17:29:37.33Z" }, + { url = "https://files.pythonhosted.org/packages/63/d1/9d7c5091d2276ed47795c131c1bf9316c3c1ab2789c22e2f59e0572ccd38/fonttools-4.61.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:b846a1fcf8beadeb9ea4f44ec5bdde393e2f1569e17d700bfc49cd69bde75881", size = 4993327, upload-time = "2025-12-12T17:29:39.781Z" }, + { url = "https://files.pythonhosted.org/packages/6f/2d/28def73837885ae32260d07660a052b99f0aa00454867d33745dfe49dbf0/fonttools-4.61.1-cp310-cp310-win32.whl", hash = "sha256:78a7d3ab09dc47ac1a363a493e6112d8cabed7ba7caad5f54dbe2f08676d1b47", size = 1502180, upload-time = "2025-12-12T17:29:42.217Z" }, + { url = "https://files.pythonhosted.org/packages/63/fa/bfdc98abb4dd2bd491033e85e3ba69a2313c850e759a6daa014bc9433b0f/fonttools-4.61.1-cp310-cp310-win_amd64.whl", hash = "sha256:eff1ac3cc66c2ac7cda1e64b4e2f3ffef474b7335f92fc3833fc632d595fcee6", size = 1550654, upload-time = "2025-12-12T17:29:44.564Z" }, + { url = "https://files.pythonhosted.org/packages/69/12/bf9f4eaa2fad039356cc627587e30ed008c03f1cebd3034376b5ee8d1d44/fonttools-4.61.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:c6604b735bb12fef8e0efd5578c9fb5d3d8532d5001ea13a19cddf295673ee09", size = 2852213, upload-time = "2025-12-12T17:29:46.675Z" }, + { url = "https://files.pythonhosted.org/packages/ac/49/4138d1acb6261499bedde1c07f8c2605d1d8f9d77a151e5507fd3ef084b6/fonttools-4.61.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5ce02f38a754f207f2f06557523cd39a06438ba3aafc0639c477ac409fc64e37", size = 2401689, upload-time = "2025-12-12T17:29:48.769Z" }, + { url = "https://files.pythonhosted.org/packages/e5/fe/e6ce0fe20a40e03aef906af60aa87668696f9e4802fa283627d0b5ed777f/fonttools-4.61.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77efb033d8d7ff233385f30c62c7c79271c8885d5c9657d967ede124671bbdfb", size = 5058809, upload-time = "2025-12-12T17:29:51.701Z" }, + { url = "https://files.pythonhosted.org/packages/79/61/1ca198af22f7dd22c17ab86e9024ed3c06299cfdb08170640e9996d501a0/fonttools-4.61.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:75c1a6dfac6abd407634420c93864a1e274ebc1c7531346d9254c0d8f6ca00f9", size = 5036039, upload-time = "2025-12-12T17:29:53.659Z" }, + { url = "https://files.pythonhosted.org/packages/99/cc/fa1801e408586b5fce4da9f5455af8d770f4fc57391cd5da7256bb364d38/fonttools-4.61.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:0de30bfe7745c0d1ffa2b0b7048fb7123ad0d71107e10ee090fa0b16b9452e87", size = 5034714, upload-time = "2025-12-12T17:29:55.592Z" }, + { url = "https://files.pythonhosted.org/packages/bf/aa/b7aeafe65adb1b0a925f8f25725e09f078c635bc22754f3fecb7456955b0/fonttools-4.61.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:58b0ee0ab5b1fc9921eccfe11d1435added19d6494dde14e323f25ad2bc30c56", size = 5158648, upload-time = "2025-12-12T17:29:57.861Z" }, + { url = "https://files.pythonhosted.org/packages/99/f9/08ea7a38663328881384c6e7777bbefc46fd7d282adfd87a7d2b84ec9d50/fonttools-4.61.1-cp311-cp311-win32.whl", hash = "sha256:f79b168428351d11e10c5aeb61a74e1851ec221081299f4cf56036a95431c43a", size = 2280681, upload-time = "2025-12-12T17:29:59.943Z" }, + { url = "https://files.pythonhosted.org/packages/07/ad/37dd1ae5fa6e01612a1fbb954f0927681f282925a86e86198ccd7b15d515/fonttools-4.61.1-cp311-cp311-win_amd64.whl", hash = "sha256:fe2efccb324948a11dd09d22136fe2ac8a97d6c1347cf0b58a911dcd529f66b7", size = 2331951, upload-time = "2025-12-12T17:30:02.254Z" }, + { url = "https://files.pythonhosted.org/packages/6f/16/7decaa24a1bd3a70c607b2e29f0adc6159f36a7e40eaba59846414765fd4/fonttools-4.61.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:f3cb4a569029b9f291f88aafc927dd53683757e640081ca8c412781ea144565e", size = 2851593, upload-time = "2025-12-12T17:30:04.225Z" }, + { url = "https://files.pythonhosted.org/packages/94/98/3c4cb97c64713a8cf499b3245c3bf9a2b8fd16a3e375feff2aed78f96259/fonttools-4.61.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41a7170d042e8c0024703ed13b71893519a1a6d6e18e933e3ec7507a2c26a4b2", size = 2400231, upload-time = "2025-12-12T17:30:06.47Z" }, + { url = "https://files.pythonhosted.org/packages/b7/37/82dbef0f6342eb01f54bca073ac1498433d6ce71e50c3c3282b655733b31/fonttools-4.61.1-cp312-cp312-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:10d88e55330e092940584774ee5e8a6971b01fc2f4d3466a1d6c158230880796", size = 4954103, upload-time = "2025-12-12T17:30:08.432Z" }, + { url = "https://files.pythonhosted.org/packages/6c/44/f3aeac0fa98e7ad527f479e161aca6c3a1e47bb6996b053d45226fe37bf2/fonttools-4.61.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:15acc09befd16a0fb8a8f62bc147e1a82817542d72184acca9ce6e0aeda9fa6d", size = 5004295, upload-time = "2025-12-12T17:30:10.56Z" }, + { url = "https://files.pythonhosted.org/packages/14/e8/7424ced75473983b964d09f6747fa09f054a6d656f60e9ac9324cf40c743/fonttools-4.61.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e6bcdf33aec38d16508ce61fd81838f24c83c90a1d1b8c68982857038673d6b8", size = 4944109, upload-time = "2025-12-12T17:30:12.874Z" }, + { url = "https://files.pythonhosted.org/packages/c8/8b/6391b257fa3d0b553d73e778f953a2f0154292a7a7a085e2374b111e5410/fonttools-4.61.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5fade934607a523614726119164ff621e8c30e8fa1ffffbbd358662056ba69f0", size = 5093598, upload-time = "2025-12-12T17:30:15.79Z" }, + { url = "https://files.pythonhosted.org/packages/d9/71/fd2ea96cdc512d92da5678a1c98c267ddd4d8c5130b76d0f7a80f9a9fde8/fonttools-4.61.1-cp312-cp312-win32.whl", hash = "sha256:75da8f28eff26defba42c52986de97b22106cb8f26515b7c22443ebc9c2d3261", size = 2269060, upload-time = "2025-12-12T17:30:18.058Z" }, + { url = "https://files.pythonhosted.org/packages/80/3b/a3e81b71aed5a688e89dfe0e2694b26b78c7d7f39a5ffd8a7d75f54a12a8/fonttools-4.61.1-cp312-cp312-win_amd64.whl", hash = "sha256:497c31ce314219888c0e2fce5ad9178ca83fe5230b01a5006726cdf3ac9f24d9", size = 2319078, upload-time = "2025-12-12T17:30:22.862Z" }, + { url = "https://files.pythonhosted.org/packages/4b/cf/00ba28b0990982530addb8dc3e9e6f2fa9cb5c20df2abdda7baa755e8fe1/fonttools-4.61.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c56c488ab471628ff3bfa80964372fc13504ece601e0d97a78ee74126b2045c", size = 2846454, upload-time = "2025-12-12T17:30:24.938Z" }, + { url = "https://files.pythonhosted.org/packages/5a/ca/468c9a8446a2103ae645d14fee3f610567b7042aba85031c1c65e3ef7471/fonttools-4.61.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:dc492779501fa723b04d0ab1f5be046797fee17d27700476edc7ee9ae535a61e", size = 2398191, upload-time = "2025-12-12T17:30:27.343Z" }, + { url = "https://files.pythonhosted.org/packages/a3/4b/d67eedaed19def5967fade3297fed8161b25ba94699efc124b14fb68cdbc/fonttools-4.61.1-cp313-cp313-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:64102ca87e84261419c3747a0d20f396eb024bdbeb04c2bfb37e2891f5fadcb5", size = 4928410, upload-time = "2025-12-12T17:30:29.771Z" }, + { url = "https://files.pythonhosted.org/packages/b0/8d/6fb3494dfe61a46258cd93d979cf4725ded4eb46c2a4ca35e4490d84daea/fonttools-4.61.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4c1b526c8d3f615a7b1867f38a9410849c8f4aef078535742198e942fba0e9bd", size = 4984460, upload-time = "2025-12-12T17:30:32.073Z" }, + { url = "https://files.pythonhosted.org/packages/f7/f1/a47f1d30b3dc00d75e7af762652d4cbc3dff5c2697a0dbd5203c81afd9c3/fonttools-4.61.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:41ed4b5ec103bd306bb68f81dc166e77409e5209443e5773cb4ed837bcc9b0d3", size = 4925800, upload-time = "2025-12-12T17:30:34.339Z" }, + { url = "https://files.pythonhosted.org/packages/a7/01/e6ae64a0981076e8a66906fab01539799546181e32a37a0257b77e4aa88b/fonttools-4.61.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b501c862d4901792adaec7c25b1ecc749e2662543f68bb194c42ba18d6eec98d", size = 5067859, upload-time = "2025-12-12T17:30:36.593Z" }, + { url = "https://files.pythonhosted.org/packages/73/aa/28e40b8d6809a9b5075350a86779163f074d2b617c15d22343fce81918db/fonttools-4.61.1-cp313-cp313-win32.whl", hash = "sha256:4d7092bb38c53bbc78e9255a59158b150bcdc115a1e3b3ce0b5f267dc35dd63c", size = 2267821, upload-time = "2025-12-12T17:30:38.478Z" }, + { url = "https://files.pythonhosted.org/packages/1a/59/453c06d1d83dc0951b69ef692d6b9f1846680342927df54e9a1ca91c6f90/fonttools-4.61.1-cp313-cp313-win_amd64.whl", hash = "sha256:21e7c8d76f62ab13c9472ccf74515ca5b9a761d1bde3265152a6dc58700d895b", size = 2318169, upload-time = "2025-12-12T17:30:40.951Z" }, + { url = "https://files.pythonhosted.org/packages/32/8f/4e7bf82c0cbb738d3c2206c920ca34ca74ef9dabde779030145d28665104/fonttools-4.61.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:fff4f534200a04b4a36e7ae3cb74493afe807b517a09e99cb4faa89a34ed6ecd", size = 2846094, upload-time = "2025-12-12T17:30:43.511Z" }, + { url = "https://files.pythonhosted.org/packages/71/09/d44e45d0a4f3a651f23a1e9d42de43bc643cce2971b19e784cc67d823676/fonttools-4.61.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d9203500f7c63545b4ce3799319fe4d9feb1a1b89b28d3cb5abd11b9dd64147e", size = 2396589, upload-time = "2025-12-12T17:30:45.681Z" }, + { url = "https://files.pythonhosted.org/packages/89/18/58c64cafcf8eb677a99ef593121f719e6dcbdb7d1c594ae5a10d4997ca8a/fonttools-4.61.1-cp314-cp314-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:fa646ecec9528bef693415c79a86e733c70a4965dd938e9a226b0fc64c9d2e6c", size = 4877892, upload-time = "2025-12-12T17:30:47.709Z" }, + { url = "https://files.pythonhosted.org/packages/8a/ec/9e6b38c7ba1e09eb51db849d5450f4c05b7e78481f662c3b79dbde6f3d04/fonttools-4.61.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:11f35ad7805edba3aac1a3710d104592df59f4b957e30108ae0ba6c10b11dd75", size = 4972884, upload-time = "2025-12-12T17:30:49.656Z" }, + { url = "https://files.pythonhosted.org/packages/5e/87/b5339da8e0256734ba0dbbf5b6cdebb1dd79b01dc8c270989b7bcd465541/fonttools-4.61.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b931ae8f62db78861b0ff1ac017851764602288575d65b8e8ff1963fed419063", size = 4924405, upload-time = "2025-12-12T17:30:51.735Z" }, + { url = "https://files.pythonhosted.org/packages/0b/47/e3409f1e1e69c073a3a6fd8cb886eb18c0bae0ee13db2c8d5e7f8495e8b7/fonttools-4.61.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b148b56f5de675ee16d45e769e69f87623a4944f7443850bf9a9376e628a89d2", size = 5035553, upload-time = "2025-12-12T17:30:54.823Z" }, + { url = "https://files.pythonhosted.org/packages/bf/b6/1f6600161b1073a984294c6c031e1a56ebf95b6164249eecf30012bb2e38/fonttools-4.61.1-cp314-cp314-win32.whl", hash = "sha256:9b666a475a65f4e839d3d10473fad6d47e0a9db14a2f4a224029c5bfde58ad2c", size = 2271915, upload-time = "2025-12-12T17:30:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/52/7b/91e7b01e37cc8eb0e1f770d08305b3655e4f002fc160fb82b3390eabacf5/fonttools-4.61.1-cp314-cp314-win_amd64.whl", hash = "sha256:4f5686e1fe5fce75d82d93c47a438a25bf0d1319d2843a926f741140b2b16e0c", size = 2323487, upload-time = "2025-12-12T17:30:59.804Z" }, + { url = "https://files.pythonhosted.org/packages/39/5c/908ad78e46c61c3e3ed70c3b58ff82ab48437faf84ec84f109592cabbd9f/fonttools-4.61.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:e76ce097e3c57c4bcb67c5aa24a0ecdbd9f74ea9219997a707a4061fbe2707aa", size = 2929571, upload-time = "2025-12-12T17:31:02.574Z" }, + { url = "https://files.pythonhosted.org/packages/bd/41/975804132c6dea64cdbfbaa59f3518a21c137a10cccf962805b301ac6ab2/fonttools-4.61.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:9cfef3ab326780c04d6646f68d4b4742aae222e8b8ea1d627c74e38afcbc9d91", size = 2435317, upload-time = "2025-12-12T17:31:04.974Z" }, + { url = "https://files.pythonhosted.org/packages/b0/5a/aef2a0a8daf1ebaae4cfd83f84186d4a72ee08fd6a8451289fcd03ffa8a4/fonttools-4.61.1-cp314-cp314t-manylinux1_x86_64.manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:a75c301f96db737e1c5ed5fd7d77d9c34466de16095a266509e13da09751bd19", size = 4882124, upload-time = "2025-12-12T17:31:07.456Z" }, + { url = "https://files.pythonhosted.org/packages/80/33/d6db3485b645b81cea538c9d1c9219d5805f0877fda18777add4671c5240/fonttools-4.61.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:91669ccac46bbc1d09e9273546181919064e8df73488ea087dcac3e2968df9ba", size = 5100391, upload-time = "2025-12-12T17:31:09.732Z" }, + { url = "https://files.pythonhosted.org/packages/6c/d6/675ba631454043c75fcf76f0ca5463eac8eb0666ea1d7badae5fea001155/fonttools-4.61.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c33ab3ca9d3ccd581d58e989d67554e42d8d4ded94ab3ade3508455fe70e65f7", size = 4978800, upload-time = "2025-12-12T17:31:11.681Z" }, + { url = "https://files.pythonhosted.org/packages/7f/33/d3ec753d547a8d2bdaedd390d4a814e8d5b45a093d558f025c6b990b554c/fonttools-4.61.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:664c5a68ec406f6b1547946683008576ef8b38275608e1cee6c061828171c118", size = 5006426, upload-time = "2025-12-12T17:31:13.764Z" }, + { url = "https://files.pythonhosted.org/packages/b4/40/cc11f378b561a67bea850ab50063366a0d1dd3f6d0a30ce0f874b0ad5664/fonttools-4.61.1-cp314-cp314t-win32.whl", hash = "sha256:aed04cabe26f30c1647ef0e8fbb207516fd40fe9472e9439695f5c6998e60ac5", size = 2335377, upload-time = "2025-12-12T17:31:16.49Z" }, + { url = "https://files.pythonhosted.org/packages/e4/ff/c9a2b66b39f8628531ea58b320d66d951267c98c6a38684daa8f50fb02f8/fonttools-4.61.1-cp314-cp314t-win_amd64.whl", hash = "sha256:2180f14c141d2f0f3da43f3a81bc8aa4684860f6b0e6f9e165a4831f24e6a23b", size = 2400613, upload-time = "2025-12-12T17:31:18.769Z" }, + { url = "https://files.pythonhosted.org/packages/c7/4e/ce75a57ff3aebf6fc1f4e9d508b8e5810618a33d900ad6c19eb30b290b97/fonttools-4.61.1-py3-none-any.whl", hash = "sha256:17d2bf5d541add43822bcf0c43d7d847b160c9bb01d15d5007d84e2217aaa371", size = 1148996, upload-time = "2025-12-12T17:31:21.03Z" }, +] + [[package]] name = "freezegun" version = "1.5.5" @@ -1176,8 +1468,7 @@ name = "imageio" version = "2.37.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy" }, { name = "pillow" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a3/6f/606be632e37bf8d05b253e8626c2291d74c691ddc7bcdf7d6aaf33b32f6a/imageio-2.37.2.tar.gz", hash = "sha256:0212ef2727ac9caa5ca4b2c75ae89454312f440a756fcfc8ef1993e718f50f8a", size = 389600, upload-time = "2025-11-04T14:29:39.898Z" } @@ -1256,6 +1547,153 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b4/09/726f168acad366b11e420df31bf1c702a54d373a83f968d94141a8c3fde0/jaraco_functools-4.3.0-py3-none-any.whl", hash = "sha256:227ff8ed6f7b8f62c56deff101545fa7543cf2c8e7b82a7c2116e672f29c26e8", size = 10408, upload-time = "2025-08-18T20:05:08.69Z" }, ] +[[package]] +name = "jax" +version = "0.5.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "jaxlib", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" }, + { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", marker = "python_full_version < '3.11' or python_full_version >= '3.13'" }, + { name = "opt-einsum", marker = "python_full_version < '3.11' or python_full_version >= '3.13'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/13/e5/dabb73ab10330e9535aba14fc668b04a46fcd8e78f06567c4f4f1adce340/jax-0.5.3.tar.gz", hash = "sha256:f17fcb0fd61dc289394af6ce4de2dada2312f2689bb0d73642c6f026a95fbb2c", size = 2072748, upload-time = "2025-03-19T18:23:40.901Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/bb/fdc6513a9aada13fd21e9860e2adee5f6eea2b4f0a145b219288875acb26/jax-0.5.3-py3-none-any.whl", hash = "sha256:1483dc237b4f47e41755d69429e8c3c138736716147cd43bb2b99b259d4e3c41", size = 2406371, upload-time = "2025-03-19T18:23:38.952Z" }, +] + +[[package]] +name = "jax" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "numpy", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "opt-einsum", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bc/e8/b393ee314d3b042bd66b986d38e52f4e6046590399d916381265c20467d3/jax-0.7.1.tar.gz", hash = "sha256:118f56338c503361d2791f069d24339d8d44a8db442ed851d2e591222fb7a56d", size = 2428411, upload-time = "2025-08-20T15:55:46.098Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/81/793d78c91b0546b3b1f08e55fdd97437174171cd7d70e46098f1a4d94b7b/jax-0.7.1-py3-none-any.whl", hash = "sha256:056e576e0e58465506125699f48111ac8891cce4c9ebf034704c42b219dfd4a6", size = 2827341, upload-time = "2025-08-20T15:55:44.576Z" }, +] + +[[package]] +name = "jaxlib" +version = "0.5.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", marker = "python_full_version < '3.11' or python_full_version >= '3.13'" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/12/b1da8468ad843b30976b0e87c6b344ee621fb75ef8bbd39156a303f59059/jaxlib-0.5.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:48ff5c89fb8a0fe04d475e9ddc074b4879a91d7ab68a51cec5cd1e87f81e6c47", size = 63694868, upload-time = "2025-03-19T18:23:52.193Z" }, + { url = "https://files.pythonhosted.org/packages/0e/a5/378d71e8bcffbb229a0952d713a2ed766c959a04777abc0ee01b5aac29b7/jaxlib-0.5.3-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:972400db4af6e85270d81db5e6e620d31395f0472e510c50dfcd4cb3f72b7220", size = 95766664, upload-time = "2025-03-19T18:23:58.12Z" }, + { url = "https://files.pythonhosted.org/packages/f1/86/1edf85f425532cbba0180d969f396590dd266909e4dfb0e95f8ee9a8e5fe/jaxlib-0.5.3-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:52be6c9775aff738a61170d8c047505c75bb799a45518e66a7a0908127b11785", size = 105118562, upload-time = "2025-03-19T18:24:03.384Z" }, + { url = "https://files.pythonhosted.org/packages/61/84/427cd89dd7904a4c923a3fc5494daec8d42d824c1a40d7a5d1c985e2f5ac/jaxlib-0.5.3-cp310-cp310-win_amd64.whl", hash = "sha256:b41a6fcaeb374fabc4ee7e74cfed60843bdab607cd54f60a68b7f7655cde2b66", size = 65766784, upload-time = "2025-03-19T18:24:09.025Z" }, + { url = "https://files.pythonhosted.org/packages/c2/f2/d9397f264141f2289e229b2faf3b3ddb6397b014a09abe234367814f9697/jaxlib-0.5.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:b62bd8b29e5a4f9bfaa57c8daf6e04820b2c994f448f3dec602d64255545e9f2", size = 63696815, upload-time = "2025-03-19T18:24:14.662Z" }, + { url = "https://files.pythonhosted.org/packages/e8/91/04bf391a21ccfb299b9952f91d5c082e5f9877221e5d98592875af4a50e4/jaxlib-0.5.3-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:a4666f81d72c060ed3e581ded116a9caa9b0a70a148a54cb12a1d3afca3624b5", size = 95770114, upload-time = "2025-03-19T18:24:19.498Z" }, + { url = "https://files.pythonhosted.org/packages/67/de/50debb40944baa5ba459604578f8c721be9f38c78ef9e8902895566e6a66/jaxlib-0.5.3-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:29e1530fc81833216f1e28b578d0c59697654f72ee31c7a44ed7753baf5ac466", size = 105119259, upload-time = "2025-03-19T18:24:25.39Z" }, + { url = "https://files.pythonhosted.org/packages/20/91/d73c842d1e5cc6b914bb521006d668fbfda4c53cd4424ce9c3a097f6c071/jaxlib-0.5.3-cp311-cp311-win_amd64.whl", hash = "sha256:8eb54e38d789557579f900ea3d70f104a440f8555a9681ed45f4a122dcbfd92e", size = 65765739, upload-time = "2025-03-19T18:24:30.264Z" }, + { url = "https://files.pythonhosted.org/packages/d5/a5/646af791ccf75641b4df84fb6cb6e3914b0df87ec5fa5f82397fd5dc30ee/jaxlib-0.5.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d394dbde4a1c6bd67501cfb29d3819a10b900cb534cc0fc603319f7092f24cfa", size = 63711839, upload-time = "2025-03-19T18:24:34.555Z" }, + { url = "https://files.pythonhosted.org/packages/53/8c/cbd861e40f0efe7923962ade21919fddcea43fae2794634833e800009b14/jaxlib-0.5.3-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:bddf6360377aa1c792e47fd87f307c342e331e5ff3582f940b1bca00f6b4bc73", size = 95764647, upload-time = "2025-03-19T18:24:39.376Z" }, + { url = "https://files.pythonhosted.org/packages/3e/03/bace4acec295febca9329b3d2dd927b8ac74841e620e0d675f76109b805b/jaxlib-0.5.3-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:5a5e88ab1cd6fdf78d69abe3544e8f09cce200dd339bb85fbe3c2ea67f2a5e68", size = 105132789, upload-time = "2025-03-19T18:24:45.232Z" }, + { url = "https://files.pythonhosted.org/packages/79/f8/34568ec75f53d55b68649b6e1d6befd976fb9646e607954477264f5379ce/jaxlib-0.5.3-cp312-cp312-win_amd64.whl", hash = "sha256:520665929649f29f7d948d4070dbaf3e032a4c1f7c11f2863eac73320fcee784", size = 65789714, upload-time = "2025-03-19T18:24:51.218Z" }, + { url = "https://files.pythonhosted.org/packages/b4/d0/ed6007cd17dc0f37f950f89e785092d9f0541f3fa6021d029657955206b5/jaxlib-0.5.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:31321c25282a06a6dfc940507bc14d0a0ac838d8ced6c07aa00a7fae34ce7b3f", size = 63710483, upload-time = "2025-03-19T18:24:55.41Z" }, + { url = "https://files.pythonhosted.org/packages/36/8f/cafdf24170084de897ffe2a030241c2ba72d12eede85b940a81a94cab156/jaxlib-0.5.3-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:e904b92dedfbc7e545725a8d7676987030ae9c069001d94701bc109c6dab4100", size = 95765995, upload-time = "2025-03-19T18:25:00.062Z" }, + { url = "https://files.pythonhosted.org/packages/86/c7/fc0755ebd999c7c66ac4203d99f958d5ffc0a34eb270f57932ca0213bb54/jaxlib-0.5.3-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:bb7593cb7fffcb13963f22fa5229ed960b8fb4ae5ec3b0820048cbd67f1e8e31", size = 105130796, upload-time = "2025-03-19T18:25:05.574Z" }, + { url = "https://files.pythonhosted.org/packages/83/98/e32da21a490dc408d172ba246d6c47428482fe50d771c3f813e5fc063781/jaxlib-0.5.3-cp313-cp313-win_amd64.whl", hash = "sha256:8019f73a10b1290f988dd3768c684f3a8a147239091c3b790ce7e47e3bbc00bd", size = 65792205, upload-time = "2025-03-19T18:25:10.684Z" }, + { url = "https://files.pythonhosted.org/packages/88/c6/0d69ed0d408c811959a471563afa99baecacdc56ed1799002e309520b565/jaxlib-0.5.3-cp313-cp313t-manylinux2014_x86_64.whl", hash = "sha256:4c9a9d4cda091a3ef068ace8379fff9e98eea2fc51dbdd7c3386144a1bdf715d", size = 105318736, upload-time = "2025-03-25T15:00:12.514Z" }, +] + +[[package]] +name = "jaxlib" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "numpy", marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/af/5058d545e95f99a54289648f5430cc3c23263dd70a1391e7491f24ed328d/jaxlib-0.7.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3f32c3e4c167b7327c342e82d3df84079714ea0b43718be871d039999670b3c9", size = 57686934, upload-time = "2025-08-20T15:55:58.989Z" }, + { url = "https://files.pythonhosted.org/packages/e8/77/ef7f6cd03e699da7d9755f88741c29b3015654473fc9d5f906da19edcb47/jaxlib-0.7.1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:9fb189c3b39470c4394ffcb18b71e47cffc5bf85e8fcb1e33692686b0c3e04dd", size = 85134885, upload-time = "2025-08-20T15:56:03.484Z" }, + { url = "https://files.pythonhosted.org/packages/4d/72/304018d46703f337787f010735f70d17212f86778fcba8bb5cf678f8e460/jaxlib-0.7.1-cp311-cp311-manylinux_2_27_x86_64.whl", hash = "sha256:eaf5f68f53bf4dcb93b6512538547667625588e4f3ccaeef048788fd18d8c0d5", size = 81147868, upload-time = "2025-08-20T15:56:07.214Z" }, + { url = "https://files.pythonhosted.org/packages/f7/b7/0f0df407518691099d659ba6e19db01320dfb58e49d80594eaddd57d77c1/jaxlib-0.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:ab4510fbaeafac6c794ab335f23e71200d824c48f6a0ab20553db8deab8805c5", size = 61185342, upload-time = "2025-08-20T15:56:10.452Z" }, + { url = "https://files.pythonhosted.org/packages/ef/1f/10543d7a3f7e76dd4bbdc77134890ac2f41bc8570c565961464f6320009b/jaxlib-0.7.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:127c07c727703e5d59f84f655169bec849f4422e52f8546349cecc30a8a13e1d", size = 57682851, upload-time = "2025-08-20T15:56:13.395Z" }, + { url = "https://files.pythonhosted.org/packages/de/4d/76ee71959311fe3da9951aa6f55af8f98eb3572bb322f5a7c89faf7ab933/jaxlib-0.7.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:f0f1f52956b8c2518ab000a4d3d8c21be777e1d47f926ba03640e391061a41ee", size = 85133707, upload-time = "2025-08-20T15:56:16.908Z" }, + { url = "https://files.pythonhosted.org/packages/0d/50/e37d02e250f5feb755112ec95b1c012a36d48a99209277267037d100f630/jaxlib-0.7.1-cp312-cp312-manylinux_2_27_x86_64.whl", hash = "sha256:74abd3135797f82440dd3711a35cba16c430d1bba65474b85bb70e41733a52e9", size = 81156916, upload-time = "2025-08-20T15:56:20.41Z" }, + { url = "https://files.pythonhosted.org/packages/5a/97/c6c28dfe57cccffd85512615416024b52dd327d78270204caba9311e71f1/jaxlib-0.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:c4023863b14f280516f24ecb7539b4300a3236ea81ed69ad82595beceed1ba1f", size = 61212445, upload-time = "2025-08-20T15:56:23.929Z" }, + { url = "https://files.pythonhosted.org/packages/35/ad/ab61b9bf72bc2903ee54d02e8d1e1486a4860878712c80c4a52025bfe003/jaxlib-0.7.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2a27379e5ed29765980ef32d4d77f57dd0e1dd965803ac674554044ac23cd3ec", size = 57681905, upload-time = "2025-08-20T15:56:26.82Z" }, + { url = "https://files.pythonhosted.org/packages/12/e8/fbc318afd21ea7b232ec6a3a1b0138da0d65db1d6f1cea8af458a7d6a482/jaxlib-0.7.1-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:ce6a1ba6019764870c27507aca18e526998ad3ad4bea2dd61e19d0499c3b8b04", size = 85133036, upload-time = "2025-08-20T15:56:30.085Z" }, + { url = "https://files.pythonhosted.org/packages/56/69/209e5b81dd89da84f729a684c126cc3c22bb8d6516f8c968444c7d86050f/jaxlib-0.7.1-cp313-cp313-manylinux_2_27_x86_64.whl", hash = "sha256:b350f519a86eff5a4b1ee014c7faa36585f47f3d63787d1f3e9bdffe9cc41a66", size = 81155944, upload-time = "2025-08-20T15:56:33.917Z" }, + { url = "https://files.pythonhosted.org/packages/95/46/685ecf0fa3c6d844ac33c0dea59601b9cc6b7236d0e4eb52dc3b7f6f52bb/jaxlib-0.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:407347dd5846248ee40db33da26f2b18c7d135249ff06c0271a2e33efd8fb3fe", size = 61213006, upload-time = "2025-08-20T15:56:37.796Z" }, + { url = "https://files.pythonhosted.org/packages/4f/ca/c70e5439eec1abc1a93bad521452dbddeefa68128f256991e6845e9fc56a/jaxlib-0.7.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:66b33cc15040af5b8d16d3006811eb31372e9f4cfe09393d6cea91795366bfa4", size = 57787598, upload-time = "2025-08-20T15:56:41.372Z" }, + { url = "https://files.pythonhosted.org/packages/1f/54/6d54d6d6a2019cbffa7316de038002b0250568199f2b4138326eb27e3910/jaxlib-0.7.1-cp313-cp313t-manylinux2014_aarch64.whl", hash = "sha256:58558fa29fd7d6342c066f837e58fcba335182837a959affc128660c089702de", size = 85286808, upload-time = "2025-08-20T15:56:46.087Z" }, + { url = "https://files.pythonhosted.org/packages/79/e2/a0ce0ac2aa9b12a83f3bb27c570361017743c5cf82f1044771daa4c865a6/jaxlib-0.7.1-cp313-cp313t-manylinux_2_27_x86_64.whl", hash = "sha256:944f7555960d69f1d1c435fff0a76e4edd6a878fe47fe1781f7a0d63b61072e5", size = 81252495, upload-time = "2025-08-20T15:56:51.172Z" }, + { url = "https://files.pythonhosted.org/packages/7d/e1/9a5317a520aa964944761bf1d050a6a20d0792d9af5005e25804a1ce0e84/jaxlib-0.7.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:72dd9c3e95457a5a54f00e47ec14863b5e61540b944c0df13bf10c259b4d5d73", size = 57696029, upload-time = "2025-08-20T15:56:55.563Z" }, + { url = "https://files.pythonhosted.org/packages/98/ae/52cb0552b6e1b1008bc68b45ad51091086cccdc1a7e3d32ba9856d7fa879/jaxlib-0.7.1-cp314-cp314-manylinux2014_aarch64.whl", hash = "sha256:023df41212ae4fda869338370f9532bfcd98ccfaee909bb95ea540d6053df547", size = 85146969, upload-time = "2025-08-20T15:57:00.519Z" }, + { url = "https://files.pythonhosted.org/packages/c7/3a/a63b6bd9cac259c22ae324fcb1d5c74e053e8acad23f13ad39ffa6fd11e6/jaxlib-0.7.1-cp314-cp314-manylinux_2_27_x86_64.whl", hash = "sha256:d52817a42c130d0c330f48edcb3a3e455dc984b6ce53f3182c37aa0fe960109b", size = 81167909, upload-time = "2025-08-20T15:57:05.113Z" }, + { url = "https://files.pythonhosted.org/packages/5d/e3/75cc67d24e14279c634ed4c4d501ee62f68d8260aac7d87fd267488aa90a/jaxlib-0.7.1-cp314-cp314-win_amd64.whl", hash = "sha256:38ac46ec17c0a86b428590d04559629357fc6dc3a3c76569f022b982c36fc1af", size = 63566625, upload-time = "2025-08-20T15:57:09.127Z" }, + { url = "https://files.pythonhosted.org/packages/a8/94/f2adf6979a89fcba2ff68b6add47e316bd1274d6ac1284c991d19d8dc2e0/jaxlib-0.7.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6057a602632bd3299d6c7ffbbb3f1ef2c7573ccbed9eb06cc92042b96e2ca5d4", size = 57787077, upload-time = "2025-08-20T15:57:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/f2/64/e581007dabf29be6d916fb201f1af264edabc3f74097be9f02606eae8836/jaxlib-0.7.1-cp314-cp314t-manylinux2014_aarch64.whl", hash = "sha256:9766817cfa51743a48fac78c087605c30bf1a91caf11371ca8c41261e6f3a0c8", size = 85285337, upload-time = "2025-08-20T15:57:16.286Z" }, + { url = "https://files.pythonhosted.org/packages/ee/0e/2e9e5c82f0b8a5a2be470c5185811ab2d780f2da448339e09f71af06f36a/jaxlib-0.7.1-cp314-cp314t-manylinux_2_27_x86_64.whl", hash = "sha256:1ea54e6182d85b82496fbc58bbe51042bea3ee3e1e0d444b3cff446c245bebd5", size = 81256695, upload-time = "2025-08-20T15:57:20.786Z" }, +] + [[package]] name = "jeepney" version = "0.9.0" @@ -1310,6 +1748,126 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d3/32/da7f44bcb1105d3e88a0b74ebdca50c59121d2ddf71c9e34ba47df7f3a56/keyring-25.6.0-py3-none-any.whl", hash = "sha256:552a3f7af126ece7ed5c89753650eec89c7eaae8617d0aa4d9ad2b75111266bd", size = 39085, upload-time = "2024-12-25T15:26:44.377Z" }, ] +[[package]] +name = "kiwisolver" +version = "1.4.9" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5c/3c/85844f1b0feb11ee581ac23fe5fce65cd049a200c1446708cc1b7f922875/kiwisolver-1.4.9.tar.gz", hash = "sha256:c3b22c26c6fd6811b0ae8363b95ca8ce4ea3c202d3d0975b2914310ceb1bcc4d", size = 97564, upload-time = "2025-08-10T21:27:49.279Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/5d/8ce64e36d4e3aac5ca96996457dcf33e34e6051492399a3f1fec5657f30b/kiwisolver-1.4.9-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b4b4d74bda2b8ebf4da5bd42af11d02d04428b2c32846e4c2c93219df8a7987b", size = 124159, upload-time = "2025-08-10T21:25:35.472Z" }, + { url = "https://files.pythonhosted.org/packages/96/1e/22f63ec454874378175a5f435d6ea1363dd33fb2af832c6643e4ccea0dc8/kiwisolver-1.4.9-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:fb3b8132019ea572f4611d770991000d7f58127560c4889729248eb5852a102f", size = 66578, upload-time = "2025-08-10T21:25:36.73Z" }, + { url = "https://files.pythonhosted.org/packages/41/4c/1925dcfff47a02d465121967b95151c82d11027d5ec5242771e580e731bd/kiwisolver-1.4.9-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:84fd60810829c27ae375114cd379da1fa65e6918e1da405f356a775d49a62bcf", size = 65312, upload-time = "2025-08-10T21:25:37.658Z" }, + { url = "https://files.pythonhosted.org/packages/d4/42/0f333164e6307a0687d1eb9ad256215aae2f4bd5d28f4653d6cd319a3ba3/kiwisolver-1.4.9-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:b78efa4c6e804ecdf727e580dbb9cba85624d2e1c6b5cb059c66290063bd99a9", size = 1628458, upload-time = "2025-08-10T21:25:39.067Z" }, + { url = "https://files.pythonhosted.org/packages/86/b6/2dccb977d651943995a90bfe3495c2ab2ba5cd77093d9f2318a20c9a6f59/kiwisolver-1.4.9-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4efec7bcf21671db6a3294ff301d2fc861c31faa3c8740d1a94689234d1b415", size = 1225640, upload-time = "2025-08-10T21:25:40.489Z" }, + { url = "https://files.pythonhosted.org/packages/50/2b/362ebd3eec46c850ccf2bfe3e30f2fc4c008750011f38a850f088c56a1c6/kiwisolver-1.4.9-cp310-cp310-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:90f47e70293fc3688b71271100a1a5453aa9944a81d27ff779c108372cf5567b", size = 1244074, upload-time = "2025-08-10T21:25:42.221Z" }, + { url = "https://files.pythonhosted.org/packages/6f/bb/f09a1e66dab8984773d13184a10a29fe67125337649d26bdef547024ed6b/kiwisolver-1.4.9-cp310-cp310-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8fdca1def57a2e88ef339de1737a1449d6dbf5fab184c54a1fca01d541317154", size = 1293036, upload-time = "2025-08-10T21:25:43.801Z" }, + { url = "https://files.pythonhosted.org/packages/ea/01/11ecf892f201cafda0f68fa59212edaea93e96c37884b747c181303fccd1/kiwisolver-1.4.9-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:9cf554f21be770f5111a1690d42313e140355e687e05cf82cb23d0a721a64a48", size = 2175310, upload-time = "2025-08-10T21:25:45.045Z" }, + { url = "https://files.pythonhosted.org/packages/7f/5f/bfe11d5b934f500cc004314819ea92427e6e5462706a498c1d4fc052e08f/kiwisolver-1.4.9-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:fc1795ac5cd0510207482c3d1d3ed781143383b8cfd36f5c645f3897ce066220", size = 2270943, upload-time = "2025-08-10T21:25:46.393Z" }, + { url = "https://files.pythonhosted.org/packages/3d/de/259f786bf71f1e03e73d87e2db1a9a3bcab64d7b4fd780167123161630ad/kiwisolver-1.4.9-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:ccd09f20ccdbbd341b21a67ab50a119b64a403b09288c27481575105283c1586", size = 2440488, upload-time = "2025-08-10T21:25:48.074Z" }, + { url = "https://files.pythonhosted.org/packages/1b/76/c989c278faf037c4d3421ec07a5c452cd3e09545d6dae7f87c15f54e4edf/kiwisolver-1.4.9-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:540c7c72324d864406a009d72f5d6856f49693db95d1fbb46cf86febef873634", size = 2246787, upload-time = "2025-08-10T21:25:49.442Z" }, + { url = "https://files.pythonhosted.org/packages/a2/55/c2898d84ca440852e560ca9f2a0d28e6e931ac0849b896d77231929900e7/kiwisolver-1.4.9-cp310-cp310-win_amd64.whl", hash = "sha256:ede8c6d533bc6601a47ad4046080d36b8fc99f81e6f1c17b0ac3c2dc91ac7611", size = 73730, upload-time = "2025-08-10T21:25:51.102Z" }, + { url = "https://files.pythonhosted.org/packages/e8/09/486d6ac523dd33b80b368247f238125d027964cfacb45c654841e88fb2ae/kiwisolver-1.4.9-cp310-cp310-win_arm64.whl", hash = "sha256:7b4da0d01ac866a57dd61ac258c5607b4cd677f63abaec7b148354d2b2cdd536", size = 65036, upload-time = "2025-08-10T21:25:52.063Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ab/c80b0d5a9d8a1a65f4f815f2afff9798b12c3b9f31f1d304dd233dd920e2/kiwisolver-1.4.9-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:eb14a5da6dc7642b0f3a18f13654847cd8b7a2550e2645a5bda677862b03ba16", size = 124167, upload-time = "2025-08-10T21:25:53.403Z" }, + { url = "https://files.pythonhosted.org/packages/a0/c0/27fe1a68a39cf62472a300e2879ffc13c0538546c359b86f149cc19f6ac3/kiwisolver-1.4.9-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:39a219e1c81ae3b103643d2aedb90f1ef22650deb266ff12a19e7773f3e5f089", size = 66579, upload-time = "2025-08-10T21:25:54.79Z" }, + { url = "https://files.pythonhosted.org/packages/31/a2/a12a503ac1fd4943c50f9822678e8015a790a13b5490354c68afb8489814/kiwisolver-1.4.9-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2405a7d98604b87f3fc28b1716783534b1b4b8510d8142adca34ee0bc3c87543", size = 65309, upload-time = "2025-08-10T21:25:55.76Z" }, + { url = "https://files.pythonhosted.org/packages/66/e1/e533435c0be77c3f64040d68d7a657771194a63c279f55573188161e81ca/kiwisolver-1.4.9-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:dc1ae486f9abcef254b5618dfb4113dd49f94c68e3e027d03cf0143f3f772b61", size = 1435596, upload-time = "2025-08-10T21:25:56.861Z" }, + { url = "https://files.pythonhosted.org/packages/67/1e/51b73c7347f9aabdc7215aa79e8b15299097dc2f8e67dee2b095faca9cb0/kiwisolver-1.4.9-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a1f570ce4d62d718dce3f179ee78dac3b545ac16c0c04bb363b7607a949c0d1", size = 1246548, upload-time = "2025-08-10T21:25:58.246Z" }, + { url = "https://files.pythonhosted.org/packages/21/aa/72a1c5d1e430294f2d32adb9542719cfb441b5da368d09d268c7757af46c/kiwisolver-1.4.9-cp311-cp311-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cb27e7b78d716c591e88e0a09a2139c6577865d7f2e152488c2cc6257f460872", size = 1263618, upload-time = "2025-08-10T21:25:59.857Z" }, + { url = "https://files.pythonhosted.org/packages/a3/af/db1509a9e79dbf4c260ce0cfa3903ea8945f6240e9e59d1e4deb731b1a40/kiwisolver-1.4.9-cp311-cp311-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:15163165efc2f627eb9687ea5f3a28137217d217ac4024893d753f46bce9de26", size = 1317437, upload-time = "2025-08-10T21:26:01.105Z" }, + { url = "https://files.pythonhosted.org/packages/e0/f2/3ea5ee5d52abacdd12013a94130436e19969fa183faa1e7c7fbc89e9a42f/kiwisolver-1.4.9-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bdee92c56a71d2b24c33a7d4c2856bd6419d017e08caa7802d2963870e315028", size = 2195742, upload-time = "2025-08-10T21:26:02.675Z" }, + { url = "https://files.pythonhosted.org/packages/6f/9b/1efdd3013c2d9a2566aa6a337e9923a00590c516add9a1e89a768a3eb2fc/kiwisolver-1.4.9-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:412f287c55a6f54b0650bd9b6dce5aceddb95864a1a90c87af16979d37c89771", size = 2290810, upload-time = "2025-08-10T21:26:04.009Z" }, + { url = "https://files.pythonhosted.org/packages/fb/e5/cfdc36109ae4e67361f9bc5b41323648cb24a01b9ade18784657e022e65f/kiwisolver-1.4.9-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:2c93f00dcba2eea70af2be5f11a830a742fe6b579a1d4e00f47760ef13be247a", size = 2461579, upload-time = "2025-08-10T21:26:05.317Z" }, + { url = "https://files.pythonhosted.org/packages/62/86/b589e5e86c7610842213994cdea5add00960076bef4ae290c5fa68589cac/kiwisolver-1.4.9-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:f117e1a089d9411663a3207ba874f31be9ac8eaa5b533787024dc07aeb74f464", size = 2268071, upload-time = "2025-08-10T21:26:06.686Z" }, + { url = "https://files.pythonhosted.org/packages/3b/c6/f8df8509fd1eee6c622febe54384a96cfaf4d43bf2ccec7a0cc17e4715c9/kiwisolver-1.4.9-cp311-cp311-win_amd64.whl", hash = "sha256:be6a04e6c79819c9a8c2373317d19a96048e5a3f90bec587787e86a1153883c2", size = 73840, upload-time = "2025-08-10T21:26:07.94Z" }, + { url = "https://files.pythonhosted.org/packages/e2/2d/16e0581daafd147bc11ac53f032a2b45eabac897f42a338d0a13c1e5c436/kiwisolver-1.4.9-cp311-cp311-win_arm64.whl", hash = "sha256:0ae37737256ba2de764ddc12aed4956460277f00c4996d51a197e72f62f5eec7", size = 65159, upload-time = "2025-08-10T21:26:09.048Z" }, + { url = "https://files.pythonhosted.org/packages/86/c9/13573a747838aeb1c76e3267620daa054f4152444d1f3d1a2324b78255b5/kiwisolver-1.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:ac5a486ac389dddcc5bef4f365b6ae3ffff2c433324fb38dd35e3fab7c957999", size = 123686, upload-time = "2025-08-10T21:26:10.034Z" }, + { url = "https://files.pythonhosted.org/packages/51/ea/2ecf727927f103ffd1739271ca19c424d0e65ea473fbaeea1c014aea93f6/kiwisolver-1.4.9-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:f2ba92255faa7309d06fe44c3a4a97efe1c8d640c2a79a5ef728b685762a6fd2", size = 66460, upload-time = "2025-08-10T21:26:11.083Z" }, + { url = "https://files.pythonhosted.org/packages/5b/5a/51f5464373ce2aeb5194508298a508b6f21d3867f499556263c64c621914/kiwisolver-1.4.9-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4a2899935e724dd1074cb568ce7ac0dce28b2cd6ab539c8e001a8578eb106d14", size = 64952, upload-time = "2025-08-10T21:26:12.058Z" }, + { url = "https://files.pythonhosted.org/packages/70/90/6d240beb0f24b74371762873e9b7f499f1e02166a2d9c5801f4dbf8fa12e/kiwisolver-1.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f6008a4919fdbc0b0097089f67a1eb55d950ed7e90ce2cc3e640abadd2757a04", size = 1474756, upload-time = "2025-08-10T21:26:13.096Z" }, + { url = "https://files.pythonhosted.org/packages/12/42/f36816eaf465220f683fb711efdd1bbf7a7005a2473d0e4ed421389bd26c/kiwisolver-1.4.9-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:67bb8b474b4181770f926f7b7d2f8c0248cbcb78b660fdd41a47054b28d2a752", size = 1276404, upload-time = "2025-08-10T21:26:14.457Z" }, + { url = "https://files.pythonhosted.org/packages/2e/64/bc2de94800adc830c476dce44e9b40fd0809cddeef1fde9fcf0f73da301f/kiwisolver-1.4.9-cp312-cp312-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2327a4a30d3ee07d2fbe2e7933e8a37c591663b96ce42a00bc67461a87d7df77", size = 1294410, upload-time = "2025-08-10T21:26:15.73Z" }, + { url = "https://files.pythonhosted.org/packages/5f/42/2dc82330a70aa8e55b6d395b11018045e58d0bb00834502bf11509f79091/kiwisolver-1.4.9-cp312-cp312-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7a08b491ec91b1d5053ac177afe5290adacf1f0f6307d771ccac5de30592d198", size = 1343631, upload-time = "2025-08-10T21:26:17.045Z" }, + { url = "https://files.pythonhosted.org/packages/22/fd/f4c67a6ed1aab149ec5a8a401c323cee7a1cbe364381bb6c9c0d564e0e20/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d8fc5c867c22b828001b6a38d2eaeb88160bf5783c6cb4a5e440efc981ce286d", size = 2224963, upload-time = "2025-08-10T21:26:18.737Z" }, + { url = "https://files.pythonhosted.org/packages/45/aa/76720bd4cb3713314677d9ec94dcc21ced3f1baf4830adde5bb9b2430a5f/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:3b3115b2581ea35bb6d1f24a4c90af37e5d9b49dcff267eeed14c3893c5b86ab", size = 2321295, upload-time = "2025-08-10T21:26:20.11Z" }, + { url = "https://files.pythonhosted.org/packages/80/19/d3ec0d9ab711242f56ae0dc2fc5d70e298bb4a1f9dfab44c027668c673a1/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:858e4c22fb075920b96a291928cb7dea5644e94c0ee4fcd5af7e865655e4ccf2", size = 2487987, upload-time = "2025-08-10T21:26:21.49Z" }, + { url = "https://files.pythonhosted.org/packages/39/e9/61e4813b2c97e86b6fdbd4dd824bf72d28bcd8d4849b8084a357bc0dd64d/kiwisolver-1.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ed0fecd28cc62c54b262e3736f8bb2512d8dcfdc2bcf08be5f47f96bf405b145", size = 2291817, upload-time = "2025-08-10T21:26:22.812Z" }, + { url = "https://files.pythonhosted.org/packages/a0/41/85d82b0291db7504da3c2defe35c9a8a5c9803a730f297bd823d11d5fb77/kiwisolver-1.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:f68208a520c3d86ea51acf688a3e3002615a7f0238002cccc17affecc86a8a54", size = 73895, upload-time = "2025-08-10T21:26:24.37Z" }, + { url = "https://files.pythonhosted.org/packages/e2/92/5f3068cf15ee5cb624a0c7596e67e2a0bb2adee33f71c379054a491d07da/kiwisolver-1.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:2c1a4f57df73965f3f14df20b80ee29e6a7930a57d2d9e8491a25f676e197c60", size = 64992, upload-time = "2025-08-10T21:26:25.732Z" }, + { url = "https://files.pythonhosted.org/packages/31/c1/c2686cda909742ab66c7388e9a1a8521a59eb89f8bcfbee28fc980d07e24/kiwisolver-1.4.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:a5d0432ccf1c7ab14f9949eec60c5d1f924f17c037e9f8b33352fa05799359b8", size = 123681, upload-time = "2025-08-10T21:26:26.725Z" }, + { url = "https://files.pythonhosted.org/packages/ca/f0/f44f50c9f5b1a1860261092e3bc91ecdc9acda848a8b8c6abfda4a24dd5c/kiwisolver-1.4.9-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:efb3a45b35622bb6c16dbfab491a8f5a391fe0e9d45ef32f4df85658232ca0e2", size = 66464, upload-time = "2025-08-10T21:26:27.733Z" }, + { url = "https://files.pythonhosted.org/packages/2d/7a/9d90a151f558e29c3936b8a47ac770235f436f2120aca41a6d5f3d62ae8d/kiwisolver-1.4.9-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1a12cf6398e8a0a001a059747a1cbf24705e18fe413bc22de7b3d15c67cffe3f", size = 64961, upload-time = "2025-08-10T21:26:28.729Z" }, + { url = "https://files.pythonhosted.org/packages/e9/e9/f218a2cb3a9ffbe324ca29a9e399fa2d2866d7f348ec3a88df87fc248fc5/kiwisolver-1.4.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b67e6efbf68e077dd71d1a6b37e43e1a99d0bff1a3d51867d45ee8908b931098", size = 1474607, upload-time = "2025-08-10T21:26:29.798Z" }, + { url = "https://files.pythonhosted.org/packages/d9/28/aac26d4c882f14de59041636292bc838db8961373825df23b8eeb807e198/kiwisolver-1.4.9-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5656aa670507437af0207645273ccdfee4f14bacd7f7c67a4306d0dcaeaf6eed", size = 1276546, upload-time = "2025-08-10T21:26:31.401Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ad/8bfc1c93d4cc565e5069162f610ba2f48ff39b7de4b5b8d93f69f30c4bed/kiwisolver-1.4.9-cp313-cp313-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bfc08add558155345129c7803b3671cf195e6a56e7a12f3dde7c57d9b417f525", size = 1294482, upload-time = "2025-08-10T21:26:32.721Z" }, + { url = "https://files.pythonhosted.org/packages/da/f1/6aca55ff798901d8ce403206d00e033191f63d82dd708a186e0ed2067e9c/kiwisolver-1.4.9-cp313-cp313-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:40092754720b174e6ccf9e845d0d8c7d8e12c3d71e7fc35f55f3813e96376f78", size = 1343720, upload-time = "2025-08-10T21:26:34.032Z" }, + { url = "https://files.pythonhosted.org/packages/d1/91/eed031876c595c81d90d0f6fc681ece250e14bf6998c3d7c419466b523b7/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:497d05f29a1300d14e02e6441cf0f5ee81c1ff5a304b0d9fb77423974684e08b", size = 2224907, upload-time = "2025-08-10T21:26:35.824Z" }, + { url = "https://files.pythonhosted.org/packages/e9/ec/4d1925f2e49617b9cca9c34bfa11adefad49d00db038e692a559454dfb2e/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:bdd1a81a1860476eb41ac4bc1e07b3f07259e6d55bbf739b79c8aaedcf512799", size = 2321334, upload-time = "2025-08-10T21:26:37.534Z" }, + { url = "https://files.pythonhosted.org/packages/43/cb/450cd4499356f68802750c6ddc18647b8ea01ffa28f50d20598e0befe6e9/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:e6b93f13371d341afee3be9f7c5964e3fe61d5fa30f6a30eb49856935dfe4fc3", size = 2488313, upload-time = "2025-08-10T21:26:39.191Z" }, + { url = "https://files.pythonhosted.org/packages/71/67/fc76242bd99f885651128a5d4fa6083e5524694b7c88b489b1b55fdc491d/kiwisolver-1.4.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:d75aa530ccfaa593da12834b86a0724f58bff12706659baa9227c2ccaa06264c", size = 2291970, upload-time = "2025-08-10T21:26:40.828Z" }, + { url = "https://files.pythonhosted.org/packages/75/bd/f1a5d894000941739f2ae1b65a32892349423ad49c2e6d0771d0bad3fae4/kiwisolver-1.4.9-cp313-cp313-win_amd64.whl", hash = "sha256:dd0a578400839256df88c16abddf9ba14813ec5f21362e1fe65022e00c883d4d", size = 73894, upload-time = "2025-08-10T21:26:42.33Z" }, + { url = "https://files.pythonhosted.org/packages/95/38/dce480814d25b99a391abbddadc78f7c117c6da34be68ca8b02d5848b424/kiwisolver-1.4.9-cp313-cp313-win_arm64.whl", hash = "sha256:d4188e73af84ca82468f09cadc5ac4db578109e52acb4518d8154698d3a87ca2", size = 64995, upload-time = "2025-08-10T21:26:43.889Z" }, + { url = "https://files.pythonhosted.org/packages/e2/37/7d218ce5d92dadc5ebdd9070d903e0c7cf7edfe03f179433ac4d13ce659c/kiwisolver-1.4.9-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:5a0f2724dfd4e3b3ac5a82436a8e6fd16baa7d507117e4279b660fe8ca38a3a1", size = 126510, upload-time = "2025-08-10T21:26:44.915Z" }, + { url = "https://files.pythonhosted.org/packages/23/b0/e85a2b48233daef4b648fb657ebbb6f8367696a2d9548a00b4ee0eb67803/kiwisolver-1.4.9-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:1b11d6a633e4ed84fc0ddafd4ebfd8ea49b3f25082c04ad12b8315c11d504dc1", size = 67903, upload-time = "2025-08-10T21:26:45.934Z" }, + { url = "https://files.pythonhosted.org/packages/44/98/f2425bc0113ad7de24da6bb4dae1343476e95e1d738be7c04d31a5d037fd/kiwisolver-1.4.9-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:61874cdb0a36016354853593cffc38e56fc9ca5aa97d2c05d3dcf6922cd55a11", size = 66402, upload-time = "2025-08-10T21:26:47.101Z" }, + { url = "https://files.pythonhosted.org/packages/98/d8/594657886df9f34c4177cc353cc28ca7e6e5eb562d37ccc233bff43bbe2a/kiwisolver-1.4.9-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:60c439763a969a6af93b4881db0eed8fadf93ee98e18cbc35bc8da868d0c4f0c", size = 1582135, upload-time = "2025-08-10T21:26:48.665Z" }, + { url = "https://files.pythonhosted.org/packages/5c/c6/38a115b7170f8b306fc929e166340c24958347308ea3012c2b44e7e295db/kiwisolver-1.4.9-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:92a2f997387a1b79a75e7803aa7ded2cfbe2823852ccf1ba3bcf613b62ae3197", size = 1389409, upload-time = "2025-08-10T21:26:50.335Z" }, + { url = "https://files.pythonhosted.org/packages/bf/3b/e04883dace81f24a568bcee6eb3001da4ba05114afa622ec9b6fafdc1f5e/kiwisolver-1.4.9-cp313-cp313t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a31d512c812daea6d8b3be3b2bfcbeb091dbb09177706569bcfc6240dcf8b41c", size = 1401763, upload-time = "2025-08-10T21:26:51.867Z" }, + { url = "https://files.pythonhosted.org/packages/9f/80/20ace48e33408947af49d7d15c341eaee69e4e0304aab4b7660e234d6288/kiwisolver-1.4.9-cp313-cp313t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:52a15b0f35dad39862d376df10c5230155243a2c1a436e39eb55623ccbd68185", size = 1453643, upload-time = "2025-08-10T21:26:53.592Z" }, + { url = "https://files.pythonhosted.org/packages/64/31/6ce4380a4cd1f515bdda976a1e90e547ccd47b67a1546d63884463c92ca9/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:a30fd6fdef1430fd9e1ba7b3398b5ee4e2887783917a687d86ba69985fb08748", size = 2330818, upload-time = "2025-08-10T21:26:55.051Z" }, + { url = "https://files.pythonhosted.org/packages/fa/e9/3f3fcba3bcc7432c795b82646306e822f3fd74df0ee81f0fa067a1f95668/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_ppc64le.whl", hash = "sha256:cc9617b46837c6468197b5945e196ee9ca43057bb7d9d1ae688101e4e1dddf64", size = 2419963, upload-time = "2025-08-10T21:26:56.421Z" }, + { url = "https://files.pythonhosted.org/packages/99/43/7320c50e4133575c66e9f7dadead35ab22d7c012a3b09bb35647792b2a6d/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_s390x.whl", hash = "sha256:0ab74e19f6a2b027ea4f845a78827969af45ce790e6cb3e1ebab71bdf9f215ff", size = 2594639, upload-time = "2025-08-10T21:26:57.882Z" }, + { url = "https://files.pythonhosted.org/packages/65/d6/17ae4a270d4a987ef8a385b906d2bdfc9fce502d6dc0d3aea865b47f548c/kiwisolver-1.4.9-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:dba5ee5d3981160c28d5490f0d1b7ed730c22470ff7f6cc26cfcfaacb9896a07", size = 2391741, upload-time = "2025-08-10T21:26:59.237Z" }, + { url = "https://files.pythonhosted.org/packages/2a/8f/8f6f491d595a9e5912971f3f863d81baddccc8a4d0c3749d6a0dd9ffc9df/kiwisolver-1.4.9-cp313-cp313t-win_arm64.whl", hash = "sha256:0749fd8f4218ad2e851e11cc4dc05c7cbc0cbc4267bdfdb31782e65aace4ee9c", size = 68646, upload-time = "2025-08-10T21:27:00.52Z" }, + { url = "https://files.pythonhosted.org/packages/6b/32/6cc0fbc9c54d06c2969faa9c1d29f5751a2e51809dd55c69055e62d9b426/kiwisolver-1.4.9-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:9928fe1eb816d11ae170885a74d074f57af3a0d65777ca47e9aeb854a1fba386", size = 123806, upload-time = "2025-08-10T21:27:01.537Z" }, + { url = "https://files.pythonhosted.org/packages/b2/dd/2bfb1d4a4823d92e8cbb420fe024b8d2167f72079b3bb941207c42570bdf/kiwisolver-1.4.9-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:d0005b053977e7b43388ddec89fa567f43d4f6d5c2c0affe57de5ebf290dc552", size = 66605, upload-time = "2025-08-10T21:27:03.335Z" }, + { url = "https://files.pythonhosted.org/packages/f7/69/00aafdb4e4509c2ca6064646cba9cd4b37933898f426756adb2cb92ebbed/kiwisolver-1.4.9-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2635d352d67458b66fd0667c14cb1d4145e9560d503219034a18a87e971ce4f3", size = 64925, upload-time = "2025-08-10T21:27:04.339Z" }, + { url = "https://files.pythonhosted.org/packages/43/dc/51acc6791aa14e5cb6d8a2e28cefb0dc2886d8862795449d021334c0df20/kiwisolver-1.4.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:767c23ad1c58c9e827b649a9ab7809fd5fd9db266a9cf02b0e926ddc2c680d58", size = 1472414, upload-time = "2025-08-10T21:27:05.437Z" }, + { url = "https://files.pythonhosted.org/packages/3d/bb/93fa64a81db304ac8a246f834d5094fae4b13baf53c839d6bb6e81177129/kiwisolver-1.4.9-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:72d0eb9fba308b8311685c2268cf7d0a0639a6cd027d8128659f72bdd8a024b4", size = 1281272, upload-time = "2025-08-10T21:27:07.063Z" }, + { url = "https://files.pythonhosted.org/packages/70/e6/6df102916960fb8d05069d4bd92d6d9a8202d5a3e2444494e7cd50f65b7a/kiwisolver-1.4.9-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f68e4f3eeca8fb22cc3d731f9715a13b652795ef657a13df1ad0c7dc0e9731df", size = 1298578, upload-time = "2025-08-10T21:27:08.452Z" }, + { url = "https://files.pythonhosted.org/packages/7c/47/e142aaa612f5343736b087864dbaebc53ea8831453fb47e7521fa8658f30/kiwisolver-1.4.9-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d84cd4061ae292d8ac367b2c3fa3aad11cb8625a95d135fe93f286f914f3f5a6", size = 1345607, upload-time = "2025-08-10T21:27:10.125Z" }, + { url = "https://files.pythonhosted.org/packages/54/89/d641a746194a0f4d1a3670fb900d0dbaa786fb98341056814bc3f058fa52/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a60ea74330b91bd22a29638940d115df9dc00af5035a9a2a6ad9399ffb4ceca5", size = 2230150, upload-time = "2025-08-10T21:27:11.484Z" }, + { url = "https://files.pythonhosted.org/packages/aa/6b/5ee1207198febdf16ac11f78c5ae40861b809cbe0e6d2a8d5b0b3044b199/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:ce6a3a4e106cf35c2d9c4fa17c05ce0b180db622736845d4315519397a77beaf", size = 2325979, upload-time = "2025-08-10T21:27:12.917Z" }, + { url = "https://files.pythonhosted.org/packages/fc/ff/b269eefd90f4ae14dcc74973d5a0f6d28d3b9bb1afd8c0340513afe6b39a/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:77937e5e2a38a7b48eef0585114fe7930346993a88060d0bf886086d2aa49ef5", size = 2491456, upload-time = "2025-08-10T21:27:14.353Z" }, + { url = "https://files.pythonhosted.org/packages/fc/d4/10303190bd4d30de547534601e259a4fbf014eed94aae3e5521129215086/kiwisolver-1.4.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:24c175051354f4a28c5d6a31c93906dc653e2bf234e8a4bbfb964892078898ce", size = 2294621, upload-time = "2025-08-10T21:27:15.808Z" }, + { url = "https://files.pythonhosted.org/packages/28/e0/a9a90416fce5c0be25742729c2ea52105d62eda6c4be4d803c2a7be1fa50/kiwisolver-1.4.9-cp314-cp314-win_amd64.whl", hash = "sha256:0763515d4df10edf6d06a3c19734e2566368980d21ebec439f33f9eb936c07b7", size = 75417, upload-time = "2025-08-10T21:27:17.436Z" }, + { url = "https://files.pythonhosted.org/packages/1f/10/6949958215b7a9a264299a7db195564e87900f709db9245e4ebdd3c70779/kiwisolver-1.4.9-cp314-cp314-win_arm64.whl", hash = "sha256:0e4e2bf29574a6a7b7f6cb5fa69293b9f96c928949ac4a53ba3f525dffb87f9c", size = 66582, upload-time = "2025-08-10T21:27:18.436Z" }, + { url = "https://files.pythonhosted.org/packages/ec/79/60e53067903d3bc5469b369fe0dfc6b3482e2133e85dae9daa9527535991/kiwisolver-1.4.9-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:d976bbb382b202f71c67f77b0ac11244021cfa3f7dfd9e562eefcea2df711548", size = 126514, upload-time = "2025-08-10T21:27:19.465Z" }, + { url = "https://files.pythonhosted.org/packages/25/d1/4843d3e8d46b072c12a38c97c57fab4608d36e13fe47d47ee96b4d61ba6f/kiwisolver-1.4.9-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:2489e4e5d7ef9a1c300a5e0196e43d9c739f066ef23270607d45aba368b91f2d", size = 67905, upload-time = "2025-08-10T21:27:20.51Z" }, + { url = "https://files.pythonhosted.org/packages/8c/ae/29ffcbd239aea8b93108de1278271ae764dfc0d803a5693914975f200596/kiwisolver-1.4.9-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:e2ea9f7ab7fbf18fffb1b5434ce7c69a07582f7acc7717720f1d69f3e806f90c", size = 66399, upload-time = "2025-08-10T21:27:21.496Z" }, + { url = "https://files.pythonhosted.org/packages/a1/ae/d7ba902aa604152c2ceba5d352d7b62106bedbccc8e95c3934d94472bfa3/kiwisolver-1.4.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b34e51affded8faee0dfdb705416153819d8ea9250bbbf7ea1b249bdeb5f1122", size = 1582197, upload-time = "2025-08-10T21:27:22.604Z" }, + { url = "https://files.pythonhosted.org/packages/f2/41/27c70d427eddb8bc7e4f16420a20fefc6f480312122a59a959fdfe0445ad/kiwisolver-1.4.9-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d8aacd3d4b33b772542b2e01beb50187536967b514b00003bdda7589722d2a64", size = 1390125, upload-time = "2025-08-10T21:27:24.036Z" }, + { url = "https://files.pythonhosted.org/packages/41/42/b3799a12bafc76d962ad69083f8b43b12bf4fe78b097b12e105d75c9b8f1/kiwisolver-1.4.9-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:7cf974dd4e35fa315563ac99d6287a1024e4dc2077b8a7d7cd3d2fb65d283134", size = 1402612, upload-time = "2025-08-10T21:27:25.773Z" }, + { url = "https://files.pythonhosted.org/packages/d2/b5/a210ea073ea1cfaca1bb5c55a62307d8252f531beb364e18aa1e0888b5a0/kiwisolver-1.4.9-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:85bd218b5ecfbee8c8a82e121802dcb519a86044c9c3b2e4aef02fa05c6da370", size = 1453990, upload-time = "2025-08-10T21:27:27.089Z" }, + { url = "https://files.pythonhosted.org/packages/5f/ce/a829eb8c033e977d7ea03ed32fb3c1781b4fa0433fbadfff29e39c676f32/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:0856e241c2d3df4efef7c04a1e46b1936b6120c9bcf36dd216e3acd84bc4fb21", size = 2331601, upload-time = "2025-08-10T21:27:29.343Z" }, + { url = "https://files.pythonhosted.org/packages/e0/4b/b5e97eb142eb9cd0072dacfcdcd31b1c66dc7352b0f7c7255d339c0edf00/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:9af39d6551f97d31a4deebeac6f45b156f9755ddc59c07b402c148f5dbb6482a", size = 2422041, upload-time = "2025-08-10T21:27:30.754Z" }, + { url = "https://files.pythonhosted.org/packages/40/be/8eb4cd53e1b85ba4edc3a9321666f12b83113a178845593307a3e7891f44/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:bb4ae2b57fc1d8cbd1cf7b1d9913803681ffa903e7488012be5b76dedf49297f", size = 2594897, upload-time = "2025-08-10T21:27:32.803Z" }, + { url = "https://files.pythonhosted.org/packages/99/dd/841e9a66c4715477ea0abc78da039832fbb09dac5c35c58dc4c41a407b8a/kiwisolver-1.4.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:aedff62918805fb62d43a4aa2ecd4482c380dc76cd31bd7c8878588a61bd0369", size = 2391835, upload-time = "2025-08-10T21:27:34.23Z" }, + { url = "https://files.pythonhosted.org/packages/0c/28/4b2e5c47a0da96896fdfdb006340ade064afa1e63675d01ea5ac222b6d52/kiwisolver-1.4.9-cp314-cp314t-win_amd64.whl", hash = "sha256:1fa333e8b2ce4d9660f2cda9c0e1b6bafcfb2457a9d259faa82289e73ec24891", size = 79988, upload-time = "2025-08-10T21:27:35.587Z" }, + { url = "https://files.pythonhosted.org/packages/80/be/3578e8afd18c88cdf9cb4cffde75a96d2be38c5a903f1ed0ceec061bd09e/kiwisolver-1.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:4a48a2ce79d65d363597ef7b567ce3d14d68783d2b2263d98db3d9477805ba32", size = 70260, upload-time = "2025-08-10T21:27:36.606Z" }, + { url = "https://files.pythonhosted.org/packages/a2/63/fde392691690f55b38d5dd7b3710f5353bf7a8e52de93a22968801ab8978/kiwisolver-1.4.9-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:4d1d9e582ad4d63062d34077a9a1e9f3c34088a2ec5135b1f7190c07cf366527", size = 60183, upload-time = "2025-08-10T21:27:37.669Z" }, + { url = "https://files.pythonhosted.org/packages/27/b1/6aad34edfdb7cced27f371866f211332bba215bfd918ad3322a58f480d8b/kiwisolver-1.4.9-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:deed0c7258ceb4c44ad5ec7d9918f9f14fd05b2be86378d86cf50e63d1e7b771", size = 58675, upload-time = "2025-08-10T21:27:39.031Z" }, + { url = "https://files.pythonhosted.org/packages/9d/1a/23d855a702bb35a76faed5ae2ba3de57d323f48b1f6b17ee2176c4849463/kiwisolver-1.4.9-pp310-pypy310_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0a590506f303f512dff6b7f75fd2fd18e16943efee932008fe7140e5fa91d80e", size = 80277, upload-time = "2025-08-10T21:27:40.129Z" }, + { url = "https://files.pythonhosted.org/packages/5a/5b/5239e3c2b8fb5afa1e8508f721bb77325f740ab6994d963e61b2b7abcc1e/kiwisolver-1.4.9-pp310-pypy310_pp73-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e09c2279a4d01f099f52d5c4b3d9e208e91edcbd1a175c9662a8b16e000fece9", size = 77994, upload-time = "2025-08-10T21:27:41.181Z" }, + { url = "https://files.pythonhosted.org/packages/f9/1c/5d4d468fb16f8410e596ed0eac02d2c68752aa7dc92997fe9d60a7147665/kiwisolver-1.4.9-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:c9e7cdf45d594ee04d5be1b24dd9d49f3d1590959b2271fb30b5ca2b262c00fb", size = 73744, upload-time = "2025-08-10T21:27:42.254Z" }, + { url = "https://files.pythonhosted.org/packages/a3/0f/36d89194b5a32c054ce93e586d4049b6c2c22887b0eb229c61c68afd3078/kiwisolver-1.4.9-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:720e05574713db64c356e86732c0f3c5252818d05f9df320f0ad8380641acea5", size = 60104, upload-time = "2025-08-10T21:27:43.287Z" }, + { url = "https://files.pythonhosted.org/packages/52/ba/4ed75f59e4658fd21fe7dde1fee0ac397c678ec3befba3fe6482d987af87/kiwisolver-1.4.9-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:17680d737d5335b552994a2008fab4c851bcd7de33094a82067ef3a576ff02fa", size = 58592, upload-time = "2025-08-10T21:27:44.314Z" }, + { url = "https://files.pythonhosted.org/packages/33/01/a8ea7c5ea32a9b45ceeaee051a04c8ed4320f5add3c51bfa20879b765b70/kiwisolver-1.4.9-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:85b5352f94e490c028926ea567fc569c52ec79ce131dadb968d3853e809518c2", size = 80281, upload-time = "2025-08-10T21:27:45.369Z" }, + { url = "https://files.pythonhosted.org/packages/da/e3/dbd2ecdce306f1d07a1aaf324817ee993aab7aee9db47ceac757deabafbe/kiwisolver-1.4.9-pp311-pypy311_pp73-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:464415881e4801295659462c49461a24fb107c140de781d55518c4b80cb6790f", size = 78009, upload-time = "2025-08-10T21:27:46.376Z" }, + { url = "https://files.pythonhosted.org/packages/da/e9/0d4add7873a73e462aeb45c036a2dead2562b825aa46ba326727b3f31016/kiwisolver-1.4.9-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:fb940820c63a9590d31d88b815e7a3aa5915cad3ce735ab45f0c730b39547de1", size = 73929, upload-time = "2025-08-10T21:27:48.236Z" }, +] + +[[package]] +name = "lazy-loader" +version = "0.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6f/6b/c875b30a1ba490860c93da4cabf479e03f584eba06fe5963f6f6644653d8/lazy_loader-0.4.tar.gz", hash = "sha256:47c75182589b91a4e1a85a136c074285a5ad4d9f39c63e0d7fb76391c4574cd1", size = 15431, upload-time = "2024-04-05T13:03:12.261Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/60/d497a310bde3f01cb805196ac61b7ad6dc5dcf8dce66634dc34364b20b4f/lazy_loader-0.4-py3-none-any.whl", hash = "sha256:342aa8e14d543a154047afb4ba8ef17f5563baad3fc610d7b15b213b0f119efc", size = 12097, upload-time = "2024-04-05T13:03:10.514Z" }, +] + [[package]] name = "lmdb" version = "1.7.5" @@ -1452,6 +2010,80 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, ] +[[package]] +name = "matplotlib" +version = "3.10.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "contourpy", version = "1.3.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "contourpy", version = "1.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "cycler" }, + { name = "fonttools" }, + { name = "kiwisolver" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "pyparsing" }, + { name = "python-dateutil" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8a/76/d3c6e3a13fe484ebe7718d14e269c9569c4eb0020a968a327acb3b9a8fe6/matplotlib-3.10.8.tar.gz", hash = "sha256:2299372c19d56bcd35cf05a2738308758d32b9eaed2371898d8f5bd33f084aa3", size = 34806269, upload-time = "2025-12-10T22:56:51.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/58/be/a30bd917018ad220c400169fba298f2bb7003c8ccbc0c3e24ae2aacad1e8/matplotlib-3.10.8-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:00270d217d6b20d14b584c521f810d60c5c78406dc289859776550df837dcda7", size = 8239828, upload-time = "2025-12-10T22:55:02.313Z" }, + { url = "https://files.pythonhosted.org/packages/58/27/ca01e043c4841078e82cf6e80a6993dfecd315c3d79f5f3153afbb8e1ec6/matplotlib-3.10.8-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:37b3c1cc42aa184b3f738cfa18c1c1d72fd496d85467a6cf7b807936d39aa656", size = 8128050, upload-time = "2025-12-10T22:55:04.997Z" }, + { url = "https://files.pythonhosted.org/packages/cb/aa/7ab67f2b729ae6a91bcf9dcac0affb95fb8c56f7fd2b2af894ae0b0cf6fa/matplotlib-3.10.8-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ee40c27c795bda6a5292e9cff9890189d32f7e3a0bf04e0e3c9430c4a00c37df", size = 8700452, upload-time = "2025-12-10T22:55:07.47Z" }, + { url = "https://files.pythonhosted.org/packages/73/ae/2d5817b0acee3c49b7e7ccfbf5b273f284957cc8e270adf36375db353190/matplotlib-3.10.8-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a48f2b74020919552ea25d222d5cc6af9ca3f4eb43a93e14d068457f545c2a17", size = 9534928, upload-time = "2025-12-10T22:55:10.566Z" }, + { url = "https://files.pythonhosted.org/packages/c9/5b/8e66653e9f7c39cb2e5cab25fce4810daffa2bff02cbf5f3077cea9e942c/matplotlib-3.10.8-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:f254d118d14a7f99d616271d6c3c27922c092dac11112670b157798b89bf4933", size = 9586377, upload-time = "2025-12-10T22:55:12.362Z" }, + { url = "https://files.pythonhosted.org/packages/e2/e2/fd0bbadf837f81edb0d208ba8f8cb552874c3b16e27cb91a31977d90875d/matplotlib-3.10.8-cp310-cp310-win_amd64.whl", hash = "sha256:f9b587c9c7274c1613a30afabf65a272114cd6cdbe67b3406f818c79d7ab2e2a", size = 8128127, upload-time = "2025-12-10T22:55:14.436Z" }, + { url = "https://files.pythonhosted.org/packages/f8/86/de7e3a1cdcfc941483af70609edc06b83e7c8a0e0dc9ac325200a3f4d220/matplotlib-3.10.8-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:6be43b667360fef5c754dda5d25a32e6307a03c204f3c0fc5468b78fa87b4160", size = 8251215, upload-time = "2025-12-10T22:55:16.175Z" }, + { url = "https://files.pythonhosted.org/packages/fd/14/baad3222f424b19ce6ad243c71de1ad9ec6b2e4eb1e458a48fdc6d120401/matplotlib-3.10.8-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a2b336e2d91a3d7006864e0990c83b216fcdca64b5a6484912902cef87313d78", size = 8139625, upload-time = "2025-12-10T22:55:17.712Z" }, + { url = "https://files.pythonhosted.org/packages/8f/a0/7024215e95d456de5883e6732e708d8187d9753a21d32f8ddb3befc0c445/matplotlib-3.10.8-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:efb30e3baaea72ce5928e32bab719ab4770099079d66726a62b11b1ef7273be4", size = 8712614, upload-time = "2025-12-10T22:55:20.8Z" }, + { url = "https://files.pythonhosted.org/packages/5a/f4/b8347351da9a5b3f41e26cf547252d861f685c6867d179a7c9d60ad50189/matplotlib-3.10.8-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d56a1efd5bfd61486c8bc968fa18734464556f0fb8e51690f4ac25d85cbbbbc2", size = 9540997, upload-time = "2025-12-10T22:55:23.258Z" }, + { url = "https://files.pythonhosted.org/packages/9e/c0/c7b914e297efe0bc36917bf216b2acb91044b91e930e878ae12981e461e5/matplotlib-3.10.8-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:238b7ce5717600615c895050239ec955d91f321c209dd110db988500558e70d6", size = 9596825, upload-time = "2025-12-10T22:55:25.217Z" }, + { url = "https://files.pythonhosted.org/packages/6f/d3/a4bbc01c237ab710a1f22b4da72f4ff6d77eb4c7735ea9811a94ae239067/matplotlib-3.10.8-cp311-cp311-win_amd64.whl", hash = "sha256:18821ace09c763ec93aef5eeff087ee493a24051936d7b9ebcad9662f66501f9", size = 8135090, upload-time = "2025-12-10T22:55:27.162Z" }, + { url = "https://files.pythonhosted.org/packages/89/dd/a0b6588f102beab33ca6f5218b31725216577b2a24172f327eaf6417d5c9/matplotlib-3.10.8-cp311-cp311-win_arm64.whl", hash = "sha256:bab485bcf8b1c7d2060b4fcb6fc368a9e6f4cd754c9c2fea281f4be21df394a2", size = 8012377, upload-time = "2025-12-10T22:55:29.185Z" }, + { url = "https://files.pythonhosted.org/packages/9e/67/f997cdcbb514012eb0d10cd2b4b332667997fb5ebe26b8d41d04962fa0e6/matplotlib-3.10.8-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:64fcc24778ca0404ce0cb7b6b77ae1f4c7231cdd60e6778f999ee05cbd581b9a", size = 8260453, upload-time = "2025-12-10T22:55:30.709Z" }, + { url = "https://files.pythonhosted.org/packages/7e/65/07d5f5c7f7c994f12c768708bd2e17a4f01a2b0f44a1c9eccad872433e2e/matplotlib-3.10.8-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b9a5ca4ac220a0cdd1ba6bcba3608547117d30468fefce49bb26f55c1a3d5c58", size = 8148321, upload-time = "2025-12-10T22:55:33.265Z" }, + { url = "https://files.pythonhosted.org/packages/3e/f3/c5195b1ae57ef85339fd7285dfb603b22c8b4e79114bae5f4f0fcf688677/matplotlib-3.10.8-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3ab4aabc72de4ff77b3ec33a6d78a68227bf1123465887f9905ba79184a1cc04", size = 8716944, upload-time = "2025-12-10T22:55:34.922Z" }, + { url = "https://files.pythonhosted.org/packages/00/f9/7638f5cc82ec8a7aa005de48622eecc3ed7c9854b96ba15bd76b7fd27574/matplotlib-3.10.8-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:24d50994d8c5816ddc35411e50a86ab05f575e2530c02752e02538122613371f", size = 9550099, upload-time = "2025-12-10T22:55:36.789Z" }, + { url = "https://files.pythonhosted.org/packages/57/61/78cd5920d35b29fd2a0fe894de8adf672ff52939d2e9b43cb83cd5ce1bc7/matplotlib-3.10.8-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:99eefd13c0dc3b3c1b4d561c1169e65fe47aab7b8158754d7c084088e2329466", size = 9613040, upload-time = "2025-12-10T22:55:38.715Z" }, + { url = "https://files.pythonhosted.org/packages/30/4e/c10f171b6e2f44d9e3a2b96efa38b1677439d79c99357600a62cc1e9594e/matplotlib-3.10.8-cp312-cp312-win_amd64.whl", hash = "sha256:dd80ecb295460a5d9d260df63c43f4afbdd832d725a531f008dad1664f458adf", size = 8142717, upload-time = "2025-12-10T22:55:41.103Z" }, + { url = "https://files.pythonhosted.org/packages/f1/76/934db220026b5fef85f45d51a738b91dea7d70207581063cd9bd8fafcf74/matplotlib-3.10.8-cp312-cp312-win_arm64.whl", hash = "sha256:3c624e43ed56313651bc18a47f838b60d7b8032ed348911c54906b130b20071b", size = 8012751, upload-time = "2025-12-10T22:55:42.684Z" }, + { url = "https://files.pythonhosted.org/packages/3d/b9/15fd5541ef4f5b9a17eefd379356cf12175fe577424e7b1d80676516031a/matplotlib-3.10.8-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:3f2e409836d7f5ac2f1c013110a4d50b9f7edc26328c108915f9075d7d7a91b6", size = 8261076, upload-time = "2025-12-10T22:55:44.648Z" }, + { url = "https://files.pythonhosted.org/packages/8d/a0/2ba3473c1b66b9c74dc7107c67e9008cb1782edbe896d4c899d39ae9cf78/matplotlib-3.10.8-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:56271f3dac49a88d7fca5060f004d9d22b865f743a12a23b1e937a0be4818ee1", size = 8148794, upload-time = "2025-12-10T22:55:46.252Z" }, + { url = "https://files.pythonhosted.org/packages/75/97/a471f1c3eb1fd6f6c24a31a5858f443891d5127e63a7788678d14e249aea/matplotlib-3.10.8-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a0a7f52498f72f13d4a25ea70f35f4cb60642b466cbb0a9be951b5bc3f45a486", size = 8718474, upload-time = "2025-12-10T22:55:47.864Z" }, + { url = "https://files.pythonhosted.org/packages/01/be/cd478f4b66f48256f42927d0acbcd63a26a893136456cd079c0cc24fbabf/matplotlib-3.10.8-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:646d95230efb9ca614a7a594d4fcacde0ac61d25e37dd51710b36477594963ce", size = 9549637, upload-time = "2025-12-10T22:55:50.048Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7c/8dc289776eae5109e268c4fb92baf870678dc048a25d4ac903683b86d5bf/matplotlib-3.10.8-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f89c151aab2e2e23cb3fe0acad1e8b82841fd265379c4cecd0f3fcb34c15e0f6", size = 9613678, upload-time = "2025-12-10T22:55:52.21Z" }, + { url = "https://files.pythonhosted.org/packages/64/40/37612487cc8a437d4dd261b32ca21fe2d79510fe74af74e1f42becb1bdb8/matplotlib-3.10.8-cp313-cp313-win_amd64.whl", hash = "sha256:e8ea3e2d4066083e264e75c829078f9e149fa119d27e19acd503de65e0b13149", size = 8142686, upload-time = "2025-12-10T22:55:54.253Z" }, + { url = "https://files.pythonhosted.org/packages/66/52/8d8a8730e968185514680c2a6625943f70269509c3dcfc0dcf7d75928cb8/matplotlib-3.10.8-cp313-cp313-win_arm64.whl", hash = "sha256:c108a1d6fa78a50646029cb6d49808ff0fc1330fda87fa6f6250c6b5369b6645", size = 8012917, upload-time = "2025-12-10T22:55:56.268Z" }, + { url = "https://files.pythonhosted.org/packages/b5/27/51fe26e1062f298af5ef66343d8ef460e090a27fea73036c76c35821df04/matplotlib-3.10.8-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:ad3d9833a64cf48cc4300f2b406c3d0f4f4724a91c0bd5640678a6ba7c102077", size = 8305679, upload-time = "2025-12-10T22:55:57.856Z" }, + { url = "https://files.pythonhosted.org/packages/2c/1e/4de865bc591ac8e3062e835f42dd7fe7a93168d519557837f0e37513f629/matplotlib-3.10.8-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:eb3823f11823deade26ce3b9f40dcb4a213da7a670013929f31d5f5ed1055b22", size = 8198336, upload-time = "2025-12-10T22:55:59.371Z" }, + { url = "https://files.pythonhosted.org/packages/c6/cb/2f7b6e75fb4dce87ef91f60cac4f6e34f4c145ab036a22318ec837971300/matplotlib-3.10.8-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d9050fee89a89ed57b4fb2c1bfac9a3d0c57a0d55aed95949eedbc42070fea39", size = 8731653, upload-time = "2025-12-10T22:56:01.032Z" }, + { url = "https://files.pythonhosted.org/packages/46/b3/bd9c57d6ba670a37ab31fb87ec3e8691b947134b201f881665b28cc039ff/matplotlib-3.10.8-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b44d07310e404ba95f8c25aa5536f154c0a8ec473303535949e52eb71d0a1565", size = 9561356, upload-time = "2025-12-10T22:56:02.95Z" }, + { url = "https://files.pythonhosted.org/packages/c0/3d/8b94a481456dfc9dfe6e39e93b5ab376e50998cddfd23f4ae3b431708f16/matplotlib-3.10.8-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:0a33deb84c15ede243aead39f77e990469fff93ad1521163305095b77b72ce4a", size = 9614000, upload-time = "2025-12-10T22:56:05.411Z" }, + { url = "https://files.pythonhosted.org/packages/bd/cd/bc06149fe5585ba800b189a6a654a75f1f127e8aab02fd2be10df7fa500c/matplotlib-3.10.8-cp313-cp313t-win_amd64.whl", hash = "sha256:3a48a78d2786784cc2413e57397981fb45c79e968d99656706018d6e62e57958", size = 8220043, upload-time = "2025-12-10T22:56:07.551Z" }, + { url = "https://files.pythonhosted.org/packages/e3/de/b22cf255abec916562cc04eef457c13e58a1990048de0c0c3604d082355e/matplotlib-3.10.8-cp313-cp313t-win_arm64.whl", hash = "sha256:15d30132718972c2c074cd14638c7f4592bd98719e2308bccea40e0538bc0cb5", size = 8062075, upload-time = "2025-12-10T22:56:09.178Z" }, + { url = "https://files.pythonhosted.org/packages/3c/43/9c0ff7a2f11615e516c3b058e1e6e8f9614ddeca53faca06da267c48345d/matplotlib-3.10.8-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:b53285e65d4fa4c86399979e956235deb900be5baa7fc1218ea67fbfaeaadd6f", size = 8262481, upload-time = "2025-12-10T22:56:10.885Z" }, + { url = "https://files.pythonhosted.org/packages/6f/ca/e8ae28649fcdf039fda5ef554b40a95f50592a3c47e6f7270c9561c12b07/matplotlib-3.10.8-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:32f8dce744be5569bebe789e46727946041199030db8aeb2954d26013a0eb26b", size = 8151473, upload-time = "2025-12-10T22:56:12.377Z" }, + { url = "https://files.pythonhosted.org/packages/f1/6f/009d129ae70b75e88cbe7e503a12a4c0670e08ed748a902c2568909e9eb5/matplotlib-3.10.8-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4cf267add95b1c88300d96ca837833d4112756045364f5c734a2276038dae27d", size = 9553896, upload-time = "2025-12-10T22:56:14.432Z" }, + { url = "https://files.pythonhosted.org/packages/f5/26/4221a741eb97967bc1fd5e4c52b9aa5a91b2f4ec05b59f6def4d820f9df9/matplotlib-3.10.8-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2cf5bd12cecf46908f286d7838b2abc6c91cda506c0445b8223a7c19a00df008", size = 9824193, upload-time = "2025-12-10T22:56:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f3/3abf75f38605772cf48a9daf5821cd4f563472f38b4b828c6fba6fa6d06e/matplotlib-3.10.8-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:41703cc95688f2516b480f7f339d8851a6035f18e100ee6a32bc0b8536a12a9c", size = 9615444, upload-time = "2025-12-10T22:56:18.155Z" }, + { url = "https://files.pythonhosted.org/packages/93/a5/de89ac80f10b8dc615807ee1133cd99ac74082581196d4d9590bea10690d/matplotlib-3.10.8-cp314-cp314-win_amd64.whl", hash = "sha256:83d282364ea9f3e52363da262ce32a09dfe241e4080dcedda3c0db059d3c1f11", size = 8272719, upload-time = "2025-12-10T22:56:20.366Z" }, + { url = "https://files.pythonhosted.org/packages/69/ce/b006495c19ccc0a137b48083168a37bd056392dee02f87dba0472f2797fe/matplotlib-3.10.8-cp314-cp314-win_arm64.whl", hash = "sha256:2c1998e92cd5999e295a731bcb2911c75f597d937341f3030cc24ef2733d78a8", size = 8144205, upload-time = "2025-12-10T22:56:22.239Z" }, + { url = "https://files.pythonhosted.org/packages/68/d9/b31116a3a855bd313c6fcdb7226926d59b041f26061c6c5b1be66a08c826/matplotlib-3.10.8-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:b5a2b97dbdc7d4f353ebf343744f1d1f1cca8aa8bfddb4262fcf4306c3761d50", size = 8305785, upload-time = "2025-12-10T22:56:24.218Z" }, + { url = "https://files.pythonhosted.org/packages/1e/90/6effe8103f0272685767ba5f094f453784057072f49b393e3ea178fe70a5/matplotlib-3.10.8-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:3f5c3e4da343bba819f0234186b9004faba952cc420fbc522dc4e103c1985908", size = 8198361, upload-time = "2025-12-10T22:56:26.787Z" }, + { url = "https://files.pythonhosted.org/packages/d7/65/a73188711bea603615fc0baecca1061429ac16940e2385433cc778a9d8e7/matplotlib-3.10.8-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5f62550b9a30afde8c1c3ae450e5eb547d579dd69b25c2fc7a1c67f934c1717a", size = 9561357, upload-time = "2025-12-10T22:56:28.953Z" }, + { url = "https://files.pythonhosted.org/packages/f4/3d/b5c5d5d5be8ce63292567f0e2c43dde9953d3ed86ac2de0a72e93c8f07a1/matplotlib-3.10.8-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:495672de149445ec1b772ff2c9ede9b769e3cb4f0d0aa7fa730d7f59e2d4e1c1", size = 9823610, upload-time = "2025-12-10T22:56:31.455Z" }, + { url = "https://files.pythonhosted.org/packages/4d/4b/e7beb6bbd49f6bae727a12b270a2654d13c397576d25bd6786e47033300f/matplotlib-3.10.8-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:595ba4d8fe983b88f0eec8c26a241e16d6376fe1979086232f481f8f3f67494c", size = 9614011, upload-time = "2025-12-10T22:56:33.85Z" }, + { url = "https://files.pythonhosted.org/packages/7c/e6/76f2813d31f032e65f6f797e3f2f6e4aab95b65015924b1c51370395c28a/matplotlib-3.10.8-cp314-cp314t-win_amd64.whl", hash = "sha256:25d380fe8b1dc32cf8f0b1b448470a77afb195438bafdf1d858bfb876f3edf7b", size = 8362801, upload-time = "2025-12-10T22:56:36.107Z" }, + { url = "https://files.pythonhosted.org/packages/5d/49/d651878698a0b67f23aa28e17f45a6d6dd3d3f933fa29087fa4ce5947b5a/matplotlib-3.10.8-cp314-cp314t-win_arm64.whl", hash = "sha256:113bb52413ea508ce954a02c10ffd0d565f9c3bc7f2eddc27dfe1731e71c7b5f", size = 8192560, upload-time = "2025-12-10T22:56:38.008Z" }, + { url = "https://files.pythonhosted.org/packages/f5/43/31d59500bb950b0d188e149a2e552040528c13d6e3d6e84d0cccac593dcd/matplotlib-3.10.8-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:f97aeb209c3d2511443f8797e3e5a569aebb040d4f8bc79aa3ee78a8fb9e3dd8", size = 8237252, upload-time = "2025-12-10T22:56:39.529Z" }, + { url = "https://files.pythonhosted.org/packages/0c/2c/615c09984f3c5f907f51c886538ad785cf72e0e11a3225de2c0f9442aecc/matplotlib-3.10.8-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:fb061f596dad3a0f52b60dc6a5dec4a0c300dec41e058a7efe09256188d170b7", size = 8124693, upload-time = "2025-12-10T22:56:41.758Z" }, + { url = "https://files.pythonhosted.org/packages/91/e1/2757277a1c56041e1fc104b51a0f7b9a4afc8eb737865d63cababe30bc61/matplotlib-3.10.8-pp310-pypy310_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:12d90df9183093fcd479f4172ac26b322b1248b15729cb57f42f71f24c7e37a3", size = 8702205, upload-time = "2025-12-10T22:56:43.415Z" }, + { url = "https://files.pythonhosted.org/packages/04/30/3afaa31c757f34b7725ab9d2ba8b48b5e89c2019c003e7d0ead143aabc5a/matplotlib-3.10.8-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:6da7c2ce169267d0d066adcf63758f0604aa6c3eebf67458930f9d9b79ad1db1", size = 8249198, upload-time = "2025-12-10T22:56:45.584Z" }, + { url = "https://files.pythonhosted.org/packages/48/2f/6334aec331f57485a642a7c8be03cb286f29111ae71c46c38b363230063c/matplotlib-3.10.8-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:9153c3292705be9f9c64498a8872118540c3f4123d1a1c840172edf262c8be4a", size = 8136817, upload-time = "2025-12-10T22:56:47.339Z" }, + { url = "https://files.pythonhosted.org/packages/73/e4/6d6f14b2a759c622f191b2d67e9075a3f56aaccb3be4bb9bb6890030d0a0/matplotlib-3.10.8-pp311-pypy311_pp73-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ae029229a57cd1e8fe542485f27e7ca7b23aa9e8944ddb4985d0bc444f1eca2", size = 8713867, upload-time = "2025-12-10T22:56:48.954Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -1461,6 +2093,132 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "mediapipe" +version = "0.10.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "absl-py" }, + { name = "attrs" }, + { name = "flatbuffers" }, + { name = "jax", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" }, + { name = "jax", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "jaxlib", version = "0.5.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or python_full_version >= '3.13'" }, + { name = "jaxlib", version = "0.7.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.13'" }, + { name = "matplotlib" }, + { name = "numpy" }, + { name = "opencv-contrib-python" }, + { name = "protobuf" }, + { name = "sentencepiece" }, + { name = "sounddevice" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/66/85/5be11ee221b04cd1496131222199ad6469bd49e377157b84b0895325afb2/mediapipe-0.10.21-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:b7cf0756f38f81b6d1cbc77bf78aeba53e4c99270a400388fdecabc3350264d8", size = 49161075, upload-time = "2025-02-06T19:06:01.165Z" }, + { url = "https://files.pythonhosted.org/packages/d4/cc/5a7b367f444b7524d907209af796b3be46d35a90f53ee40ee614b03c7a5b/mediapipe-0.10.21-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:1861f24bd64fbec87364e4d1ea1c4a648d94a8cafbd2def249039c03bd63c0be", size = 49069365, upload-time = "2025-02-06T19:06:10.207Z" }, + { url = "https://files.pythonhosted.org/packages/41/9f/838762bcfd3236d66e196b8c076c8bd26d749b9c52947a6b9201be034668/mediapipe-0.10.21-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:b838d2fc85e7b17e7dd3ceabf2d5b832f6834438c8065506cd2d4ef0aa9f83cc", size = 35622771, upload-time = "2025-02-06T19:06:16.206Z" }, + { url = "https://files.pythonhosted.org/packages/de/88/aa7fd24beb9d47283fdc271032c749be795c0cfe98eee13f624d6e1b15f4/mediapipe-0.10.21-cp310-cp310-win_amd64.whl", hash = "sha256:05cf8b57130c723f12ecb8273fea79038e86775925659cdfe61674627b9da65e", size = 50965067, upload-time = "2025-02-06T19:06:24.419Z" }, + { url = "https://files.pythonhosted.org/packages/6e/37/bd9a3cc4d9b8af54fe40a4858ddd7397e163e079dd67eccec7a948f598a9/mediapipe-0.10.21-cp311-cp311-macosx_11_0_universal2.whl", hash = "sha256:b4946abad54d62812ce4c543187373a4872985fd66cd43fe7437f4b95d53d9ec", size = 49160845, upload-time = "2025-02-06T19:06:32.207Z" }, + { url = "https://files.pythonhosted.org/packages/c5/78/b1bd10f9941eb5853a669d4a72d96010b43c45a8241a90a3c017c543697b/mediapipe-0.10.21-cp311-cp311-macosx_11_0_x86_64.whl", hash = "sha256:fabb8c4703b982dfbc918e939d340202db5d74603f7948c98eeef56166a6cac4", size = 49069767, upload-time = "2025-02-06T19:06:39.111Z" }, + { url = "https://files.pythonhosted.org/packages/ad/83/56f760fecdc60de84c529d6c05c9dfc7d972617c6632bd254d5e021e5b16/mediapipe-0.10.21-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:05dc4a9e593655a79558d05d6227d31018c2537a4bd3362b51e230cf22aecfe3", size = 35622638, upload-time = "2025-02-06T19:06:45.952Z" }, + { url = "https://files.pythonhosted.org/packages/c5/1f/0b2bc80d76d996801113c800e386f7b33aeaec663e2e8ee83695c0dfbde8/mediapipe-0.10.21-cp311-cp311-win_amd64.whl", hash = "sha256:c9ad3adcf2b606f0ebde89afce29da8852f1c93f4fcb7f8dc93d72313625aa6d", size = 50964784, upload-time = "2025-02-06T19:06:52.806Z" }, + { url = "https://files.pythonhosted.org/packages/f1/df/b1912b0d4e88c70a969258b201737c92ea268a90650c470ee842833a6690/mediapipe-0.10.21-cp312-cp312-macosx_11_0_universal2.whl", hash = "sha256:a65508f8c0e28a73f519c8a717f95949b9d5b2fb3bf4a632b181cdb71224d62e", size = 49176855, upload-time = "2025-02-06T19:07:02.525Z" }, + { url = "https://files.pythonhosted.org/packages/89/34/f490f2e1614c021679fe757b08fdedd6100fd18f21111532af8864d006fe/mediapipe-0.10.21-cp312-cp312-macosx_11_0_x86_64.whl", hash = "sha256:96bf0dafa9851c74b1f930090193f23c257ccf4bb1bdebbd2ca3a02397011c0e", size = 49080132, upload-time = "2025-02-06T19:07:11.306Z" }, + { url = "https://files.pythonhosted.org/packages/9f/99/5da7ae7f7e25847383bc2fe5a9adc7ce150dd371682f486c0666b407cad7/mediapipe-0.10.21-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:956eb1ebc275c629e61b085b2cab89c3a5b9e93bad1bb107348d98dafb5a4bb5", size = 35627340, upload-time = "2025-02-06T19:07:19.82Z" }, + { url = "https://files.pythonhosted.org/packages/b7/79/b77808f8195f229ef0c15875540dfdd36724748a4b3de53d993f23336839/mediapipe-0.10.21-cp312-cp312-win_amd64.whl", hash = "sha256:d07b5e69308411161286ea4e48be2ea3e9ab62745143fcb0cb4acced0517e341", size = 50967082, upload-time = "2025-02-06T19:07:26.979Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version >= '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fd/15/76f86faa0902836cc133939732f7611ace68cf54148487a99c539c272dc8/ml_dtypes-0.4.1.tar.gz", hash = "sha256:fad5f2de464fd09127e49b7fd1252b9006fb43d2edc1ff112d390c324af5ca7a", size = 692594, upload-time = "2024-09-13T19:07:11.624Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/9e/76b84f77c7afee3b116dc8407903a2d5004ba3059a8f3dcdcfa6ebf33fff/ml_dtypes-0.4.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:1fe8b5b5e70cd67211db94b05cfd58dace592f24489b038dc6f9fe347d2e07d5", size = 397975, upload-time = "2024-09-13T19:06:44.265Z" }, + { url = "https://files.pythonhosted.org/packages/03/7b/32650e1b2a2713a5923a0af2a8503d0d4a8fc99d1e1e0a1c40e996634460/ml_dtypes-0.4.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8c09a6d11d8475c2a9fd2bc0695628aec105f97cab3b3a3fb7c9660348ff7d24", size = 2182570, upload-time = "2024-09-13T19:06:46.189Z" }, + { url = "https://files.pythonhosted.org/packages/16/86/a9f7569e7e4f5395f927de38a13b92efa73f809285d04f2923b291783dd2/ml_dtypes-0.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9f5e8f75fa371020dd30f9196e7d73babae2abd51cf59bdd56cb4f8de7e13354", size = 2160365, upload-time = "2024-09-13T19:06:48.198Z" }, + { url = "https://files.pythonhosted.org/packages/04/1b/9a3afb437702503514f3934ec8d7904270edf013d28074f3e700e5dfbb0f/ml_dtypes-0.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:15fdd922fea57e493844e5abb930b9c0bd0af217d9edd3724479fc3d7ce70e3f", size = 126633, upload-time = "2024-09-13T19:06:50.656Z" }, + { url = "https://files.pythonhosted.org/packages/d1/76/9835c8609c29f2214359e88f29255fc4aad4ea0f613fb48aa8815ceda1b6/ml_dtypes-0.4.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2d55b588116a7085d6e074cf0cdb1d6fa3875c059dddc4d2c94a4cc81c23e975", size = 397973, upload-time = "2024-09-13T19:06:51.748Z" }, + { url = "https://files.pythonhosted.org/packages/7e/99/e68c56fac5de973007a10254b6e17a0362393724f40f66d5e4033f4962c2/ml_dtypes-0.4.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e138a9b7a48079c900ea969341a5754019a1ad17ae27ee330f7ebf43f23877f9", size = 2185134, upload-time = "2024-09-13T19:06:53.197Z" }, + { url = "https://files.pythonhosted.org/packages/28/bc/6a2344338ea7b61cd7b46fb24ec459360a5a0903b57c55b156c1e46c644a/ml_dtypes-0.4.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:74c6cfb5cf78535b103fde9ea3ded8e9f16f75bc07789054edc7776abfb3d752", size = 2163661, upload-time = "2024-09-13T19:06:54.519Z" }, + { url = "https://files.pythonhosted.org/packages/e8/d3/ddfd9878b223b3aa9a930c6100a99afca5cfab7ea703662e00323acb7568/ml_dtypes-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:274cc7193dd73b35fb26bef6c5d40ae3eb258359ee71cd82f6e96a8c948bdaa6", size = 126727, upload-time = "2024-09-13T19:06:55.897Z" }, + { url = "https://files.pythonhosted.org/packages/ba/1a/99e924f12e4b62139fbac87419698c65f956d58de0dbfa7c028fa5b096aa/ml_dtypes-0.4.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:827d3ca2097085cf0355f8fdf092b888890bb1b1455f52801a2d7756f056f54b", size = 405077, upload-time = "2024-09-13T19:06:57.538Z" }, + { url = "https://files.pythonhosted.org/packages/8f/8c/7b610bd500617854c8cc6ed7c8cfb9d48d6a5c21a1437a36a4b9bc8a3598/ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:772426b08a6172a891274d581ce58ea2789cc8abc1c002a27223f314aaf894e7", size = 2181554, upload-time = "2024-09-13T19:06:59.196Z" }, + { url = "https://files.pythonhosted.org/packages/c7/c6/f89620cecc0581dc1839e218c4315171312e46c62a62da6ace204bda91c0/ml_dtypes-0.4.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:126e7d679b8676d1a958f2651949fbfa182832c3cd08020d8facd94e4114f3e9", size = 2160488, upload-time = "2024-09-13T19:07:03.131Z" }, + { url = "https://files.pythonhosted.org/packages/ae/11/a742d3c31b2cc8557a48efdde53427fd5f9caa2fa3c9c27d826e78a66f51/ml_dtypes-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:df0fb650d5c582a9e72bb5bd96cfebb2cdb889d89daff621c8fbc60295eba66c", size = 127462, upload-time = "2024-09-13T19:07:04.916Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fe/3a/c5b855752a70267ff729c349e650263adb3c206c29d28cc8ea7ace30a1d5/ml_dtypes-0.5.4-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b95e97e470fe60ed493fd9ae3911d8da4ebac16bd21f87ffa2b7c588bf22ea2c", size = 679735, upload-time = "2025-11-17T22:31:31.367Z" }, + { url = "https://files.pythonhosted.org/packages/41/79/7433f30ee04bd4faa303844048f55e1eb939131c8e5195a00a96a0939b64/ml_dtypes-0.5.4-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b4b801ebe0b477be666696bda493a9be8356f1f0057a57f1e35cd26928823e5a", size = 5051883, upload-time = "2025-11-17T22:31:33.658Z" }, + { url = "https://files.pythonhosted.org/packages/10/b1/8938e8830b0ee2e167fc75a094dea766a1152bde46752cd9bfc57ee78a82/ml_dtypes-0.5.4-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:388d399a2152dd79a3f0456a952284a99ee5c93d3e2f8dfe25977511e0515270", size = 5030369, upload-time = "2025-11-17T22:31:35.595Z" }, + { url = "https://files.pythonhosted.org/packages/c7/a3/51886727bd16e2f47587997b802dd56398692ce8c6c03c2e5bb32ecafe26/ml_dtypes-0.5.4-cp310-cp310-win_amd64.whl", hash = "sha256:4ff7f3e7ca2972e7de850e7b8fcbb355304271e2933dd90814c1cb847414d6e2", size = 210738, upload-time = "2025-11-17T22:31:37.43Z" }, + { url = "https://files.pythonhosted.org/packages/c6/5e/712092cfe7e5eb667b8ad9ca7c54442f21ed7ca8979745f1000e24cf8737/ml_dtypes-0.5.4-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6c7ecb74c4bd71db68a6bea1edf8da8c34f3d9fe218f038814fd1d310ac76c90", size = 679734, upload-time = "2025-11-17T22:31:39.223Z" }, + { url = "https://files.pythonhosted.org/packages/4f/cf/912146dfd4b5c0eea956836c01dcd2fce6c9c844b2691f5152aca196ce4f/ml_dtypes-0.5.4-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bc11d7e8c44a65115d05e2ab9989d1e045125d7be8e05a071a48bc76eb6d6040", size = 5056165, upload-time = "2025-11-17T22:31:41.071Z" }, + { url = "https://files.pythonhosted.org/packages/a9/80/19189ea605017473660e43762dc853d2797984b3c7bf30ce656099add30c/ml_dtypes-0.5.4-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:19b9a53598f21e453ea2fbda8aa783c20faff8e1eeb0d7ab899309a0053f1483", size = 5034975, upload-time = "2025-11-17T22:31:42.758Z" }, + { url = "https://files.pythonhosted.org/packages/b4/24/70bd59276883fdd91600ca20040b41efd4902a923283c4d6edcb1de128d2/ml_dtypes-0.5.4-cp311-cp311-win_amd64.whl", hash = "sha256:7c23c54a00ae43edf48d44066a7ec31e05fdc2eee0be2b8b50dd1903a1db94bb", size = 210742, upload-time = "2025-11-17T22:31:44.068Z" }, + { url = "https://files.pythonhosted.org/packages/a0/c9/64230ef14e40aa3f1cb254ef623bf812735e6bec7772848d19131111ac0d/ml_dtypes-0.5.4-cp311-cp311-win_arm64.whl", hash = "sha256:557a31a390b7e9439056644cb80ed0735a6e3e3bb09d67fd5687e4b04238d1de", size = 160709, upload-time = "2025-11-17T22:31:46.557Z" }, + { url = "https://files.pythonhosted.org/packages/a8/b8/3c70881695e056f8a32f8b941126cf78775d9a4d7feba8abcb52cb7b04f2/ml_dtypes-0.5.4-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:a174837a64f5b16cab6f368171a1a03a27936b31699d167684073ff1c4237dac", size = 676927, upload-time = "2025-11-17T22:31:48.182Z" }, + { url = "https://files.pythonhosted.org/packages/54/0f/428ef6881782e5ebb7eca459689448c0394fa0a80bea3aa9262cba5445ea/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7f7c643e8b1320fd958bf098aa7ecf70623a42ec5154e3be3be673f4c34d900", size = 5028464, upload-time = "2025-11-17T22:31:50.135Z" }, + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, + { url = "https://files.pythonhosted.org/packages/f5/f0/0cfadd537c5470378b1b32bd859cf2824972174b51b873c9d95cfd7475a5/ml_dtypes-0.5.4-cp312-cp312-win_amd64.whl", hash = "sha256:c1a953995cccb9e25a4ae19e34316671e4e2edaebe4cf538229b1fc7109087b7", size = 212222, upload-time = "2025-11-17T22:31:53.742Z" }, + { url = "https://files.pythonhosted.org/packages/16/2e/9acc86985bfad8f2c2d30291b27cd2bb4c74cea08695bd540906ed744249/ml_dtypes-0.5.4-cp312-cp312-win_arm64.whl", hash = "sha256:9bad06436568442575beb2d03389aa7456c690a5b05892c471215bfd8cf39460", size = 160793, upload-time = "2025-11-17T22:31:55.358Z" }, + { url = "https://files.pythonhosted.org/packages/d9/a1/4008f14bbc616cfb1ac5b39ea485f9c63031c4634ab3f4cf72e7541f816a/ml_dtypes-0.5.4-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:8c760d85a2f82e2bed75867079188c9d18dae2ee77c25a54d60e9cc79be1bc48", size = 676888, upload-time = "2025-11-17T22:31:56.907Z" }, + { url = "https://files.pythonhosted.org/packages/d3/b7/dff378afc2b0d5a7d6cd9d3209b60474d9819d1189d347521e1688a60a53/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ce756d3a10d0c4067172804c9cc276ba9cc0ff47af9078ad439b075d1abdc29b", size = 5036993, upload-time = "2025-11-17T22:31:58.497Z" }, + { url = "https://files.pythonhosted.org/packages/eb/33/40cd74219417e78b97c47802037cf2d87b91973e18bb968a7da48a96ea44/ml_dtypes-0.5.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:533ce891ba774eabf607172254f2e7260ba5f57bdd64030c9a4fcfbd99815d0d", size = 5010956, upload-time = "2025-11-17T22:31:59.931Z" }, + { url = "https://files.pythonhosted.org/packages/e1/8b/200088c6859d8221454825959df35b5244fa9bdf263fd0249ac5fb75e281/ml_dtypes-0.5.4-cp313-cp313-win_amd64.whl", hash = "sha256:f21c9219ef48ca5ee78402d5cc831bd58ea27ce89beda894428bc67a52da5328", size = 212224, upload-time = "2025-11-17T22:32:01.349Z" }, + { url = "https://files.pythonhosted.org/packages/8f/75/dfc3775cb36367816e678f69a7843f6f03bd4e2bcd79941e01ea960a068e/ml_dtypes-0.5.4-cp313-cp313-win_arm64.whl", hash = "sha256:35f29491a3e478407f7047b8a4834e4640a77d2737e0b294d049746507af5175", size = 160798, upload-time = "2025-11-17T22:32:02.864Z" }, + { url = "https://files.pythonhosted.org/packages/4f/74/e9ddb35fd1dd43b1106c20ced3f53c2e8e7fc7598c15638e9f80677f81d4/ml_dtypes-0.5.4-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:304ad47faa395415b9ccbcc06a0350800bc50eda70f0e45326796e27c62f18b6", size = 702083, upload-time = "2025-11-17T22:32:04.08Z" }, + { url = "https://files.pythonhosted.org/packages/74/f5/667060b0aed1aa63166b22897fdf16dca9eb704e6b4bbf86848d5a181aa7/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6a0df4223b514d799b8a1629c65ddc351b3efa833ccf7f8ea0cf654a61d1e35d", size = 5354111, upload-time = "2025-11-17T22:32:05.546Z" }, + { url = "https://files.pythonhosted.org/packages/40/49/0f8c498a28c0efa5f5c95a9e374c83ec1385ca41d0e85e7cf40e5d519a21/ml_dtypes-0.5.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:531eff30e4d368cb6255bc2328d070e35836aa4f282a0fb5f3a0cd7260257298", size = 5366453, upload-time = "2025-11-17T22:32:07.115Z" }, + { url = "https://files.pythonhosted.org/packages/8c/27/12607423d0a9c6bbbcc780ad19f1f6baa2b68b18ce4bddcdc122c4c68dc9/ml_dtypes-0.5.4-cp313-cp313t-win_amd64.whl", hash = "sha256:cb73dccfc991691c444acc8c0012bee8f2470da826a92e3a20bb333b1a7894e6", size = 225612, upload-time = "2025-11-17T22:32:08.615Z" }, + { url = "https://files.pythonhosted.org/packages/e5/80/5a5929e92c72936d5b19872c5fb8fc09327c1da67b3b68c6a13139e77e20/ml_dtypes-0.5.4-cp313-cp313t-win_arm64.whl", hash = "sha256:3bbbe120b915090d9dd1375e4684dd17a20a2491ef25d640a908281da85e73f1", size = 164145, upload-time = "2025-11-17T22:32:09.782Z" }, + { url = "https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22", size = 673781, upload-time = "2025-11-17T22:32:11.364Z" }, + { url = "https://files.pythonhosted.org/packages/04/f9/067b84365c7e83bda15bba2b06c6ca250ce27b20630b1128c435fb7a09aa/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465", size = 5036145, upload-time = "2025-11-17T22:32:12.783Z" }, + { url = "https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f", size = 5010230, upload-time = "2025-11-17T22:32:14.38Z" }, + { url = "https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56", size = 221032, upload-time = "2025-11-17T22:32:15.763Z" }, + { url = "https://files.pythonhosted.org/packages/76/a3/9c912fe6ea747bb10fe2f8f54d027eb265db05dfb0c6335e3e063e74e6e8/ml_dtypes-0.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049", size = 163353, upload-time = "2025-11-17T22:32:16.932Z" }, + { url = "https://files.pythonhosted.org/packages/cd/02/48aa7d84cc30ab4ee37624a2fd98c56c02326785750cd212bc0826c2f15b/ml_dtypes-0.5.4-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9", size = 702085, upload-time = "2025-11-17T22:32:18.175Z" }, + { url = "https://files.pythonhosted.org/packages/5a/e7/85cb99fe80a7a5513253ec7faa88a65306be071163485e9a626fce1b6e84/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7", size = 5355358, upload-time = "2025-11-17T22:32:19.7Z" }, + { url = "https://files.pythonhosted.org/packages/79/2b/a826ba18d2179a56e144aef69e57fb2ab7c464ef0b2111940ee8a3a223a2/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf", size = 5366332, upload-time = "2025-11-17T22:32:21.193Z" }, + { url = "https://files.pythonhosted.org/packages/84/44/f4d18446eacb20ea11e82f133ea8f86e2bf2891785b67d9da8d0ab0ef525/ml_dtypes-0.5.4-cp314-cp314t-win_amd64.whl", hash = "sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1", size = 236612, upload-time = "2025-11-17T22:32:22.579Z" }, + { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" }, +] + [[package]] name = "more-itertools" version = "10.8.0" @@ -1622,9 +2380,11 @@ name = "networkx" version = "3.4.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.11' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", "python_full_version < '3.11' and sys_platform == 'win32'", - "python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/fd/1d/06475e1cd5264c0b870ea2cc6fdb3e37177c1e565c43f56ff17a10e3937f/networkx-3.4.2.tar.gz", hash = "sha256:307c3669428c5362aab27c8a1260aa8f47c4e91d3891f48be0141738d8d053e1", size = 2151368, upload-time = "2024-10-21T12:39:38.695Z" } wheels = [ @@ -1636,12 +2396,21 @@ name = "networkx" version = "3.5" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", - "python_full_version >= '3.12' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.11.*' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/6c/4f/ccdb8ad3a38e583f214547fd2f7ff1fc160c43a75af88e6aec213404b96a/networkx-3.5.tar.gz", hash = "sha256:d4c6f9cf81f52d69230866796b82afbccdec3db7ae4fbd1b65ea750feed50037", size = 2471065, upload-time = "2025-05-29T11:35:07.804Z" } wheels = [ @@ -1692,158 +2461,34 @@ wheels = [ [[package]] name = "numpy" -version = "2.2.6" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version < '3.11' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'win32'", - "python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32'", -] -sdist = { url = "https://files.pythonhosted.org/packages/76/21/7d2a95e4bba9dc13d043ee156a356c0a8f0c6309dff6b21b4d71a073b8a8/numpy-2.2.6.tar.gz", hash = "sha256:e29554e2bef54a90aa5cc07da6ce955accb83f21ab5de01a62c8478897b264fd", size = 20276440, upload-time = "2025-05-17T22:38:04.611Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9a/3e/ed6db5be21ce87955c0cbd3009f2803f59fa08df21b5df06862e2d8e2bdd/numpy-2.2.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b412caa66f72040e6d268491a59f2c43bf03eb6c96dd8f0307829feb7fa2b6fb", size = 21165245, upload-time = "2025-05-17T21:27:58.555Z" }, - { url = "https://files.pythonhosted.org/packages/22/c2/4b9221495b2a132cc9d2eb862e21d42a009f5a60e45fc44b00118c174bff/numpy-2.2.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e41fd67c52b86603a91c1a505ebaef50b3314de0213461c7a6e99c9a3beff90", size = 14360048, upload-time = "2025-05-17T21:28:21.406Z" }, - { url = "https://files.pythonhosted.org/packages/fd/77/dc2fcfc66943c6410e2bf598062f5959372735ffda175b39906d54f02349/numpy-2.2.6-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:37e990a01ae6ec7fe7fa1c26c55ecb672dd98b19c3d0e1d1f326fa13cb38d163", size = 5340542, upload-time = "2025-05-17T21:28:30.931Z" }, - { url = "https://files.pythonhosted.org/packages/7a/4f/1cb5fdc353a5f5cc7feb692db9b8ec2c3d6405453f982435efc52561df58/numpy-2.2.6-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:5a6429d4be8ca66d889b7cf70f536a397dc45ba6faeb5f8c5427935d9592e9cf", size = 6878301, upload-time = "2025-05-17T21:28:41.613Z" }, - { url = "https://files.pythonhosted.org/packages/eb/17/96a3acd228cec142fcb8723bd3cc39c2a474f7dcf0a5d16731980bcafa95/numpy-2.2.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:efd28d4e9cd7d7a8d39074a4d44c63eda73401580c5c76acda2ce969e0a38e83", size = 14297320, upload-time = "2025-05-17T21:29:02.78Z" }, - { url = "https://files.pythonhosted.org/packages/b4/63/3de6a34ad7ad6646ac7d2f55ebc6ad439dbbf9c4370017c50cf403fb19b5/numpy-2.2.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc7b73d02efb0e18c000e9ad8b83480dfcd5dfd11065997ed4c6747470ae8915", size = 16801050, upload-time = "2025-05-17T21:29:27.675Z" }, - { url = "https://files.pythonhosted.org/packages/07/b6/89d837eddef52b3d0cec5c6ba0456c1bf1b9ef6a6672fc2b7873c3ec4e2e/numpy-2.2.6-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:74d4531beb257d2c3f4b261bfb0fc09e0f9ebb8842d82a7b4209415896adc680", size = 15807034, upload-time = "2025-05-17T21:29:51.102Z" }, - { url = "https://files.pythonhosted.org/packages/01/c8/dc6ae86e3c61cfec1f178e5c9f7858584049b6093f843bca541f94120920/numpy-2.2.6-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8fc377d995680230e83241d8a96def29f204b5782f371c532579b4f20607a289", size = 18614185, upload-time = "2025-05-17T21:30:18.703Z" }, - { url = "https://files.pythonhosted.org/packages/5b/c5/0064b1b7e7c89137b471ccec1fd2282fceaae0ab3a9550f2568782d80357/numpy-2.2.6-cp310-cp310-win32.whl", hash = "sha256:b093dd74e50a8cba3e873868d9e93a85b78e0daf2e98c6797566ad8044e8363d", size = 6527149, upload-time = "2025-05-17T21:30:29.788Z" }, - { url = "https://files.pythonhosted.org/packages/a3/dd/4b822569d6b96c39d1215dbae0582fd99954dcbcf0c1a13c61783feaca3f/numpy-2.2.6-cp310-cp310-win_amd64.whl", hash = "sha256:f0fd6321b839904e15c46e0d257fdd101dd7f530fe03fd6359c1ea63738703f3", size = 12904620, upload-time = "2025-05-17T21:30:48.994Z" }, - { url = "https://files.pythonhosted.org/packages/da/a8/4f83e2aa666a9fbf56d6118faaaf5f1974d456b1823fda0a176eff722839/numpy-2.2.6-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f9f1adb22318e121c5c69a09142811a201ef17ab257a1e66ca3025065b7f53ae", size = 21176963, upload-time = "2025-05-17T21:31:19.36Z" }, - { url = "https://files.pythonhosted.org/packages/b3/2b/64e1affc7972decb74c9e29e5649fac940514910960ba25cd9af4488b66c/numpy-2.2.6-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c820a93b0255bc360f53eca31a0e676fd1101f673dda8da93454a12e23fc5f7a", size = 14406743, upload-time = "2025-05-17T21:31:41.087Z" }, - { url = "https://files.pythonhosted.org/packages/4a/9f/0121e375000b5e50ffdd8b25bf78d8e1a5aa4cca3f185d41265198c7b834/numpy-2.2.6-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3d70692235e759f260c3d837193090014aebdf026dfd167834bcba43e30c2a42", size = 5352616, upload-time = "2025-05-17T21:31:50.072Z" }, - { url = "https://files.pythonhosted.org/packages/31/0d/b48c405c91693635fbe2dcd7bc84a33a602add5f63286e024d3b6741411c/numpy-2.2.6-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:481b49095335f8eed42e39e8041327c05b0f6f4780488f61286ed3c01368d491", size = 6889579, upload-time = "2025-05-17T21:32:01.712Z" }, - { url = "https://files.pythonhosted.org/packages/52/b8/7f0554d49b565d0171eab6e99001846882000883998e7b7d9f0d98b1f934/numpy-2.2.6-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b64d8d4d17135e00c8e346e0a738deb17e754230d7e0810ac5012750bbd85a5a", size = 14312005, upload-time = "2025-05-17T21:32:23.332Z" }, - { url = "https://files.pythonhosted.org/packages/b3/dd/2238b898e51bd6d389b7389ffb20d7f4c10066d80351187ec8e303a5a475/numpy-2.2.6-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba10f8411898fc418a521833e014a77d3ca01c15b0c6cdcce6a0d2897e6dbbdf", size = 16821570, upload-time = "2025-05-17T21:32:47.991Z" }, - { url = "https://files.pythonhosted.org/packages/83/6c/44d0325722cf644f191042bf47eedad61c1e6df2432ed65cbe28509d404e/numpy-2.2.6-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:bd48227a919f1bafbdda0583705e547892342c26fb127219d60a5c36882609d1", size = 15818548, upload-time = "2025-05-17T21:33:11.728Z" }, - { url = "https://files.pythonhosted.org/packages/ae/9d/81e8216030ce66be25279098789b665d49ff19eef08bfa8cb96d4957f422/numpy-2.2.6-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:9551a499bf125c1d4f9e250377c1ee2eddd02e01eac6644c080162c0c51778ab", size = 18620521, upload-time = "2025-05-17T21:33:39.139Z" }, - { url = "https://files.pythonhosted.org/packages/6a/fd/e19617b9530b031db51b0926eed5345ce8ddc669bb3bc0044b23e275ebe8/numpy-2.2.6-cp311-cp311-win32.whl", hash = "sha256:0678000bb9ac1475cd454c6b8c799206af8107e310843532b04d49649c717a47", size = 6525866, upload-time = "2025-05-17T21:33:50.273Z" }, - { url = "https://files.pythonhosted.org/packages/31/0a/f354fb7176b81747d870f7991dc763e157a934c717b67b58456bc63da3df/numpy-2.2.6-cp311-cp311-win_amd64.whl", hash = "sha256:e8213002e427c69c45a52bbd94163084025f533a55a59d6f9c5b820774ef3303", size = 12907455, upload-time = "2025-05-17T21:34:09.135Z" }, - { url = "https://files.pythonhosted.org/packages/82/5d/c00588b6cf18e1da539b45d3598d3557084990dcc4331960c15ee776ee41/numpy-2.2.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:41c5a21f4a04fa86436124d388f6ed60a9343a6f767fced1a8a71c3fbca038ff", size = 20875348, upload-time = "2025-05-17T21:34:39.648Z" }, - { url = "https://files.pythonhosted.org/packages/66/ee/560deadcdde6c2f90200450d5938f63a34b37e27ebff162810f716f6a230/numpy-2.2.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:de749064336d37e340f640b05f24e9e3dd678c57318c7289d222a8a2f543e90c", size = 14119362, upload-time = "2025-05-17T21:35:01.241Z" }, - { url = "https://files.pythonhosted.org/packages/3c/65/4baa99f1c53b30adf0acd9a5519078871ddde8d2339dc5a7fde80d9d87da/numpy-2.2.6-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:894b3a42502226a1cac872f840030665f33326fc3dac8e57c607905773cdcde3", size = 5084103, upload-time = "2025-05-17T21:35:10.622Z" }, - { url = "https://files.pythonhosted.org/packages/cc/89/e5a34c071a0570cc40c9a54eb472d113eea6d002e9ae12bb3a8407fb912e/numpy-2.2.6-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:71594f7c51a18e728451bb50cc60a3ce4e6538822731b2933209a1f3614e9282", size = 6625382, upload-time = "2025-05-17T21:35:21.414Z" }, - { url = "https://files.pythonhosted.org/packages/f8/35/8c80729f1ff76b3921d5c9487c7ac3de9b2a103b1cd05e905b3090513510/numpy-2.2.6-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2618db89be1b4e05f7a1a847a9c1c0abd63e63a1607d892dd54668dd92faf87", size = 14018462, upload-time = "2025-05-17T21:35:42.174Z" }, - { url = "https://files.pythonhosted.org/packages/8c/3d/1e1db36cfd41f895d266b103df00ca5b3cbe965184df824dec5c08c6b803/numpy-2.2.6-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd83c01228a688733f1ded5201c678f0c53ecc1006ffbc404db9f7a899ac6249", size = 16527618, upload-time = "2025-05-17T21:36:06.711Z" }, - { url = "https://files.pythonhosted.org/packages/61/c6/03ed30992602c85aa3cd95b9070a514f8b3c33e31124694438d88809ae36/numpy-2.2.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:37c0ca431f82cd5fa716eca9506aefcabc247fb27ba69c5062a6d3ade8cf8f49", size = 15505511, upload-time = "2025-05-17T21:36:29.965Z" }, - { url = "https://files.pythonhosted.org/packages/b7/25/5761d832a81df431e260719ec45de696414266613c9ee268394dd5ad8236/numpy-2.2.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fe27749d33bb772c80dcd84ae7e8df2adc920ae8297400dabec45f0dedb3f6de", size = 18313783, upload-time = "2025-05-17T21:36:56.883Z" }, - { url = "https://files.pythonhosted.org/packages/57/0a/72d5a3527c5ebffcd47bde9162c39fae1f90138c961e5296491ce778e682/numpy-2.2.6-cp312-cp312-win32.whl", hash = "sha256:4eeaae00d789f66c7a25ac5f34b71a7035bb474e679f410e5e1a94deb24cf2d4", size = 6246506, upload-time = "2025-05-17T21:37:07.368Z" }, - { url = "https://files.pythonhosted.org/packages/36/fa/8c9210162ca1b88529ab76b41ba02d433fd54fecaf6feb70ef9f124683f1/numpy-2.2.6-cp312-cp312-win_amd64.whl", hash = "sha256:c1f9540be57940698ed329904db803cf7a402f3fc200bfe599334c9bd84a40b2", size = 12614190, upload-time = "2025-05-17T21:37:26.213Z" }, - { url = "https://files.pythonhosted.org/packages/f9/5c/6657823f4f594f72b5471f1db1ab12e26e890bb2e41897522d134d2a3e81/numpy-2.2.6-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0811bb762109d9708cca4d0b13c4f67146e3c3b7cf8d34018c722adb2d957c84", size = 20867828, upload-time = "2025-05-17T21:37:56.699Z" }, - { url = "https://files.pythonhosted.org/packages/dc/9e/14520dc3dadf3c803473bd07e9b2bd1b69bc583cb2497b47000fed2fa92f/numpy-2.2.6-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:287cc3162b6f01463ccd86be154f284d0893d2b3ed7292439ea97eafa8170e0b", size = 14143006, upload-time = "2025-05-17T21:38:18.291Z" }, - { url = "https://files.pythonhosted.org/packages/4f/06/7e96c57d90bebdce9918412087fc22ca9851cceaf5567a45c1f404480e9e/numpy-2.2.6-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:f1372f041402e37e5e633e586f62aa53de2eac8d98cbfb822806ce4bbefcb74d", size = 5076765, upload-time = "2025-05-17T21:38:27.319Z" }, - { url = "https://files.pythonhosted.org/packages/73/ed/63d920c23b4289fdac96ddbdd6132e9427790977d5457cd132f18e76eae0/numpy-2.2.6-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:55a4d33fa519660d69614a9fad433be87e5252f4b03850642f88993f7b2ca566", size = 6617736, upload-time = "2025-05-17T21:38:38.141Z" }, - { url = "https://files.pythonhosted.org/packages/85/c5/e19c8f99d83fd377ec8c7e0cf627a8049746da54afc24ef0a0cb73d5dfb5/numpy-2.2.6-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92729c95468a2f4f15e9bb94c432a9229d0d50de67304399627a943201baa2f", size = 14010719, upload-time = "2025-05-17T21:38:58.433Z" }, - { url = "https://files.pythonhosted.org/packages/19/49/4df9123aafa7b539317bf6d342cb6d227e49f7a35b99c287a6109b13dd93/numpy-2.2.6-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1bc23a79bfabc5d056d106f9befb8d50c31ced2fbc70eedb8155aec74a45798f", size = 16526072, upload-time = "2025-05-17T21:39:22.638Z" }, - { url = "https://files.pythonhosted.org/packages/b2/6c/04b5f47f4f32f7c2b0e7260442a8cbcf8168b0e1a41ff1495da42f42a14f/numpy-2.2.6-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e3143e4451880bed956e706a3220b4e5cf6172ef05fcc397f6f36a550b1dd868", size = 15503213, upload-time = "2025-05-17T21:39:45.865Z" }, - { url = "https://files.pythonhosted.org/packages/17/0a/5cd92e352c1307640d5b6fec1b2ffb06cd0dabe7d7b8227f97933d378422/numpy-2.2.6-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:b4f13750ce79751586ae2eb824ba7e1e8dba64784086c98cdbbcc6a42112ce0d", size = 18316632, upload-time = "2025-05-17T21:40:13.331Z" }, - { url = "https://files.pythonhosted.org/packages/f0/3b/5cba2b1d88760ef86596ad0f3d484b1cbff7c115ae2429678465057c5155/numpy-2.2.6-cp313-cp313-win32.whl", hash = "sha256:5beb72339d9d4fa36522fc63802f469b13cdbe4fdab4a288f0c441b74272ebfd", size = 6244532, upload-time = "2025-05-17T21:43:46.099Z" }, - { url = "https://files.pythonhosted.org/packages/cb/3b/d58c12eafcb298d4e6d0d40216866ab15f59e55d148a5658bb3132311fcf/numpy-2.2.6-cp313-cp313-win_amd64.whl", hash = "sha256:b0544343a702fa80c95ad5d3d608ea3599dd54d4632df855e4c8d24eb6ecfa1c", size = 12610885, upload-time = "2025-05-17T21:44:05.145Z" }, - { url = "https://files.pythonhosted.org/packages/6b/9e/4bf918b818e516322db999ac25d00c75788ddfd2d2ade4fa66f1f38097e1/numpy-2.2.6-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0bca768cd85ae743b2affdc762d617eddf3bcf8724435498a1e80132d04879e6", size = 20963467, upload-time = "2025-05-17T21:40:44Z" }, - { url = "https://files.pythonhosted.org/packages/61/66/d2de6b291507517ff2e438e13ff7b1e2cdbdb7cb40b3ed475377aece69f9/numpy-2.2.6-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:fc0c5673685c508a142ca65209b4e79ed6740a4ed6b2267dbba90f34b0b3cfda", size = 14225144, upload-time = "2025-05-17T21:41:05.695Z" }, - { url = "https://files.pythonhosted.org/packages/e4/25/480387655407ead912e28ba3a820bc69af9adf13bcbe40b299d454ec011f/numpy-2.2.6-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:5bd4fc3ac8926b3819797a7c0e2631eb889b4118a9898c84f585a54d475b7e40", size = 5200217, upload-time = "2025-05-17T21:41:15.903Z" }, - { url = "https://files.pythonhosted.org/packages/aa/4a/6e313b5108f53dcbf3aca0c0f3e9c92f4c10ce57a0a721851f9785872895/numpy-2.2.6-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:fee4236c876c4e8369388054d02d0e9bb84821feb1a64dd59e137e6511a551f8", size = 6712014, upload-time = "2025-05-17T21:41:27.321Z" }, - { url = "https://files.pythonhosted.org/packages/b7/30/172c2d5c4be71fdf476e9de553443cf8e25feddbe185e0bd88b096915bcc/numpy-2.2.6-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e1dda9c7e08dc141e0247a5b8f49cf05984955246a327d4c48bda16821947b2f", size = 14077935, upload-time = "2025-05-17T21:41:49.738Z" }, - { url = "https://files.pythonhosted.org/packages/12/fb/9e743f8d4e4d3c710902cf87af3512082ae3d43b945d5d16563f26ec251d/numpy-2.2.6-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f447e6acb680fd307f40d3da4852208af94afdfab89cf850986c3ca00562f4fa", size = 16600122, upload-time = "2025-05-17T21:42:14.046Z" }, - { url = "https://files.pythonhosted.org/packages/12/75/ee20da0e58d3a66f204f38916757e01e33a9737d0b22373b3eb5a27358f9/numpy-2.2.6-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:389d771b1623ec92636b0786bc4ae56abafad4a4c513d36a55dce14bd9ce8571", size = 15586143, upload-time = "2025-05-17T21:42:37.464Z" }, - { url = "https://files.pythonhosted.org/packages/76/95/bef5b37f29fc5e739947e9ce5179ad402875633308504a52d188302319c8/numpy-2.2.6-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8e9ace4a37db23421249ed236fdcdd457d671e25146786dfc96835cd951aa7c1", size = 18385260, upload-time = "2025-05-17T21:43:05.189Z" }, - { url = "https://files.pythonhosted.org/packages/09/04/f2f83279d287407cf36a7a8053a5abe7be3622a4363337338f2585e4afda/numpy-2.2.6-cp313-cp313t-win32.whl", hash = "sha256:038613e9fb8c72b0a41f025a7e4c3f0b7a1b5d768ece4796b674c8f3fe13efff", size = 6377225, upload-time = "2025-05-17T21:43:16.254Z" }, - { url = "https://files.pythonhosted.org/packages/67/0e/35082d13c09c02c011cf21570543d202ad929d961c02a147493cb0c2bdf5/numpy-2.2.6-cp313-cp313t-win_amd64.whl", hash = "sha256:6031dd6dfecc0cf9f668681a37648373bddd6421fff6c66ec1624eed0180ee06", size = 12771374, upload-time = "2025-05-17T21:43:35.479Z" }, - { url = "https://files.pythonhosted.org/packages/9e/3b/d94a75f4dbf1ef5d321523ecac21ef23a3cd2ac8b78ae2aac40873590229/numpy-2.2.6-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:0b605b275d7bd0c640cad4e5d30fa701a8d59302e127e5f79138ad62762c3e3d", size = 21040391, upload-time = "2025-05-17T21:44:35.948Z" }, - { url = "https://files.pythonhosted.org/packages/17/f4/09b2fa1b58f0fb4f7c7963a1649c64c4d315752240377ed74d9cd878f7b5/numpy-2.2.6-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:7befc596a7dc9da8a337f79802ee8adb30a552a94f792b9c9d18c840055907db", size = 6786754, upload-time = "2025-05-17T21:44:47.446Z" }, - { url = "https://files.pythonhosted.org/packages/af/30/feba75f143bdc868a1cc3f44ccfa6c4b9ec522b36458e738cd00f67b573f/numpy-2.2.6-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ce47521a4754c8f4593837384bd3424880629f718d87c5d44f8ed763edd63543", size = 16643476, upload-time = "2025-05-17T21:45:11.871Z" }, - { url = "https://files.pythonhosted.org/packages/37/48/ac2a9584402fb6c0cd5b5d1a91dcf176b15760130dd386bbafdbfe3640bf/numpy-2.2.6-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d042d24c90c41b54fd506da306759e06e568864df8ec17ccc17e9e884634fd00", size = 12812666, upload-time = "2025-05-17T21:45:31.426Z" }, -] - -[[package]] -name = "numpy" -version = "2.3.4" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", - "python_full_version >= '3.12' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.11.*' and sys_platform != 'linux' and sys_platform != 'win32'", -] -sdist = { url = "https://files.pythonhosted.org/packages/b5/f4/098d2270d52b41f1bd7db9fc288aaa0400cb48c2a3e2af6fa365d9720947/numpy-2.3.4.tar.gz", hash = "sha256:a7d018bfedb375a8d979ac758b120ba846a7fe764911a64465fd87b8729f4a6a", size = 20582187, upload-time = "2025-10-15T16:18:11.77Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/60/e7/0e07379944aa8afb49a556a2b54587b828eb41dc9adc56fb7615b678ca53/numpy-2.3.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:e78aecd2800b32e8347ce49316d3eaf04aed849cd5b38e0af39f829a4e59f5eb", size = 21259519, upload-time = "2025-10-15T16:15:19.012Z" }, - { url = "https://files.pythonhosted.org/packages/d0/cb/5a69293561e8819b09e34ed9e873b9a82b5f2ade23dce4c51dc507f6cfe1/numpy-2.3.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7fd09cc5d65bda1e79432859c40978010622112e9194e581e3415a3eccc7f43f", size = 14452796, upload-time = "2025-10-15T16:15:23.094Z" }, - { url = "https://files.pythonhosted.org/packages/e4/04/ff11611200acd602a1e5129e36cfd25bf01ad8e5cf927baf2e90236eb02e/numpy-2.3.4-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:1b219560ae2c1de48ead517d085bc2d05b9433f8e49d0955c82e8cd37bd7bf36", size = 5381639, upload-time = "2025-10-15T16:15:25.572Z" }, - { url = "https://files.pythonhosted.org/packages/ea/77/e95c757a6fe7a48d28a009267408e8aa382630cc1ad1db7451b3bc21dbb4/numpy-2.3.4-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:bafa7d87d4c99752d07815ed7a2c0964f8ab311eb8168f41b910bd01d15b6032", size = 6914296, upload-time = "2025-10-15T16:15:27.079Z" }, - { url = "https://files.pythonhosted.org/packages/a3/d2/137c7b6841c942124eae921279e5c41b1c34bab0e6fc60c7348e69afd165/numpy-2.3.4-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:36dc13af226aeab72b7abad501d370d606326a0029b9f435eacb3b8c94b8a8b7", size = 14591904, upload-time = "2025-10-15T16:15:29.044Z" }, - { url = "https://files.pythonhosted.org/packages/bb/32/67e3b0f07b0aba57a078c4ab777a9e8e6bc62f24fb53a2337f75f9691699/numpy-2.3.4-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a7b2f9a18b5ff9824a6af80de4f37f4ec3c2aab05ef08f51c77a093f5b89adda", size = 16939602, upload-time = "2025-10-15T16:15:31.106Z" }, - { url = "https://files.pythonhosted.org/packages/95/22/9639c30e32c93c4cee3ccdb4b09c2d0fbff4dcd06d36b357da06146530fb/numpy-2.3.4-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:9984bd645a8db6ca15d850ff996856d8762c51a2239225288f08f9050ca240a0", size = 16372661, upload-time = "2025-10-15T16:15:33.546Z" }, - { url = "https://files.pythonhosted.org/packages/12/e9/a685079529be2b0156ae0c11b13d6be647743095bb51d46589e95be88086/numpy-2.3.4-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:64c5825affc76942973a70acf438a8ab618dbd692b84cd5ec40a0a0509edc09a", size = 18884682, upload-time = "2025-10-15T16:15:36.105Z" }, - { url = "https://files.pythonhosted.org/packages/cf/85/f6f00d019b0cc741e64b4e00ce865a57b6bed945d1bbeb1ccadbc647959b/numpy-2.3.4-cp311-cp311-win32.whl", hash = "sha256:ed759bf7a70342f7817d88376eb7142fab9fef8320d6019ef87fae05a99874e1", size = 6570076, upload-time = "2025-10-15T16:15:38.225Z" }, - { url = "https://files.pythonhosted.org/packages/7d/10/f8850982021cb90e2ec31990291f9e830ce7d94eef432b15066e7cbe0bec/numpy-2.3.4-cp311-cp311-win_amd64.whl", hash = "sha256:faba246fb30ea2a526c2e9645f61612341de1a83fb1e0c5edf4ddda5a9c10996", size = 13089358, upload-time = "2025-10-15T16:15:40.404Z" }, - { url = "https://files.pythonhosted.org/packages/d1/ad/afdd8351385edf0b3445f9e24210a9c3971ef4de8fd85155462fc4321d79/numpy-2.3.4-cp311-cp311-win_arm64.whl", hash = "sha256:4c01835e718bcebe80394fd0ac66c07cbb90147ebbdad3dcecd3f25de2ae7e2c", size = 10462292, upload-time = "2025-10-15T16:15:42.896Z" }, - { url = "https://files.pythonhosted.org/packages/96/7a/02420400b736f84317e759291b8edaeee9dc921f72b045475a9cbdb26b17/numpy-2.3.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:ef1b5a3e808bc40827b5fa2c8196151a4c5abe110e1726949d7abddfe5c7ae11", size = 20957727, upload-time = "2025-10-15T16:15:44.9Z" }, - { url = "https://files.pythonhosted.org/packages/18/90/a014805d627aa5750f6f0e878172afb6454552da929144b3c07fcae1bb13/numpy-2.3.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2f91f496a87235c6aaf6d3f3d89b17dba64996abadccb289f48456cff931ca9", size = 14187262, upload-time = "2025-10-15T16:15:47.761Z" }, - { url = "https://files.pythonhosted.org/packages/c7/e4/0a94b09abe89e500dc748e7515f21a13e30c5c3fe3396e6d4ac108c25fca/numpy-2.3.4-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:f77e5b3d3da652b474cc80a14084927a5e86a5eccf54ca8ca5cbd697bf7f2667", size = 5115992, upload-time = "2025-10-15T16:15:50.144Z" }, - { url = "https://files.pythonhosted.org/packages/88/dd/db77c75b055c6157cbd4f9c92c4458daef0dd9cbe6d8d2fe7f803cb64c37/numpy-2.3.4-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:8ab1c5f5ee40d6e01cbe96de5863e39b215a4d24e7d007cad56c7184fdf4aeef", size = 6648672, upload-time = "2025-10-15T16:15:52.442Z" }, - { url = "https://files.pythonhosted.org/packages/e1/e6/e31b0d713719610e406c0ea3ae0d90760465b086da8783e2fd835ad59027/numpy-2.3.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:77b84453f3adcb994ddbd0d1c5d11db2d6bda1a2b7fd5ac5bd4649d6f5dc682e", size = 14284156, upload-time = "2025-10-15T16:15:54.351Z" }, - { url = "https://files.pythonhosted.org/packages/f9/58/30a85127bfee6f108282107caf8e06a1f0cc997cb6b52cdee699276fcce4/numpy-2.3.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4121c5beb58a7f9e6dfdee612cb24f4df5cd4db6e8261d7f4d7450a997a65d6a", size = 16641271, upload-time = "2025-10-15T16:15:56.67Z" }, - { url = "https://files.pythonhosted.org/packages/06/f2/2e06a0f2adf23e3ae29283ad96959267938d0efd20a2e25353b70065bfec/numpy-2.3.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:65611ecbb00ac9846efe04db15cbe6186f562f6bb7e5e05f077e53a599225d16", size = 16059531, upload-time = "2025-10-15T16:15:59.412Z" }, - { url = "https://files.pythonhosted.org/packages/b0/e7/b106253c7c0d5dc352b9c8fab91afd76a93950998167fa3e5afe4ef3a18f/numpy-2.3.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:dabc42f9c6577bcc13001b8810d300fe814b4cfbe8a92c873f269484594f9786", size = 18578983, upload-time = "2025-10-15T16:16:01.804Z" }, - { url = "https://files.pythonhosted.org/packages/73/e3/04ecc41e71462276ee867ccbef26a4448638eadecf1bc56772c9ed6d0255/numpy-2.3.4-cp312-cp312-win32.whl", hash = "sha256:a49d797192a8d950ca59ee2d0337a4d804f713bb5c3c50e8db26d49666e351dc", size = 6291380, upload-time = "2025-10-15T16:16:03.938Z" }, - { url = "https://files.pythonhosted.org/packages/3d/a8/566578b10d8d0e9955b1b6cd5db4e9d4592dd0026a941ff7994cedda030a/numpy-2.3.4-cp312-cp312-win_amd64.whl", hash = "sha256:985f1e46358f06c2a09921e8921e2c98168ed4ae12ccd6e5e87a4f1857923f32", size = 12787999, upload-time = "2025-10-15T16:16:05.801Z" }, - { url = "https://files.pythonhosted.org/packages/58/22/9c903a957d0a8071b607f5b1bff0761d6e608b9a965945411f867d515db1/numpy-2.3.4-cp312-cp312-win_arm64.whl", hash = "sha256:4635239814149e06e2cb9db3dd584b2fa64316c96f10656983b8026a82e6e4db", size = 10197412, upload-time = "2025-10-15T16:16:07.854Z" }, - { url = "https://files.pythonhosted.org/packages/57/7e/b72610cc91edf138bc588df5150957a4937221ca6058b825b4725c27be62/numpy-2.3.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c090d4860032b857d94144d1a9976b8e36709e40386db289aaf6672de2a81966", size = 20950335, upload-time = "2025-10-15T16:16:10.304Z" }, - { url = "https://files.pythonhosted.org/packages/3e/46/bdd3370dcea2f95ef14af79dbf81e6927102ddf1cc54adc0024d61252fd9/numpy-2.3.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a13fc473b6db0be619e45f11f9e81260f7302f8d180c49a22b6e6120022596b3", size = 14179878, upload-time = "2025-10-15T16:16:12.595Z" }, - { url = "https://files.pythonhosted.org/packages/ac/01/5a67cb785bda60f45415d09c2bc245433f1c68dd82eef9c9002c508b5a65/numpy-2.3.4-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:3634093d0b428e6c32c3a69b78e554f0cd20ee420dcad5a9f3b2a63762ce4197", size = 5108673, upload-time = "2025-10-15T16:16:14.877Z" }, - { url = "https://files.pythonhosted.org/packages/c2/cd/8428e23a9fcebd33988f4cb61208fda832800ca03781f471f3727a820704/numpy-2.3.4-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:043885b4f7e6e232d7df4f51ffdef8c36320ee9d5f227b380ea636722c7ed12e", size = 6641438, upload-time = "2025-10-15T16:16:16.805Z" }, - { url = "https://files.pythonhosted.org/packages/3e/d1/913fe563820f3c6b079f992458f7331278dcd7ba8427e8e745af37ddb44f/numpy-2.3.4-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4ee6a571d1e4f0ea6d5f22d6e5fbd6ed1dc2b18542848e1e7301bd190500c9d7", size = 14281290, upload-time = "2025-10-15T16:16:18.764Z" }, - { url = "https://files.pythonhosted.org/packages/9e/7e/7d306ff7cb143e6d975cfa7eb98a93e73495c4deabb7d1b5ecf09ea0fd69/numpy-2.3.4-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fc8a63918b04b8571789688b2780ab2b4a33ab44bfe8ccea36d3eba51228c953", size = 16636543, upload-time = "2025-10-15T16:16:21.072Z" }, - { url = "https://files.pythonhosted.org/packages/47/6a/8cfc486237e56ccfb0db234945552a557ca266f022d281a2f577b98e955c/numpy-2.3.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:40cc556d5abbc54aabe2b1ae287042d7bdb80c08edede19f0c0afb36ae586f37", size = 16056117, upload-time = "2025-10-15T16:16:23.369Z" }, - { url = "https://files.pythonhosted.org/packages/b1/0e/42cb5e69ea901e06ce24bfcc4b5664a56f950a70efdcf221f30d9615f3f3/numpy-2.3.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:ecb63014bb7f4ce653f8be7f1df8cbc6093a5a2811211770f6606cc92b5a78fd", size = 18577788, upload-time = "2025-10-15T16:16:27.496Z" }, - { url = "https://files.pythonhosted.org/packages/86/92/41c3d5157d3177559ef0a35da50f0cda7fa071f4ba2306dd36818591a5bc/numpy-2.3.4-cp313-cp313-win32.whl", hash = "sha256:e8370eb6925bb8c1c4264fec52b0384b44f675f191df91cbe0140ec9f0955646", size = 6282620, upload-time = "2025-10-15T16:16:29.811Z" }, - { url = "https://files.pythonhosted.org/packages/09/97/fd421e8bc50766665ad35536c2bb4ef916533ba1fdd053a62d96cc7c8b95/numpy-2.3.4-cp313-cp313-win_amd64.whl", hash = "sha256:56209416e81a7893036eea03abcb91c130643eb14233b2515c90dcac963fe99d", size = 12784672, upload-time = "2025-10-15T16:16:31.589Z" }, - { url = "https://files.pythonhosted.org/packages/ad/df/5474fb2f74970ca8eb978093969b125a84cc3d30e47f82191f981f13a8a0/numpy-2.3.4-cp313-cp313-win_arm64.whl", hash = "sha256:a700a4031bc0fd6936e78a752eefb79092cecad2599ea9c8039c548bc097f9bc", size = 10196702, upload-time = "2025-10-15T16:16:33.902Z" }, - { url = "https://files.pythonhosted.org/packages/11/83/66ac031464ec1767ea3ed48ce40f615eb441072945e98693bec0bcd056cc/numpy-2.3.4-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:86966db35c4040fdca64f0816a1c1dd8dbd027d90fca5a57e00e1ca4cd41b879", size = 21049003, upload-time = "2025-10-15T16:16:36.101Z" }, - { url = "https://files.pythonhosted.org/packages/5f/99/5b14e0e686e61371659a1d5bebd04596b1d72227ce36eed121bb0aeab798/numpy-2.3.4-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:838f045478638b26c375ee96ea89464d38428c69170360b23a1a50fa4baa3562", size = 14302980, upload-time = "2025-10-15T16:16:39.124Z" }, - { url = "https://files.pythonhosted.org/packages/2c/44/e9486649cd087d9fc6920e3fc3ac2aba10838d10804b1e179fb7cbc4e634/numpy-2.3.4-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:d7315ed1dab0286adca467377c8381cd748f3dc92235f22a7dfc42745644a96a", size = 5231472, upload-time = "2025-10-15T16:16:41.168Z" }, - { url = "https://files.pythonhosted.org/packages/3e/51/902b24fa8887e5fe2063fd61b1895a476d0bbf46811ab0c7fdf4bd127345/numpy-2.3.4-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:84f01a4d18b2cc4ade1814a08e5f3c907b079c847051d720fad15ce37aa930b6", size = 6739342, upload-time = "2025-10-15T16:16:43.777Z" }, - { url = "https://files.pythonhosted.org/packages/34/f1/4de9586d05b1962acdcdb1dc4af6646361a643f8c864cef7c852bf509740/numpy-2.3.4-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:817e719a868f0dacde4abdfc5c1910b301877970195db9ab6a5e2c4bd5b121f7", size = 14354338, upload-time = "2025-10-15T16:16:46.081Z" }, - { url = "https://files.pythonhosted.org/packages/1f/06/1c16103b425de7969d5a76bdf5ada0804b476fed05d5f9e17b777f1cbefd/numpy-2.3.4-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:85e071da78d92a214212cacea81c6da557cab307f2c34b5f85b628e94803f9c0", size = 16702392, upload-time = "2025-10-15T16:16:48.455Z" }, - { url = "https://files.pythonhosted.org/packages/34/b2/65f4dc1b89b5322093572b6e55161bb42e3e0487067af73627f795cc9d47/numpy-2.3.4-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:2ec646892819370cf3558f518797f16597b4e4669894a2ba712caccc9da53f1f", size = 16134998, upload-time = "2025-10-15T16:16:51.114Z" }, - { url = "https://files.pythonhosted.org/packages/d4/11/94ec578896cdb973aaf56425d6c7f2aff4186a5c00fac15ff2ec46998b46/numpy-2.3.4-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:035796aaaddfe2f9664b9a9372f089cfc88bd795a67bd1bfe15e6e770934cf64", size = 18651574, upload-time = "2025-10-15T16:16:53.429Z" }, - { url = "https://files.pythonhosted.org/packages/62/b7/7efa763ab33dbccf56dade36938a77345ce8e8192d6b39e470ca25ff3cd0/numpy-2.3.4-cp313-cp313t-win32.whl", hash = "sha256:fea80f4f4cf83b54c3a051f2f727870ee51e22f0248d3114b8e755d160b38cfb", size = 6413135, upload-time = "2025-10-15T16:16:55.992Z" }, - { url = "https://files.pythonhosted.org/packages/43/70/aba4c38e8400abcc2f345e13d972fb36c26409b3e644366db7649015f291/numpy-2.3.4-cp313-cp313t-win_amd64.whl", hash = "sha256:15eea9f306b98e0be91eb344a94c0e630689ef302e10c2ce5f7e11905c704f9c", size = 12928582, upload-time = "2025-10-15T16:16:57.943Z" }, - { url = "https://files.pythonhosted.org/packages/67/63/871fad5f0073fc00fbbdd7232962ea1ac40eeaae2bba66c76214f7954236/numpy-2.3.4-cp313-cp313t-win_arm64.whl", hash = "sha256:b6c231c9c2fadbae4011ca5e7e83e12dc4a5072f1a1d85a0a7b3ed754d145a40", size = 10266691, upload-time = "2025-10-15T16:17:00.048Z" }, - { url = "https://files.pythonhosted.org/packages/72/71/ae6170143c115732470ae3a2d01512870dd16e0953f8a6dc89525696069b/numpy-2.3.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:81c3e6d8c97295a7360d367f9f8553973651b76907988bb6066376bc2252f24e", size = 20955580, upload-time = "2025-10-15T16:17:02.509Z" }, - { url = "https://files.pythonhosted.org/packages/af/39/4be9222ffd6ca8a30eda033d5f753276a9c3426c397bb137d8e19dedd200/numpy-2.3.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:7c26b0b2bf58009ed1f38a641f3db4be8d960a417ca96d14e5b06df1506d41ff", size = 14188056, upload-time = "2025-10-15T16:17:04.873Z" }, - { url = "https://files.pythonhosted.org/packages/6c/3d/d85f6700d0a4aa4f9491030e1021c2b2b7421b2b38d01acd16734a2bfdc7/numpy-2.3.4-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:62b2198c438058a20b6704351b35a1d7db881812d8512d67a69c9de1f18ca05f", size = 5116555, upload-time = "2025-10-15T16:17:07.499Z" }, - { url = "https://files.pythonhosted.org/packages/bf/04/82c1467d86f47eee8a19a464c92f90a9bb68ccf14a54c5224d7031241ffb/numpy-2.3.4-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:9d729d60f8d53a7361707f4b68a9663c968882dd4f09e0d58c044c8bf5faee7b", size = 6643581, upload-time = "2025-10-15T16:17:09.774Z" }, - { url = "https://files.pythonhosted.org/packages/0c/d3/c79841741b837e293f48bd7db89d0ac7a4f2503b382b78a790ef1dc778a5/numpy-2.3.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bd0c630cf256b0a7fd9d0a11c9413b42fef5101219ce6ed5a09624f5a65392c7", size = 14299186, upload-time = "2025-10-15T16:17:11.937Z" }, - { url = "https://files.pythonhosted.org/packages/e8/7e/4a14a769741fbf237eec5a12a2cbc7a4c4e061852b6533bcb9e9a796c908/numpy-2.3.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d5e081bc082825f8b139f9e9fe42942cb4054524598aaeb177ff476cc76d09d2", size = 16638601, upload-time = "2025-10-15T16:17:14.391Z" }, - { url = "https://files.pythonhosted.org/packages/93/87/1c1de269f002ff0a41173fe01dcc925f4ecff59264cd8f96cf3b60d12c9b/numpy-2.3.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:15fb27364ed84114438fff8aaf998c9e19adbeba08c0b75409f8c452a8692c52", size = 16074219, upload-time = "2025-10-15T16:17:17.058Z" }, - { url = "https://files.pythonhosted.org/packages/cd/28/18f72ee77408e40a76d691001ae599e712ca2a47ddd2c4f695b16c65f077/numpy-2.3.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:85d9fb2d8cd998c84d13a79a09cc0c1091648e848e4e6249b0ccd7f6b487fa26", size = 18576702, upload-time = "2025-10-15T16:17:19.379Z" }, - { url = "https://files.pythonhosted.org/packages/c3/76/95650169b465ececa8cf4b2e8f6df255d4bf662775e797ade2025cc51ae6/numpy-2.3.4-cp314-cp314-win32.whl", hash = "sha256:e73d63fd04e3a9d6bc187f5455d81abfad05660b212c8804bf3b407e984cd2bc", size = 6337136, upload-time = "2025-10-15T16:17:22.886Z" }, - { url = "https://files.pythonhosted.org/packages/dc/89/a231a5c43ede5d6f77ba4a91e915a87dea4aeea76560ba4d2bf185c683f0/numpy-2.3.4-cp314-cp314-win_amd64.whl", hash = "sha256:3da3491cee49cf16157e70f607c03a217ea6647b1cea4819c4f48e53d49139b9", size = 12920542, upload-time = "2025-10-15T16:17:24.783Z" }, - { url = "https://files.pythonhosted.org/packages/0d/0c/ae9434a888f717c5ed2ff2393b3f344f0ff6f1c793519fa0c540461dc530/numpy-2.3.4-cp314-cp314-win_arm64.whl", hash = "sha256:6d9cd732068e8288dbe2717177320723ccec4fb064123f0caf9bbd90ab5be868", size = 10480213, upload-time = "2025-10-15T16:17:26.935Z" }, - { url = "https://files.pythonhosted.org/packages/83/4b/c4a5f0841f92536f6b9592694a5b5f68c9ab37b775ff342649eadf9055d3/numpy-2.3.4-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:22758999b256b595cf0b1d102b133bb61866ba5ceecf15f759623b64c020c9ec", size = 21052280, upload-time = "2025-10-15T16:17:29.638Z" }, - { url = "https://files.pythonhosted.org/packages/3e/80/90308845fc93b984d2cc96d83e2324ce8ad1fd6efea81b324cba4b673854/numpy-2.3.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:9cb177bc55b010b19798dc5497d540dea67fd13a8d9e882b2dae71de0cf09eb3", size = 14302930, upload-time = "2025-10-15T16:17:32.384Z" }, - { url = "https://files.pythonhosted.org/packages/3d/4e/07439f22f2a3b247cec4d63a713faae55e1141a36e77fb212881f7cda3fb/numpy-2.3.4-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:0f2bcc76f1e05e5ab58893407c63d90b2029908fa41f9f1cc51eecce936c3365", size = 5231504, upload-time = "2025-10-15T16:17:34.515Z" }, - { url = "https://files.pythonhosted.org/packages/ab/de/1e11f2547e2fe3d00482b19721855348b94ada8359aef5d40dd57bfae9df/numpy-2.3.4-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:8dc20bde86802df2ed8397a08d793da0ad7a5fd4ea3ac85d757bf5dd4ad7c252", size = 6739405, upload-time = "2025-10-15T16:17:36.128Z" }, - { url = "https://files.pythonhosted.org/packages/3b/40/8cd57393a26cebe2e923005db5134a946c62fa56a1087dc7c478f3e30837/numpy-2.3.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5e199c087e2aa71c8f9ce1cb7a8e10677dc12457e7cc1be4798632da37c3e86e", size = 14354866, upload-time = "2025-10-15T16:17:38.884Z" }, - { url = "https://files.pythonhosted.org/packages/93/39/5b3510f023f96874ee6fea2e40dfa99313a00bf3ab779f3c92978f34aace/numpy-2.3.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:85597b2d25ddf655495e2363fe044b0ae999b75bc4d630dc0d886484b03a5eb0", size = 16703296, upload-time = "2025-10-15T16:17:41.564Z" }, - { url = "https://files.pythonhosted.org/packages/41/0d/19bb163617c8045209c1996c4e427bccbc4bbff1e2c711f39203c8ddbb4a/numpy-2.3.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:04a69abe45b49c5955923cf2c407843d1c85013b424ae8a560bba16c92fe44a0", size = 16136046, upload-time = "2025-10-15T16:17:43.901Z" }, - { url = "https://files.pythonhosted.org/packages/e2/c1/6dba12fdf68b02a21ac411c9df19afa66bed2540f467150ca64d246b463d/numpy-2.3.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:e1708fac43ef8b419c975926ce1eaf793b0c13b7356cfab6ab0dc34c0a02ac0f", size = 18652691, upload-time = "2025-10-15T16:17:46.247Z" }, - { url = "https://files.pythonhosted.org/packages/f8/73/f85056701dbbbb910c51d846c58d29fd46b30eecd2b6ba760fc8b8a1641b/numpy-2.3.4-cp314-cp314t-win32.whl", hash = "sha256:863e3b5f4d9915aaf1b8ec79ae560ad21f0b8d5e3adc31e73126491bb86dee1d", size = 6485782, upload-time = "2025-10-15T16:17:48.872Z" }, - { url = "https://files.pythonhosted.org/packages/17/90/28fa6f9865181cb817c2471ee65678afa8a7e2a1fb16141473d5fa6bacc3/numpy-2.3.4-cp314-cp314t-win_amd64.whl", hash = "sha256:962064de37b9aef801d33bc579690f8bfe6c5e70e29b61783f60bcba838a14d6", size = 13113301, upload-time = "2025-10-15T16:17:50.938Z" }, - { url = "https://files.pythonhosted.org/packages/54/23/08c002201a8e7e1f9afba93b97deceb813252d9cfd0d3351caed123dcf97/numpy-2.3.4-cp314-cp314t-win_arm64.whl", hash = "sha256:8b5a9a39c45d852b62693d9b3f3e0fe052541f804296ff401a72a1b60edafb29", size = 10547532, upload-time = "2025-10-15T16:17:53.48Z" }, - { url = "https://files.pythonhosted.org/packages/b1/b6/64898f51a86ec88ca1257a59c1d7fd077b60082a119affefcdf1dd0df8ca/numpy-2.3.4-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:6e274603039f924c0fe5cb73438fa9246699c78a6df1bd3decef9ae592ae1c05", size = 21131552, upload-time = "2025-10-15T16:17:55.845Z" }, - { url = "https://files.pythonhosted.org/packages/ce/4c/f135dc6ebe2b6a3c77f4e4838fa63d350f85c99462012306ada1bd4bc460/numpy-2.3.4-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d149aee5c72176d9ddbc6803aef9c0f6d2ceeea7626574fc68518da5476fa346", size = 14377796, upload-time = "2025-10-15T16:17:58.308Z" }, - { url = "https://files.pythonhosted.org/packages/d0/a4/f33f9c23fcc13dd8412fc8614559b5b797e0aba9d8e01dfa8bae10c84004/numpy-2.3.4-pp311-pypy311_pp73-macosx_14_0_arm64.whl", hash = "sha256:6d34ed9db9e6395bb6cd33286035f73a59b058169733a9db9f85e650b88df37e", size = 5306904, upload-time = "2025-10-15T16:18:00.596Z" }, - { url = "https://files.pythonhosted.org/packages/28/af/c44097f25f834360f9fb960fa082863e0bad14a42f36527b2a121abdec56/numpy-2.3.4-pp311-pypy311_pp73-macosx_14_0_x86_64.whl", hash = "sha256:fdebe771ca06bb8d6abce84e51dca9f7921fe6ad34a0c914541b063e9a68928b", size = 6819682, upload-time = "2025-10-15T16:18:02.32Z" }, - { url = "https://files.pythonhosted.org/packages/c5/8c/cd283b54c3c2b77e188f63e23039844f56b23bba1712318288c13fe86baf/numpy-2.3.4-pp311-pypy311_pp73-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:957e92defe6c08211eb77902253b14fe5b480ebc5112bc741fd5e9cd0608f847", size = 14422300, upload-time = "2025-10-15T16:18:04.271Z" }, - { url = "https://files.pythonhosted.org/packages/b0/f0/8404db5098d92446b3e3695cf41c6f0ecb703d701cb0b7566ee2177f2eee/numpy-2.3.4-pp311-pypy311_pp73-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:13b9062e4f5c7ee5c7e5be96f29ba71bc5a37fed3d1d77c37390ae00724d296d", size = 16760806, upload-time = "2025-10-15T16:18:06.668Z" }, - { url = "https://files.pythonhosted.org/packages/95/8e/2844c3959ce9a63acc7c8e50881133d86666f0420bcde695e115ced0920f/numpy-2.3.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:81b3a59793523e552c4a96109dde028aa4448ae06ccac5a76ff6532a85558a7f", size = 12973130, upload-time = "2025-10-15T16:18:09.397Z" }, +version = "1.26.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/6e/09db70a523a96d25e115e71cc56a6f9031e7b8cd166c1ac8438307c14058/numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010", size = 15786129, upload-time = "2024-02-06T00:26:44.495Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/94/ace0fdea5241a27d13543ee117cbc65868e82213fb31a8eb7fe9ff23f313/numpy-1.26.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9ff0f4f29c51e2803569d7a51c2304de5554655a60c5d776e35b4a41413830d0", size = 20631468, upload-time = "2024-02-05T23:48:01.194Z" }, + { url = "https://files.pythonhosted.org/packages/20/f7/b24208eba89f9d1b58c1668bc6c8c4fd472b20c45573cb767f59d49fb0f6/numpy-1.26.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2e4ee3380d6de9c9ec04745830fd9e2eccb3e6cf790d39d7b98ffd19b0dd754a", size = 13966411, upload-time = "2024-02-05T23:48:29.038Z" }, + { url = "https://files.pythonhosted.org/packages/fc/a5/4beee6488160798683eed5bdb7eead455892c3b4e1f78d79d8d3f3b084ac/numpy-1.26.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d209d8969599b27ad20994c8e41936ee0964e6da07478d6c35016bc386b66ad4", size = 14219016, upload-time = "2024-02-05T23:48:54.098Z" }, + { url = "https://files.pythonhosted.org/packages/4b/d7/ecf66c1cd12dc28b4040b15ab4d17b773b87fa9d29ca16125de01adb36cd/numpy-1.26.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ffa75af20b44f8dba823498024771d5ac50620e6915abac414251bd971b4529f", size = 18240889, upload-time = "2024-02-05T23:49:25.361Z" }, + { url = "https://files.pythonhosted.org/packages/24/03/6f229fe3187546435c4f6f89f6d26c129d4f5bed40552899fcf1f0bf9e50/numpy-1.26.4-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:62b8e4b1e28009ef2846b4c7852046736bab361f7aeadeb6a5b89ebec3c7055a", size = 13876746, upload-time = "2024-02-05T23:49:51.983Z" }, + { url = "https://files.pythonhosted.org/packages/39/fe/39ada9b094f01f5a35486577c848fe274e374bbf8d8f472e1423a0bbd26d/numpy-1.26.4-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:a4abb4f9001ad2858e7ac189089c42178fcce737e4169dc61321660f1a96c7d2", size = 18078620, upload-time = "2024-02-05T23:50:22.515Z" }, + { url = "https://files.pythonhosted.org/packages/d5/ef/6ad11d51197aad206a9ad2286dc1aac6a378059e06e8cf22cd08ed4f20dc/numpy-1.26.4-cp310-cp310-win32.whl", hash = "sha256:bfe25acf8b437eb2a8b2d49d443800a5f18508cd811fea3181723922a8a82b07", size = 5972659, upload-time = "2024-02-05T23:50:35.834Z" }, + { url = "https://files.pythonhosted.org/packages/19/77/538f202862b9183f54108557bfda67e17603fc560c384559e769321c9d92/numpy-1.26.4-cp310-cp310-win_amd64.whl", hash = "sha256:b97fe8060236edf3662adfc2c633f56a08ae30560c56310562cb4f95500022d5", size = 15808905, upload-time = "2024-02-05T23:51:03.701Z" }, + { url = "https://files.pythonhosted.org/packages/11/57/baae43d14fe163fa0e4c47f307b6b2511ab8d7d30177c491960504252053/numpy-1.26.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4c66707fabe114439db9068ee468c26bbdf909cac0fb58686a42a24de1760c71", size = 20630554, upload-time = "2024-02-05T23:51:50.149Z" }, + { url = "https://files.pythonhosted.org/packages/1a/2e/151484f49fd03944c4a3ad9c418ed193cfd02724e138ac8a9505d056c582/numpy-1.26.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:edd8b5fe47dab091176d21bb6de568acdd906d1887a4584a15a9a96a1dca06ef", size = 13997127, upload-time = "2024-02-05T23:52:15.314Z" }, + { url = "https://files.pythonhosted.org/packages/79/ae/7e5b85136806f9dadf4878bf73cf223fe5c2636818ba3ab1c585d0403164/numpy-1.26.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7ab55401287bfec946ced39700c053796e7cc0e3acbef09993a9ad2adba6ca6e", size = 14222994, upload-time = "2024-02-05T23:52:47.569Z" }, + { url = "https://files.pythonhosted.org/packages/3a/d0/edc009c27b406c4f9cbc79274d6e46d634d139075492ad055e3d68445925/numpy-1.26.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:666dbfb6ec68962c033a450943ded891bed2d54e6755e35e5835d63f4f6931d5", size = 18252005, upload-time = "2024-02-05T23:53:15.637Z" }, + { url = "https://files.pythonhosted.org/packages/09/bf/2b1aaf8f525f2923ff6cfcf134ae5e750e279ac65ebf386c75a0cf6da06a/numpy-1.26.4-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:96ff0b2ad353d8f990b63294c8986f1ec3cb19d749234014f4e7eb0112ceba5a", size = 13885297, upload-time = "2024-02-05T23:53:42.16Z" }, + { url = "https://files.pythonhosted.org/packages/df/a0/4e0f14d847cfc2a633a1c8621d00724f3206cfeddeb66d35698c4e2cf3d2/numpy-1.26.4-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:60dedbb91afcbfdc9bc0b1f3f402804070deed7392c23eb7a7f07fa857868e8a", size = 18093567, upload-time = "2024-02-05T23:54:11.696Z" }, + { url = "https://files.pythonhosted.org/packages/d2/b7/a734c733286e10a7f1a8ad1ae8c90f2d33bf604a96548e0a4a3a6739b468/numpy-1.26.4-cp311-cp311-win32.whl", hash = "sha256:1af303d6b2210eb850fcf03064d364652b7120803a0b872f5211f5234b399f20", size = 5968812, upload-time = "2024-02-05T23:54:26.453Z" }, + { url = "https://files.pythonhosted.org/packages/3f/6b/5610004206cf7f8e7ad91c5a85a8c71b2f2f8051a0c0c4d5916b76d6cbb2/numpy-1.26.4-cp311-cp311-win_amd64.whl", hash = "sha256:cd25bcecc4974d09257ffcd1f098ee778f7834c3ad767fe5db785be9a4aa9cb2", size = 15811913, upload-time = "2024-02-05T23:54:53.933Z" }, + { url = "https://files.pythonhosted.org/packages/95/12/8f2020a8e8b8383ac0177dc9570aad031a3beb12e38847f7129bacd96228/numpy-1.26.4-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:b3ce300f3644fb06443ee2222c2201dd3a89ea6040541412b8fa189341847218", size = 20335901, upload-time = "2024-02-05T23:55:32.801Z" }, + { url = "https://files.pythonhosted.org/packages/75/5b/ca6c8bd14007e5ca171c7c03102d17b4f4e0ceb53957e8c44343a9546dcc/numpy-1.26.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:03a8c78d01d9781b28a6989f6fa1bb2c4f2d51201cf99d3dd875df6fbd96b23b", size = 13685868, upload-time = "2024-02-05T23:55:56.28Z" }, + { url = "https://files.pythonhosted.org/packages/79/f8/97f10e6755e2a7d027ca783f63044d5b1bc1ae7acb12afe6a9b4286eac17/numpy-1.26.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9fad7dcb1aac3c7f0584a5a8133e3a43eeb2fe127f47e3632d43d677c66c102b", size = 13925109, upload-time = "2024-02-05T23:56:20.368Z" }, + { url = "https://files.pythonhosted.org/packages/0f/50/de23fde84e45f5c4fda2488c759b69990fd4512387a8632860f3ac9cd225/numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed", size = 17950613, upload-time = "2024-02-05T23:56:56.054Z" }, + { url = "https://files.pythonhosted.org/packages/4c/0c/9c603826b6465e82591e05ca230dfc13376da512b25ccd0894709b054ed0/numpy-1.26.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:ab47dbe5cc8210f55aa58e4805fe224dac469cde56b9f731a4c098b91917159a", size = 13572172, upload-time = "2024-02-05T23:57:21.56Z" }, + { url = "https://files.pythonhosted.org/packages/76/8c/2ba3902e1a0fc1c74962ea9bb33a534bb05984ad7ff9515bf8d07527cadd/numpy-1.26.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0", size = 17786643, upload-time = "2024-02-05T23:57:56.585Z" }, + { url = "https://files.pythonhosted.org/packages/28/4a/46d9e65106879492374999e76eb85f87b15328e06bd1550668f79f7b18c6/numpy-1.26.4-cp312-cp312-win32.whl", hash = "sha256:50193e430acfc1346175fcbdaa28ffec49947a06918b7b92130744e81e640110", size = 5677803, upload-time = "2024-02-05T23:58:08.963Z" }, + { url = "https://files.pythonhosted.org/packages/16/2e/86f24451c2d530c88daf997cb8d6ac622c1d40d19f5a031ed68a4b73a374/numpy-1.26.4-cp312-cp312-win_amd64.whl", hash = "sha256:08beddf13648eb95f8d867350f6a018a4be2e5ad54c8d8caed89ebca558b2818", size = 15517754, upload-time = "2024-02-05T23:58:36.364Z" }, ] [[package]] @@ -1883,7 +2528,7 @@ name = "nvidia-cudnn-cu12" version = "9.10.2.21" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, @@ -1894,7 +2539,7 @@ name = "nvidia-cufft-cu12" version = "11.3.3.83" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, @@ -1921,9 +2566,9 @@ name = "nvidia-cusolver-cu12" version = "11.7.3.90" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, @@ -1934,7 +2579,7 @@ name = "nvidia-cusparse-cu12" version = "12.5.8.93" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and sys_platform == 'linux'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, @@ -1985,6 +2630,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, ] +[[package]] +name = "opencv-contrib-python" +version = "4.11.0.86" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ef/51/3ceb85ecff5f26994b7aae2922b1aa38148dbfe88cab13d63bc6facbac88/opencv-contrib-python-4.11.0.86.tar.gz", hash = "sha256:4ff773dab44911da366b906621c9592d4eb96f6ad3777098933a23f064aab38e", size = 150559874, upload-time = "2025-01-16T13:53:08.425Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f3/78/b504ca8f7a312918d184e0b8093c62bc9a110d8154f658b591ef5c020d65/opencv_contrib_python-4.11.0.86-cp37-abi3-macosx_13_0_arm64.whl", hash = "sha256:d911cedc511d98f79994580b245d59fc97f57f0f9923a99945d8b92c7ac671f6", size = 46276766, upload-time = "2025-01-16T13:52:46.131Z" }, + { url = "https://files.pythonhosted.org/packages/8c/07/68e0b24217671b65c23e105bb7afd4ef4fd01507670cf5e61373d9efd6b5/opencv_contrib_python-4.11.0.86-cp37-abi3-macosx_13_0_x86_64.whl", hash = "sha256:e10a293af18aa5f842d012fa14e87345b3ee06db4c29bd592ff94b51f7ffca2b", size = 66524088, upload-time = "2025-01-16T13:55:33.887Z" }, + { url = "https://files.pythonhosted.org/packages/ae/7b/7e1471aa92f9f3c1bd8dbe624622b62add6f734db34fbbb9974e2ec70c34/opencv_contrib_python-4.11.0.86-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f21034bc8b00eb286a0a0a92b99767bf596bfe426cf4bc2e79647d64ad0dd6da", size = 47870560, upload-time = "2025-01-16T13:51:48.592Z" }, + { url = "https://files.pythonhosted.org/packages/f7/13/756b13b8d5d417a0b4c3bf6ceafb59df0ed05cec7fedc2490bbbf5e60ebc/opencv_contrib_python-4.11.0.86-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c47c0ef1098461cdc6fa1cdce4c942b8ec974c87423f4b5951443d26bb9ae407", size = 69098423, upload-time = "2025-01-16T13:52:46.84Z" }, + { url = "https://files.pythonhosted.org/packages/fd/8b/4f63d2fdcfceab528bff10c9d8d2a4e6230098e0b0af54e3e8e91b420ea0/opencv_contrib_python-4.11.0.86-cp37-abi3-win32.whl", hash = "sha256:194841c664ceaa0692410b4ed0af557425608e33db3a181ded28b87acb66748d", size = 35156028, upload-time = "2025-01-16T13:52:30.133Z" }, + { url = "https://files.pythonhosted.org/packages/0d/c6/146487546adc4726f0be591a65b466973feaa58cc3db711087e802e940fb/opencv_contrib_python-4.11.0.86-cp37-abi3-win_amd64.whl", hash = "sha256:654758a9ae8ca9a75fca7b64b19163636534f0eedffe1e14c3d7218988625c8d", size = 46185163, upload-time = "2025-01-16T13:52:39.745Z" }, +] + +[[package]] +name = "opt-einsum" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8c/b9/2ac072041e899a52f20cf9510850ff58295003aa75525e58343591b0cbfb/opt_einsum-3.4.0.tar.gz", hash = "sha256:96ca72f1b886d148241348783498194c577fa30a8faac108586b14f1ba4473ac", size = 63004, upload-time = "2024-09-26T14:33:24.483Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl", hash = "sha256:69bb92469f86a1565195ece4ac0323943e83477171b91d24c35afe028a90d7cd", size = 71932, upload-time = "2024-09-26T14:33:23.039Z" }, +] + [[package]] name = "packaging" version = "25.0" @@ -2001,8 +2672,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "accelerate" }, { name = "huggingface-hub" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy" }, { name = "packaging" }, { name = "psutil" }, { name = "pyyaml" }, @@ -2263,6 +2933,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/5a/bc7b4a4ef808fa59a816c17b20c4bef6884daebbdf627ff2a161da67da19/propcache-0.4.1-py3-none-any.whl", hash = "sha256:af2a6052aeb6cf17d3e46ee169099044fd8224cbaf75c76a2ef596e8163e2237", size = 13305, upload-time = "2025-10-08T19:49:00.792Z" }, ] +[[package]] +name = "protobuf" +version = "4.25.8" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/df/01/34c8d2b6354906d728703cb9d546a0e534de479e25f1b581e4094c4a85cc/protobuf-4.25.8.tar.gz", hash = "sha256:6135cf8affe1fc6f76cced2641e4ea8d3e59518d1f24ae41ba97bcad82d397cd", size = 380920, upload-time = "2025-05-28T14:22:25.153Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/45/ff/05f34305fe6b85bbfbecbc559d423a5985605cad5eda4f47eae9e9c9c5c5/protobuf-4.25.8-cp310-abi3-win32.whl", hash = "sha256:504435d831565f7cfac9f0714440028907f1975e4bed228e58e72ecfff58a1e0", size = 392745, upload-time = "2025-05-28T14:22:10.524Z" }, + { url = "https://files.pythonhosted.org/packages/08/35/8b8a8405c564caf4ba835b1fdf554da869954712b26d8f2a98c0e434469b/protobuf-4.25.8-cp310-abi3-win_amd64.whl", hash = "sha256:bd551eb1fe1d7e92c1af1d75bdfa572eff1ab0e5bf1736716814cdccdb2360f9", size = 413736, upload-time = "2025-05-28T14:22:13.156Z" }, + { url = "https://files.pythonhosted.org/packages/28/d7/ab27049a035b258dab43445eb6ec84a26277b16105b277cbe0a7698bdc6c/protobuf-4.25.8-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:ca809b42f4444f144f2115c4c1a747b9a404d590f18f37e9402422033e464e0f", size = 394537, upload-time = "2025-05-28T14:22:14.768Z" }, + { url = "https://files.pythonhosted.org/packages/bd/6d/a4a198b61808dd3d1ee187082ccc21499bc949d639feb948961b48be9a7e/protobuf-4.25.8-cp37-abi3-manylinux2014_aarch64.whl", hash = "sha256:9ad7ef62d92baf5a8654fbb88dac7fa5594cfa70fd3440488a5ca3bfc6d795a7", size = 294005, upload-time = "2025-05-28T14:22:16.052Z" }, + { url = "https://files.pythonhosted.org/packages/d6/c6/c9deaa6e789b6fc41b88ccbdfe7a42d2b82663248b715f55aa77fbc00724/protobuf-4.25.8-cp37-abi3-manylinux2014_x86_64.whl", hash = "sha256:83e6e54e93d2b696a92cad6e6efc924f3850f82b52e1563778dfab8b355101b0", size = 294924, upload-time = "2025-05-28T14:22:17.105Z" }, + { url = "https://files.pythonhosted.org/packages/0c/c1/6aece0ab5209981a70cd186f164c133fdba2f51e124ff92b73de7fd24d78/protobuf-4.25.8-py3-none-any.whl", hash = "sha256:15a0af558aa3b13efef102ae6e4f3efac06f1eea11afb3a57db2901447d9fb59", size = 156757, upload-time = "2025-05-28T14:22:24.135Z" }, +] + [[package]] name = "psutil" version = "7.1.3" @@ -2505,6 +3189,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/81/ef2b1dfd1862567d573a4fdbc9f969067621764fbb74338496840a1d2977/pyopenssl-25.3.0-py3-none-any.whl", hash = "sha256:1fda6fc034d5e3d179d39e59c1895c9faeaf40a79de5fc4cbbfbe0d36f4a77b6", size = 57268, upload-time = "2025-09-17T00:32:19.474Z" }, ] +[[package]] +name = "pyparsing" +version = "3.2.5" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f2/a5/181488fc2b9d093e3972d2a472855aae8a03f000592dbfce716a512b3359/pyparsing-3.2.5.tar.gz", hash = "sha256:2df8d5b7b2802ef88e8d016a2eb9c7aeaa923529cd251ed0fe4608275d4105b6", size = 1099274, upload-time = "2025-09-21T04:11:06.277Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/5e/1aa9a93198c6b64513c9d7752de7422c06402de6600a8767da1524f9570b/pyparsing-3.2.5-py3-none-any.whl", hash = "sha256:e38a4f02064cf41fe6593d328d0512495ad1f3d8a91c4f73fc401b3079a59a5e", size = 113890, upload-time = "2025-09-21T04:11:04.117Z" }, +] + [[package]] name = "pytest" version = "8.4.2" @@ -2831,9 +3524,14 @@ name = "sageattention" version = "2.2.0" source = { url = "https://github.com/daydreamlive/SageAttention/releases/download/v2.2.0-linux/sageattention-2.2.0-cp310-cp310-linux_x86_64.whl" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", ] wheels = [ { url = "https://github.com/daydreamlive/SageAttention/releases/download/v2.2.0-linux/sageattention-2.2.0-cp310-cp310-linux_x86_64.whl", hash = "sha256:30f23d0976e89738808aac3246094579807e5710f2f9741860387aa4c777f244" }, @@ -2844,7 +3542,8 @@ name = "sageattention" version = "2.2.0+cu128torch2.8.0.post3" source = { url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version < '3.11' and sys_platform == 'win32'", ] @@ -2852,6 +3551,199 @@ wheels = [ { url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl", hash = "sha256:7dabcd00e63229b28f046c5a69ec37cf4756afb375dbadd1975dadec045ae21c" }, ] +[[package]] +name = "scikit-image" +version = "0.25.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "imageio" }, + { name = "lazy-loader" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy" }, + { name = "packaging" }, + { name = "pillow" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "tifffile", version = "2025.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "tifffile", version = "2025.12.12", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c7/a8/3c0f256012b93dd2cb6fda9245e9f4bff7dc0486880b248005f15ea2255e/scikit_image-0.25.2.tar.gz", hash = "sha256:e5a37e6cd4d0c018a7a55b9d601357e3382826d3888c10d0213fc63bff977dde", size = 22693594, upload-time = "2025-02-18T18:05:24.538Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/cb/016c63f16065c2d333c8ed0337e18a5cdf9bc32d402e4f26b0db362eb0e2/scikit_image-0.25.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d3278f586793176599df6a4cf48cb6beadae35c31e58dc01a98023af3dc31c78", size = 13988922, upload-time = "2025-02-18T18:04:11.069Z" }, + { url = "https://files.pythonhosted.org/packages/30/ca/ff4731289cbed63c94a0c9a5b672976603118de78ed21910d9060c82e859/scikit_image-0.25.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:5c311069899ce757d7dbf1d03e32acb38bb06153236ae77fcd820fd62044c063", size = 13192698, upload-time = "2025-02-18T18:04:15.362Z" }, + { url = "https://files.pythonhosted.org/packages/39/6d/a2aadb1be6d8e149199bb9b540ccde9e9622826e1ab42fe01de4c35ab918/scikit_image-0.25.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be455aa7039a6afa54e84f9e38293733a2622b8c2fb3362b822d459cc5605e99", size = 14153634, upload-time = "2025-02-18T18:04:18.496Z" }, + { url = "https://files.pythonhosted.org/packages/96/08/916e7d9ee4721031b2f625db54b11d8379bd51707afaa3e5a29aecf10bc4/scikit_image-0.25.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a4c464b90e978d137330be433df4e76d92ad3c5f46a22f159520ce0fdbea8a09", size = 14767545, upload-time = "2025-02-18T18:04:22.556Z" }, + { url = "https://files.pythonhosted.org/packages/5f/ee/c53a009e3997dda9d285402f19226fbd17b5b3cb215da391c4ed084a1424/scikit_image-0.25.2-cp310-cp310-win_amd64.whl", hash = "sha256:60516257c5a2d2f74387c502aa2f15a0ef3498fbeaa749f730ab18f0a40fd054", size = 12812908, upload-time = "2025-02-18T18:04:26.364Z" }, + { url = "https://files.pythonhosted.org/packages/c4/97/3051c68b782ee3f1fb7f8f5bb7d535cf8cb92e8aae18fa9c1cdf7e15150d/scikit_image-0.25.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f4bac9196fb80d37567316581c6060763b0f4893d3aca34a9ede3825bc035b17", size = 14003057, upload-time = "2025-02-18T18:04:30.395Z" }, + { url = "https://files.pythonhosted.org/packages/19/23/257fc696c562639826065514d551b7b9b969520bd902c3a8e2fcff5b9e17/scikit_image-0.25.2-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d989d64ff92e0c6c0f2018c7495a5b20e2451839299a018e0e5108b2680f71e0", size = 13180335, upload-time = "2025-02-18T18:04:33.449Z" }, + { url = "https://files.pythonhosted.org/packages/ef/14/0c4a02cb27ca8b1e836886b9ec7c9149de03053650e9e2ed0625f248dd92/scikit_image-0.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b2cfc96b27afe9a05bc92f8c6235321d3a66499995675b27415e0d0c76625173", size = 14144783, upload-time = "2025-02-18T18:04:36.594Z" }, + { url = "https://files.pythonhosted.org/packages/dd/9b/9fb556463a34d9842491d72a421942c8baff4281025859c84fcdb5e7e602/scikit_image-0.25.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:24cc986e1f4187a12aa319f777b36008764e856e5013666a4a83f8df083c2641", size = 14785376, upload-time = "2025-02-18T18:04:39.856Z" }, + { url = "https://files.pythonhosted.org/packages/de/ec/b57c500ee85885df5f2188f8bb70398481393a69de44a00d6f1d055f103c/scikit_image-0.25.2-cp311-cp311-win_amd64.whl", hash = "sha256:b4f6b61fc2db6340696afe3db6b26e0356911529f5f6aee8c322aa5157490c9b", size = 12791698, upload-time = "2025-02-18T18:04:42.868Z" }, + { url = "https://files.pythonhosted.org/packages/35/8c/5df82881284459f6eec796a5ac2a0a304bb3384eec2e73f35cfdfcfbf20c/scikit_image-0.25.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8db8dd03663112783221bf01ccfc9512d1cc50ac9b5b0fe8f4023967564719fb", size = 13986000, upload-time = "2025-02-18T18:04:47.156Z" }, + { url = "https://files.pythonhosted.org/packages/ce/e6/93bebe1abcdce9513ffec01d8af02528b4c41fb3c1e46336d70b9ed4ef0d/scikit_image-0.25.2-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:483bd8cc10c3d8a7a37fae36dfa5b21e239bd4ee121d91cad1f81bba10cfb0ed", size = 13235893, upload-time = "2025-02-18T18:04:51.049Z" }, + { url = "https://files.pythonhosted.org/packages/53/4b/eda616e33f67129e5979a9eb33c710013caa3aa8a921991e6cc0b22cea33/scikit_image-0.25.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d1e80107bcf2bf1291acfc0bf0425dceb8890abe9f38d8e94e23497cbf7ee0d", size = 14178389, upload-time = "2025-02-18T18:04:54.245Z" }, + { url = "https://files.pythonhosted.org/packages/6b/b5/b75527c0f9532dd8a93e8e7cd8e62e547b9f207d4c11e24f0006e8646b36/scikit_image-0.25.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a17e17eb8562660cc0d31bb55643a4da996a81944b82c54805c91b3fe66f4824", size = 15003435, upload-time = "2025-02-18T18:04:57.586Z" }, + { url = "https://files.pythonhosted.org/packages/34/e3/49beb08ebccda3c21e871b607c1cb2f258c3fa0d2f609fed0a5ba741b92d/scikit_image-0.25.2-cp312-cp312-win_amd64.whl", hash = "sha256:bdd2b8c1de0849964dbc54037f36b4e9420157e67e45a8709a80d727f52c7da2", size = 12899474, upload-time = "2025-02-18T18:05:01.166Z" }, + { url = "https://files.pythonhosted.org/packages/e6/7c/9814dd1c637f7a0e44342985a76f95a55dd04be60154247679fd96c7169f/scikit_image-0.25.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:7efa888130f6c548ec0439b1a7ed7295bc10105458a421e9bf739b457730b6da", size = 13921841, upload-time = "2025-02-18T18:05:03.963Z" }, + { url = "https://files.pythonhosted.org/packages/84/06/66a2e7661d6f526740c309e9717d3bd07b473661d5cdddef4dd978edab25/scikit_image-0.25.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:dd8011efe69c3641920614d550f5505f83658fe33581e49bed86feab43a180fc", size = 13196862, upload-time = "2025-02-18T18:05:06.986Z" }, + { url = "https://files.pythonhosted.org/packages/4e/63/3368902ed79305f74c2ca8c297dfeb4307269cbe6402412668e322837143/scikit_image-0.25.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:28182a9d3e2ce3c2e251383bdda68f8d88d9fff1a3ebe1eb61206595c9773341", size = 14117785, upload-time = "2025-02-18T18:05:10.69Z" }, + { url = "https://files.pythonhosted.org/packages/cd/9b/c3da56a145f52cd61a68b8465d6a29d9503bc45bc993bb45e84371c97d94/scikit_image-0.25.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8abd3c805ce6944b941cfed0406d88faeb19bab3ed3d4b50187af55cf24d147", size = 14977119, upload-time = "2025-02-18T18:05:13.871Z" }, + { url = "https://files.pythonhosted.org/packages/8a/97/5fcf332e1753831abb99a2525180d3fb0d70918d461ebda9873f66dcc12f/scikit_image-0.25.2-cp313-cp313-win_amd64.whl", hash = "sha256:64785a8acefee460ec49a354706db0b09d1f325674107d7fa3eadb663fb56d6f", size = 12885116, upload-time = "2025-02-18T18:05:17.844Z" }, + { url = "https://files.pythonhosted.org/packages/10/cc/75e9f17e3670b5ed93c32456fda823333c6279b144cd93e2c03aa06aa472/scikit_image-0.25.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:330d061bd107d12f8d68f1d611ae27b3b813b8cdb0300a71d07b1379178dd4cd", size = 13862801, upload-time = "2025-02-18T18:05:20.783Z" }, +] + +[[package]] +name = "scipy" +version = "1.15.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0f/37/6964b830433e654ec7485e45a00fc9a27cf868d622838f6b6d9c5ec0d532/scipy-1.15.3.tar.gz", hash = "sha256:eae3cf522bc7df64b42cad3925c876e1b0b6c35c1337c93e12c0f366f55b0eaf", size = 59419214, upload-time = "2025-05-08T16:13:05.955Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/2f/4966032c5f8cc7e6a60f1b2e0ad686293b9474b65246b0c642e3ef3badd0/scipy-1.15.3-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:a345928c86d535060c9c2b25e71e87c39ab2f22fc96e9636bd74d1dbf9de448c", size = 38702770, upload-time = "2025-05-08T16:04:20.849Z" }, + { url = "https://files.pythonhosted.org/packages/a0/6e/0c3bf90fae0e910c274db43304ebe25a6b391327f3f10b5dcc638c090795/scipy-1.15.3-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:ad3432cb0f9ed87477a8d97f03b763fd1d57709f1bbde3c9369b1dff5503b253", size = 30094511, upload-time = "2025-05-08T16:04:27.103Z" }, + { url = "https://files.pythonhosted.org/packages/ea/b1/4deb37252311c1acff7f101f6453f0440794f51b6eacb1aad4459a134081/scipy-1.15.3-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:aef683a9ae6eb00728a542b796f52a5477b78252edede72b8327a886ab63293f", size = 22368151, upload-time = "2025-05-08T16:04:31.731Z" }, + { url = "https://files.pythonhosted.org/packages/38/7d/f457626e3cd3c29b3a49ca115a304cebb8cc6f31b04678f03b216899d3c6/scipy-1.15.3-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:1c832e1bd78dea67d5c16f786681b28dd695a8cb1fb90af2e27580d3d0967e92", size = 25121732, upload-time = "2025-05-08T16:04:36.596Z" }, + { url = "https://files.pythonhosted.org/packages/db/0a/92b1de4a7adc7a15dcf5bddc6e191f6f29ee663b30511ce20467ef9b82e4/scipy-1.15.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:263961f658ce2165bbd7b99fa5135195c3a12d9bef045345016b8b50c315cb82", size = 35547617, upload-time = "2025-05-08T16:04:43.546Z" }, + { url = "https://files.pythonhosted.org/packages/8e/6d/41991e503e51fc1134502694c5fa7a1671501a17ffa12716a4a9151af3df/scipy-1.15.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e2abc762b0811e09a0d3258abee2d98e0c703eee49464ce0069590846f31d40", size = 37662964, upload-time = "2025-05-08T16:04:49.431Z" }, + { url = "https://files.pythonhosted.org/packages/25/e1/3df8f83cb15f3500478c889be8fb18700813b95e9e087328230b98d547ff/scipy-1.15.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:ed7284b21a7a0c8f1b6e5977ac05396c0d008b89e05498c8b7e8f4a1423bba0e", size = 37238749, upload-time = "2025-05-08T16:04:55.215Z" }, + { url = "https://files.pythonhosted.org/packages/93/3e/b3257cf446f2a3533ed7809757039016b74cd6f38271de91682aa844cfc5/scipy-1.15.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:5380741e53df2c566f4d234b100a484b420af85deb39ea35a1cc1be84ff53a5c", size = 40022383, upload-time = "2025-05-08T16:05:01.914Z" }, + { url = "https://files.pythonhosted.org/packages/d1/84/55bc4881973d3f79b479a5a2e2df61c8c9a04fcb986a213ac9c02cfb659b/scipy-1.15.3-cp310-cp310-win_amd64.whl", hash = "sha256:9d61e97b186a57350f6d6fd72640f9e99d5a4a2b8fbf4b9ee9a841eab327dc13", size = 41259201, upload-time = "2025-05-08T16:05:08.166Z" }, + { url = "https://files.pythonhosted.org/packages/96/ab/5cc9f80f28f6a7dff646c5756e559823614a42b1939d86dd0ed550470210/scipy-1.15.3-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:993439ce220d25e3696d1b23b233dd010169b62f6456488567e830654ee37a6b", size = 38714255, upload-time = "2025-05-08T16:05:14.596Z" }, + { url = "https://files.pythonhosted.org/packages/4a/4a/66ba30abe5ad1a3ad15bfb0b59d22174012e8056ff448cb1644deccbfed2/scipy-1.15.3-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:34716e281f181a02341ddeaad584205bd2fd3c242063bd3423d61ac259ca7eba", size = 30111035, upload-time = "2025-05-08T16:05:20.152Z" }, + { url = "https://files.pythonhosted.org/packages/4b/fa/a7e5b95afd80d24313307f03624acc65801846fa75599034f8ceb9e2cbf6/scipy-1.15.3-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:3b0334816afb8b91dab859281b1b9786934392aa3d527cd847e41bb6f45bee65", size = 22384499, upload-time = "2025-05-08T16:05:24.494Z" }, + { url = "https://files.pythonhosted.org/packages/17/99/f3aaddccf3588bb4aea70ba35328c204cadd89517a1612ecfda5b2dd9d7a/scipy-1.15.3-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:6db907c7368e3092e24919b5e31c76998b0ce1684d51a90943cb0ed1b4ffd6c1", size = 25152602, upload-time = "2025-05-08T16:05:29.313Z" }, + { url = "https://files.pythonhosted.org/packages/56/c5/1032cdb565f146109212153339f9cb8b993701e9fe56b1c97699eee12586/scipy-1.15.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:721d6b4ef5dc82ca8968c25b111e307083d7ca9091bc38163fb89243e85e3889", size = 35503415, upload-time = "2025-05-08T16:05:34.699Z" }, + { url = "https://files.pythonhosted.org/packages/bd/37/89f19c8c05505d0601ed5650156e50eb881ae3918786c8fd7262b4ee66d3/scipy-1.15.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:39cb9c62e471b1bb3750066ecc3a3f3052b37751c7c3dfd0fd7e48900ed52982", size = 37652622, upload-time = "2025-05-08T16:05:40.762Z" }, + { url = "https://files.pythonhosted.org/packages/7e/31/be59513aa9695519b18e1851bb9e487de66f2d31f835201f1b42f5d4d475/scipy-1.15.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:795c46999bae845966368a3c013e0e00947932d68e235702b5c3f6ea799aa8c9", size = 37244796, upload-time = "2025-05-08T16:05:48.119Z" }, + { url = "https://files.pythonhosted.org/packages/10/c0/4f5f3eeccc235632aab79b27a74a9130c6c35df358129f7ac8b29f562ac7/scipy-1.15.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:18aaacb735ab38b38db42cb01f6b92a2d0d4b6aabefeb07f02849e47f8fb3594", size = 40047684, upload-time = "2025-05-08T16:05:54.22Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a7/0ddaf514ce8a8714f6ed243a2b391b41dbb65251affe21ee3077ec45ea9a/scipy-1.15.3-cp311-cp311-win_amd64.whl", hash = "sha256:ae48a786a28412d744c62fd7816a4118ef97e5be0bee968ce8f0a2fba7acf3bb", size = 41246504, upload-time = "2025-05-08T16:06:00.437Z" }, + { url = "https://files.pythonhosted.org/packages/37/4b/683aa044c4162e10ed7a7ea30527f2cbd92e6999c10a8ed8edb253836e9c/scipy-1.15.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6ac6310fdbfb7aa6612408bd2f07295bcbd3fda00d2d702178434751fe48e019", size = 38766735, upload-time = "2025-05-08T16:06:06.471Z" }, + { url = "https://files.pythonhosted.org/packages/7b/7e/f30be3d03de07f25dc0ec926d1681fed5c732d759ac8f51079708c79e680/scipy-1.15.3-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:185cd3d6d05ca4b44a8f1595af87f9c372bb6acf9c808e99aa3e9aa03bd98cf6", size = 30173284, upload-time = "2025-05-08T16:06:11.686Z" }, + { url = "https://files.pythonhosted.org/packages/07/9c/0ddb0d0abdabe0d181c1793db51f02cd59e4901da6f9f7848e1f96759f0d/scipy-1.15.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:05dc6abcd105e1a29f95eada46d4a3f251743cfd7d3ae8ddb4088047f24ea477", size = 22446958, upload-time = "2025-05-08T16:06:15.97Z" }, + { url = "https://files.pythonhosted.org/packages/af/43/0bce905a965f36c58ff80d8bea33f1f9351b05fad4beaad4eae34699b7a1/scipy-1.15.3-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:06efcba926324df1696931a57a176c80848ccd67ce6ad020c810736bfd58eb1c", size = 25242454, upload-time = "2025-05-08T16:06:20.394Z" }, + { url = "https://files.pythonhosted.org/packages/56/30/a6f08f84ee5b7b28b4c597aca4cbe545535c39fe911845a96414700b64ba/scipy-1.15.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c05045d8b9bfd807ee1b9f38761993297b10b245f012b11b13b91ba8945f7e45", size = 35210199, upload-time = "2025-05-08T16:06:26.159Z" }, + { url = "https://files.pythonhosted.org/packages/0b/1f/03f52c282437a168ee2c7c14a1a0d0781a9a4a8962d84ac05c06b4c5b555/scipy-1.15.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:271e3713e645149ea5ea3e97b57fdab61ce61333f97cfae392c28ba786f9bb49", size = 37309455, upload-time = "2025-05-08T16:06:32.778Z" }, + { url = "https://files.pythonhosted.org/packages/89/b1/fbb53137f42c4bf630b1ffdfc2151a62d1d1b903b249f030d2b1c0280af8/scipy-1.15.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:6cfd56fc1a8e53f6e89ba3a7a7251f7396412d655bca2aa5611c8ec9a6784a1e", size = 36885140, upload-time = "2025-05-08T16:06:39.249Z" }, + { url = "https://files.pythonhosted.org/packages/2e/2e/025e39e339f5090df1ff266d021892694dbb7e63568edcfe43f892fa381d/scipy-1.15.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0ff17c0bb1cb32952c09217d8d1eed9b53d1463e5f1dd6052c7857f83127d539", size = 39710549, upload-time = "2025-05-08T16:06:45.729Z" }, + { url = "https://files.pythonhosted.org/packages/e6/eb/3bf6ea8ab7f1503dca3a10df2e4b9c3f6b3316df07f6c0ded94b281c7101/scipy-1.15.3-cp312-cp312-win_amd64.whl", hash = "sha256:52092bc0472cfd17df49ff17e70624345efece4e1a12b23783a1ac59a1b728ed", size = 40966184, upload-time = "2025-05-08T16:06:52.623Z" }, + { url = "https://files.pythonhosted.org/packages/73/18/ec27848c9baae6e0d6573eda6e01a602e5649ee72c27c3a8aad673ebecfd/scipy-1.15.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:2c620736bcc334782e24d173c0fdbb7590a0a436d2fdf39310a8902505008759", size = 38728256, upload-time = "2025-05-08T16:06:58.696Z" }, + { url = "https://files.pythonhosted.org/packages/74/cd/1aef2184948728b4b6e21267d53b3339762c285a46a274ebb7863c9e4742/scipy-1.15.3-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:7e11270a000969409d37ed399585ee530b9ef6aa99d50c019de4cb01e8e54e62", size = 30109540, upload-time = "2025-05-08T16:07:04.209Z" }, + { url = "https://files.pythonhosted.org/packages/5b/d8/59e452c0a255ec352bd0a833537a3bc1bfb679944c4938ab375b0a6b3a3e/scipy-1.15.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:8c9ed3ba2c8a2ce098163a9bdb26f891746d02136995df25227a20e71c396ebb", size = 22383115, upload-time = "2025-05-08T16:07:08.998Z" }, + { url = "https://files.pythonhosted.org/packages/08/f5/456f56bbbfccf696263b47095291040655e3cbaf05d063bdc7c7517f32ac/scipy-1.15.3-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:0bdd905264c0c9cfa74a4772cdb2070171790381a5c4d312c973382fc6eaf730", size = 25163884, upload-time = "2025-05-08T16:07:14.091Z" }, + { url = "https://files.pythonhosted.org/packages/a2/66/a9618b6a435a0f0c0b8a6d0a2efb32d4ec5a85f023c2b79d39512040355b/scipy-1.15.3-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:79167bba085c31f38603e11a267d862957cbb3ce018d8b38f79ac043bc92d825", size = 35174018, upload-time = "2025-05-08T16:07:19.427Z" }, + { url = "https://files.pythonhosted.org/packages/b5/09/c5b6734a50ad4882432b6bb7c02baf757f5b2f256041da5df242e2d7e6b6/scipy-1.15.3-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c9deabd6d547aee2c9a81dee6cc96c6d7e9a9b1953f74850c179f91fdc729cb7", size = 37269716, upload-time = "2025-05-08T16:07:25.712Z" }, + { url = "https://files.pythonhosted.org/packages/77/0a/eac00ff741f23bcabd352731ed9b8995a0a60ef57f5fd788d611d43d69a1/scipy-1.15.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:dde4fc32993071ac0c7dd2d82569e544f0bdaff66269cb475e0f369adad13f11", size = 36872342, upload-time = "2025-05-08T16:07:31.468Z" }, + { url = "https://files.pythonhosted.org/packages/fe/54/4379be86dd74b6ad81551689107360d9a3e18f24d20767a2d5b9253a3f0a/scipy-1.15.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:f77f853d584e72e874d87357ad70f44b437331507d1c311457bed8ed2b956126", size = 39670869, upload-time = "2025-05-08T16:07:38.002Z" }, + { url = "https://files.pythonhosted.org/packages/87/2e/892ad2862ba54f084ffe8cc4a22667eaf9c2bcec6d2bff1d15713c6c0703/scipy-1.15.3-cp313-cp313-win_amd64.whl", hash = "sha256:b90ab29d0c37ec9bf55424c064312930ca5f4bde15ee8619ee44e69319aab163", size = 40988851, upload-time = "2025-05-08T16:08:33.671Z" }, + { url = "https://files.pythonhosted.org/packages/1b/e9/7a879c137f7e55b30d75d90ce3eb468197646bc7b443ac036ae3fe109055/scipy-1.15.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:3ac07623267feb3ae308487c260ac684b32ea35fd81e12845039952f558047b8", size = 38863011, upload-time = "2025-05-08T16:07:44.039Z" }, + { url = "https://files.pythonhosted.org/packages/51/d1/226a806bbd69f62ce5ef5f3ffadc35286e9fbc802f606a07eb83bf2359de/scipy-1.15.3-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:6487aa99c2a3d509a5227d9a5e889ff05830a06b2ce08ec30df6d79db5fcd5c5", size = 30266407, upload-time = "2025-05-08T16:07:49.891Z" }, + { url = "https://files.pythonhosted.org/packages/e5/9b/f32d1d6093ab9eeabbd839b0f7619c62e46cc4b7b6dbf05b6e615bbd4400/scipy-1.15.3-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:50f9e62461c95d933d5c5ef4a1f2ebf9a2b4e83b0db374cb3f1de104d935922e", size = 22540030, upload-time = "2025-05-08T16:07:54.121Z" }, + { url = "https://files.pythonhosted.org/packages/e7/29/c278f699b095c1a884f29fda126340fcc201461ee8bfea5c8bdb1c7c958b/scipy-1.15.3-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:14ed70039d182f411ffc74789a16df3835e05dc469b898233a245cdfd7f162cb", size = 25218709, upload-time = "2025-05-08T16:07:58.506Z" }, + { url = "https://files.pythonhosted.org/packages/24/18/9e5374b617aba742a990581373cd6b68a2945d65cc588482749ef2e64467/scipy-1.15.3-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0a769105537aa07a69468a0eefcd121be52006db61cdd8cac8a0e68980bbb723", size = 34809045, upload-time = "2025-05-08T16:08:03.929Z" }, + { url = "https://files.pythonhosted.org/packages/e1/fe/9c4361e7ba2927074360856db6135ef4904d505e9b3afbbcb073c4008328/scipy-1.15.3-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9db984639887e3dffb3928d118145ffe40eff2fa40cb241a306ec57c219ebbbb", size = 36703062, upload-time = "2025-05-08T16:08:09.558Z" }, + { url = "https://files.pythonhosted.org/packages/b7/8e/038ccfe29d272b30086b25a4960f757f97122cb2ec42e62b460d02fe98e9/scipy-1.15.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:40e54d5c7e7ebf1aa596c374c49fa3135f04648a0caabcb66c52884b943f02b4", size = 36393132, upload-time = "2025-05-08T16:08:15.34Z" }, + { url = "https://files.pythonhosted.org/packages/10/7e/5c12285452970be5bdbe8352c619250b97ebf7917d7a9a9e96b8a8140f17/scipy-1.15.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:5e721fed53187e71d0ccf382b6bf977644c533e506c4d33c3fb24de89f5c3ed5", size = 38979503, upload-time = "2025-05-08T16:08:21.513Z" }, + { url = "https://files.pythonhosted.org/packages/81/06/0a5e5349474e1cbc5757975b21bd4fad0e72ebf138c5592f191646154e06/scipy-1.15.3-cp313-cp313t-win_amd64.whl", hash = "sha256:76ad1fb5f8752eabf0fa02e4cc0336b4e8f021e2d5f061ed37d6d264db35e3ca", size = 40308097, upload-time = "2025-05-08T16:08:27.627Z" }, +] + +[[package]] +name = "scipy" +version = "1.16.3" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0a/ca/d8ace4f98322d01abcd52d381134344bf7b431eba7ed8b42bdea5a3c2ac9/scipy-1.16.3.tar.gz", hash = "sha256:01e87659402762f43bd2fee13370553a17ada367d42e7487800bf2916535aecb", size = 30597883, upload-time = "2025-10-28T17:38:54.068Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/5f/6f37d7439de1455ce9c5a556b8d1db0979f03a796c030bafdf08d35b7bf9/scipy-1.16.3-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:40be6cf99e68b6c4321e9f8782e7d5ff8265af28ef2cd56e9c9b2638fa08ad97", size = 36630881, upload-time = "2025-10-28T17:31:47.104Z" }, + { url = "https://files.pythonhosted.org/packages/7c/89/d70e9f628749b7e4db2aa4cd89735502ff3f08f7b9b27d2e799485987cd9/scipy-1.16.3-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:8be1ca9170fcb6223cc7c27f4305d680ded114a1567c0bd2bfcbf947d1b17511", size = 28941012, upload-time = "2025-10-28T17:31:53.411Z" }, + { url = "https://files.pythonhosted.org/packages/a8/a8/0e7a9a6872a923505dbdf6bb93451edcac120363131c19013044a1e7cb0c/scipy-1.16.3-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:bea0a62734d20d67608660f69dcda23e7f90fb4ca20974ab80b6ed40df87a005", size = 20931935, upload-time = "2025-10-28T17:31:57.361Z" }, + { url = "https://files.pythonhosted.org/packages/bd/c7/020fb72bd79ad798e4dbe53938543ecb96b3a9ac3fe274b7189e23e27353/scipy-1.16.3-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:2a207a6ce9c24f1951241f4693ede2d393f59c07abc159b2cb2be980820e01fb", size = 23534466, upload-time = "2025-10-28T17:32:01.875Z" }, + { url = "https://files.pythonhosted.org/packages/be/a0/668c4609ce6dbf2f948e167836ccaf897f95fb63fa231c87da7558a374cd/scipy-1.16.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:532fb5ad6a87e9e9cd9c959b106b73145a03f04c7d57ea3e6f6bb60b86ab0876", size = 33593618, upload-time = "2025-10-28T17:32:06.902Z" }, + { url = "https://files.pythonhosted.org/packages/ca/6e/8942461cf2636cdae083e3eb72622a7fbbfa5cf559c7d13ab250a5dbdc01/scipy-1.16.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0151a0749efeaaab78711c78422d413c583b8cdd2011a3c1d6c794938ee9fdb2", size = 35899798, upload-time = "2025-10-28T17:32:12.665Z" }, + { url = "https://files.pythonhosted.org/packages/79/e8/d0f33590364cdbd67f28ce79368b373889faa4ee959588beddf6daef9abe/scipy-1.16.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:b7180967113560cca57418a7bc719e30366b47959dd845a93206fbed693c867e", size = 36226154, upload-time = "2025-10-28T17:32:17.961Z" }, + { url = "https://files.pythonhosted.org/packages/39/c1/1903de608c0c924a1749c590064e65810f8046e437aba6be365abc4f7557/scipy-1.16.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:deb3841c925eeddb6afc1e4e4a45e418d19ec7b87c5df177695224078e8ec733", size = 38878540, upload-time = "2025-10-28T17:32:23.907Z" }, + { url = "https://files.pythonhosted.org/packages/f1/d0/22ec7036ba0b0a35bccb7f25ab407382ed34af0b111475eb301c16f8a2e5/scipy-1.16.3-cp311-cp311-win_amd64.whl", hash = "sha256:53c3844d527213631e886621df5695d35e4f6a75f620dca412bcd292f6b87d78", size = 38722107, upload-time = "2025-10-28T17:32:29.921Z" }, + { url = "https://files.pythonhosted.org/packages/7b/60/8a00e5a524bb3bf8898db1650d350f50e6cffb9d7a491c561dc9826c7515/scipy-1.16.3-cp311-cp311-win_arm64.whl", hash = "sha256:9452781bd879b14b6f055b26643703551320aa8d79ae064a71df55c00286a184", size = 25506272, upload-time = "2025-10-28T17:32:34.577Z" }, + { url = "https://files.pythonhosted.org/packages/40/41/5bf55c3f386b1643812f3a5674edf74b26184378ef0f3e7c7a09a7e2ca7f/scipy-1.16.3-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:81fc5827606858cf71446a5e98715ba0e11f0dbc83d71c7409d05486592a45d6", size = 36659043, upload-time = "2025-10-28T17:32:40.285Z" }, + { url = "https://files.pythonhosted.org/packages/1e/0f/65582071948cfc45d43e9870bf7ca5f0e0684e165d7c9ef4e50d783073eb/scipy-1.16.3-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:c97176013d404c7346bf57874eaac5187d969293bf40497140b0a2b2b7482e07", size = 28898986, upload-time = "2025-10-28T17:32:45.325Z" }, + { url = "https://files.pythonhosted.org/packages/96/5e/36bf3f0ac298187d1ceadde9051177d6a4fe4d507e8f59067dc9dd39e650/scipy-1.16.3-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2b71d93c8a9936046866acebc915e2af2e292b883ed6e2cbe5c34beb094b82d9", size = 20889814, upload-time = "2025-10-28T17:32:49.277Z" }, + { url = "https://files.pythonhosted.org/packages/80/35/178d9d0c35394d5d5211bbff7ac4f2986c5488b59506fef9e1de13ea28d3/scipy-1.16.3-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:3d4a07a8e785d80289dfe66b7c27d8634a773020742ec7187b85ccc4b0e7b686", size = 23565795, upload-time = "2025-10-28T17:32:53.337Z" }, + { url = "https://files.pythonhosted.org/packages/fa/46/d1146ff536d034d02f83c8afc3c4bab2eddb634624d6529a8512f3afc9da/scipy-1.16.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0553371015692a898e1aa858fed67a3576c34edefa6b7ebdb4e9dde49ce5c203", size = 33349476, upload-time = "2025-10-28T17:32:58.353Z" }, + { url = "https://files.pythonhosted.org/packages/79/2e/415119c9ab3e62249e18c2b082c07aff907a273741b3f8160414b0e9193c/scipy-1.16.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:72d1717fd3b5e6ec747327ce9bda32d5463f472c9dce9f54499e81fbd50245a1", size = 35676692, upload-time = "2025-10-28T17:33:03.88Z" }, + { url = "https://files.pythonhosted.org/packages/27/82/df26e44da78bf8d2aeaf7566082260cfa15955a5a6e96e6a29935b64132f/scipy-1.16.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:1fb2472e72e24d1530debe6ae078db70fb1605350c88a3d14bc401d6306dbffe", size = 36019345, upload-time = "2025-10-28T17:33:09.773Z" }, + { url = "https://files.pythonhosted.org/packages/82/31/006cbb4b648ba379a95c87262c2855cd0d09453e500937f78b30f02fa1cd/scipy-1.16.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c5192722cffe15f9329a3948c4b1db789fbb1f05c97899187dcf009b283aea70", size = 38678975, upload-time = "2025-10-28T17:33:15.809Z" }, + { url = "https://files.pythonhosted.org/packages/c2/7f/acbd28c97e990b421af7d6d6cd416358c9c293fc958b8529e0bd5d2a2a19/scipy-1.16.3-cp312-cp312-win_amd64.whl", hash = "sha256:56edc65510d1331dae01ef9b658d428e33ed48b4f77b1d51caf479a0253f96dc", size = 38555926, upload-time = "2025-10-28T17:33:21.388Z" }, + { url = "https://files.pythonhosted.org/packages/ce/69/c5c7807fd007dad4f48e0a5f2153038dc96e8725d3345b9ee31b2b7bed46/scipy-1.16.3-cp312-cp312-win_arm64.whl", hash = "sha256:a8a26c78ef223d3e30920ef759e25625a0ecdd0d60e5a8818b7513c3e5384cf2", size = 25463014, upload-time = "2025-10-28T17:33:25.975Z" }, + { url = "https://files.pythonhosted.org/packages/72/f1/57e8327ab1508272029e27eeef34f2302ffc156b69e7e233e906c2a5c379/scipy-1.16.3-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:d2ec56337675e61b312179a1ad124f5f570c00f920cc75e1000025451b88241c", size = 36617856, upload-time = "2025-10-28T17:33:31.375Z" }, + { url = "https://files.pythonhosted.org/packages/44/13/7e63cfba8a7452eb756306aa2fd9b37a29a323b672b964b4fdeded9a3f21/scipy-1.16.3-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:16b8bc35a4cc24db80a0ec836a9286d0e31b2503cb2fd7ff7fb0e0374a97081d", size = 28874306, upload-time = "2025-10-28T17:33:36.516Z" }, + { url = "https://files.pythonhosted.org/packages/15/65/3a9400efd0228a176e6ec3454b1fa998fbbb5a8defa1672c3f65706987db/scipy-1.16.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:5803c5fadd29de0cf27fa08ccbfe7a9e5d741bf63e4ab1085437266f12460ff9", size = 20865371, upload-time = "2025-10-28T17:33:42.094Z" }, + { url = "https://files.pythonhosted.org/packages/33/d7/eda09adf009a9fb81827194d4dd02d2e4bc752cef16737cc4ef065234031/scipy-1.16.3-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:b81c27fc41954319a943d43b20e07c40bdcd3ff7cf013f4fb86286faefe546c4", size = 23524877, upload-time = "2025-10-28T17:33:48.483Z" }, + { url = "https://files.pythonhosted.org/packages/7d/6b/3f911e1ebc364cb81320223a3422aab7d26c9c7973109a9cd0f27c64c6c0/scipy-1.16.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:0c3b4dd3d9b08dbce0f3440032c52e9e2ab9f96ade2d3943313dfe51a7056959", size = 33342103, upload-time = "2025-10-28T17:33:56.495Z" }, + { url = "https://files.pythonhosted.org/packages/21/f6/4bfb5695d8941e5c570a04d9fcd0d36bce7511b7d78e6e75c8f9791f82d0/scipy-1.16.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7dc1360c06535ea6116a2220f760ae572db9f661aba2d88074fe30ec2aa1ff88", size = 35697297, upload-time = "2025-10-28T17:34:04.722Z" }, + { url = "https://files.pythonhosted.org/packages/04/e1/6496dadbc80d8d896ff72511ecfe2316b50313bfc3ebf07a3f580f08bd8c/scipy-1.16.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:663b8d66a8748051c3ee9c96465fb417509315b99c71550fda2591d7dd634234", size = 36021756, upload-time = "2025-10-28T17:34:13.482Z" }, + { url = "https://files.pythonhosted.org/packages/fe/bd/a8c7799e0136b987bda3e1b23d155bcb31aec68a4a472554df5f0937eef7/scipy-1.16.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:eab43fae33a0c39006a88096cd7b4f4ef545ea0447d250d5ac18202d40b6611d", size = 38696566, upload-time = "2025-10-28T17:34:22.384Z" }, + { url = "https://files.pythonhosted.org/packages/cd/01/1204382461fcbfeb05b6161b594f4007e78b6eba9b375382f79153172b4d/scipy-1.16.3-cp313-cp313-win_amd64.whl", hash = "sha256:062246acacbe9f8210de8e751b16fc37458213f124bef161a5a02c7a39284304", size = 38529877, upload-time = "2025-10-28T17:35:51.076Z" }, + { url = "https://files.pythonhosted.org/packages/7f/14/9d9fbcaa1260a94f4bb5b64ba9213ceb5d03cd88841fe9fd1ffd47a45b73/scipy-1.16.3-cp313-cp313-win_arm64.whl", hash = "sha256:50a3dbf286dbc7d84f176f9a1574c705f277cb6565069f88f60db9eafdbe3ee2", size = 25455366, upload-time = "2025-10-28T17:35:59.014Z" }, + { url = "https://files.pythonhosted.org/packages/e2/a3/9ec205bd49f42d45d77f1730dbad9ccf146244c1647605cf834b3a8c4f36/scipy-1.16.3-cp313-cp313t-macosx_10_14_x86_64.whl", hash = "sha256:fb4b29f4cf8cc5a8d628bc8d8e26d12d7278cd1f219f22698a378c3d67db5e4b", size = 37027931, upload-time = "2025-10-28T17:34:31.451Z" }, + { url = "https://files.pythonhosted.org/packages/25/06/ca9fd1f3a4589cbd825b1447e5db3a8ebb969c1eaf22c8579bd286f51b6d/scipy-1.16.3-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:8d09d72dc92742988b0e7750bddb8060b0c7079606c0d24a8cc8e9c9c11f9079", size = 29400081, upload-time = "2025-10-28T17:34:39.087Z" }, + { url = "https://files.pythonhosted.org/packages/6a/56/933e68210d92657d93fb0e381683bc0e53a965048d7358ff5fbf9e6a1b17/scipy-1.16.3-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:03192a35e661470197556de24e7cb1330d84b35b94ead65c46ad6f16f6b28f2a", size = 21391244, upload-time = "2025-10-28T17:34:45.234Z" }, + { url = "https://files.pythonhosted.org/packages/a8/7e/779845db03dc1418e215726329674b40576879b91814568757ff0014ad65/scipy-1.16.3-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:57d01cb6f85e34f0946b33caa66e892aae072b64b034183f3d87c4025802a119", size = 23929753, upload-time = "2025-10-28T17:34:51.793Z" }, + { url = "https://files.pythonhosted.org/packages/4c/4b/f756cf8161d5365dcdef9e5f460ab226c068211030a175d2fc7f3f41ca64/scipy-1.16.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:96491a6a54e995f00a28a3c3badfff58fd093bf26cd5fb34a2188c8c756a3a2c", size = 33496912, upload-time = "2025-10-28T17:34:59.8Z" }, + { url = "https://files.pythonhosted.org/packages/09/b5/222b1e49a58668f23839ca1542a6322bb095ab8d6590d4f71723869a6c2c/scipy-1.16.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cd13e354df9938598af2be05822c323e97132d5e6306b83a3b4ee6724c6e522e", size = 35802371, upload-time = "2025-10-28T17:35:08.173Z" }, + { url = "https://files.pythonhosted.org/packages/c1/8d/5964ef68bb31829bde27611f8c9deeac13764589fe74a75390242b64ca44/scipy-1.16.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:63d3cdacb8a824a295191a723ee5e4ea7768ca5ca5f2838532d9f2e2b3ce2135", size = 36190477, upload-time = "2025-10-28T17:35:16.7Z" }, + { url = "https://files.pythonhosted.org/packages/ab/f2/b31d75cb9b5fa4dd39a0a931ee9b33e7f6f36f23be5ef560bf72e0f92f32/scipy-1.16.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:e7efa2681ea410b10dde31a52b18b0154d66f2485328830e45fdf183af5aefc6", size = 38796678, upload-time = "2025-10-28T17:35:26.354Z" }, + { url = "https://files.pythonhosted.org/packages/b4/1e/b3723d8ff64ab548c38d87055483714fefe6ee20e0189b62352b5e015bb1/scipy-1.16.3-cp313-cp313t-win_amd64.whl", hash = "sha256:2d1ae2cf0c350e7705168ff2429962a89ad90c2d49d1dd300686d8b2a5af22fc", size = 38640178, upload-time = "2025-10-28T17:35:35.304Z" }, + { url = "https://files.pythonhosted.org/packages/8e/f3/d854ff38789aca9b0cc23008d607ced9de4f7ab14fa1ca4329f86b3758ca/scipy-1.16.3-cp313-cp313t-win_arm64.whl", hash = "sha256:0c623a54f7b79dd88ef56da19bc2873afec9673a48f3b85b18e4d402bdd29a5a", size = 25803246, upload-time = "2025-10-28T17:35:42.155Z" }, + { url = "https://files.pythonhosted.org/packages/99/f6/99b10fd70f2d864c1e29a28bbcaa0c6340f9d8518396542d9ea3b4aaae15/scipy-1.16.3-cp314-cp314-macosx_10_14_x86_64.whl", hash = "sha256:875555ce62743e1d54f06cdf22c1e0bc47b91130ac40fe5d783b6dfa114beeb6", size = 36606469, upload-time = "2025-10-28T17:36:08.741Z" }, + { url = "https://files.pythonhosted.org/packages/4d/74/043b54f2319f48ea940dd025779fa28ee360e6b95acb7cd188fad4391c6b/scipy-1.16.3-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:bb61878c18a470021fb515a843dc7a76961a8daceaaaa8bad1332f1bf4b54657", size = 28872043, upload-time = "2025-10-28T17:36:16.599Z" }, + { url = "https://files.pythonhosted.org/packages/4d/e1/24b7e50cc1c4ee6ffbcb1f27fe9f4c8b40e7911675f6d2d20955f41c6348/scipy-1.16.3-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f2622206f5559784fa5c4b53a950c3c7c1cf3e84ca1b9c4b6c03f062f289ca26", size = 20862952, upload-time = "2025-10-28T17:36:22.966Z" }, + { url = "https://files.pythonhosted.org/packages/dd/3a/3e8c01a4d742b730df368e063787c6808597ccb38636ed821d10b39ca51b/scipy-1.16.3-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:7f68154688c515cdb541a31ef8eb66d8cd1050605be9dcd74199cbd22ac739bc", size = 23508512, upload-time = "2025-10-28T17:36:29.731Z" }, + { url = "https://files.pythonhosted.org/packages/1f/60/c45a12b98ad591536bfe5330cb3cfe1850d7570259303563b1721564d458/scipy-1.16.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8b3c820ddb80029fe9f43d61b81d8b488d3ef8ca010d15122b152db77dc94c22", size = 33413639, upload-time = "2025-10-28T17:36:37.982Z" }, + { url = "https://files.pythonhosted.org/packages/71/bc/35957d88645476307e4839712642896689df442f3e53b0fa016ecf8a3357/scipy-1.16.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d3837938ae715fc0fe3c39c0202de3a8853aff22ca66781ddc2ade7554b7e2cc", size = 35704729, upload-time = "2025-10-28T17:36:46.547Z" }, + { url = "https://files.pythonhosted.org/packages/3b/15/89105e659041b1ca11c386e9995aefacd513a78493656e57789f9d9eab61/scipy-1.16.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:aadd23f98f9cb069b3bd64ddc900c4d277778242e961751f77a8cb5c4b946fb0", size = 36086251, upload-time = "2025-10-28T17:36:55.161Z" }, + { url = "https://files.pythonhosted.org/packages/1a/87/c0ea673ac9c6cc50b3da2196d860273bc7389aa69b64efa8493bdd25b093/scipy-1.16.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b7c5f1bda1354d6a19bc6af73a649f8285ca63ac6b52e64e658a5a11d4d69800", size = 38716681, upload-time = "2025-10-28T17:37:04.1Z" }, + { url = "https://files.pythonhosted.org/packages/91/06/837893227b043fb9b0d13e4bd7586982d8136cb249ffb3492930dab905b8/scipy-1.16.3-cp314-cp314-win_amd64.whl", hash = "sha256:e5d42a9472e7579e473879a1990327830493a7047506d58d73fc429b84c1d49d", size = 39358423, upload-time = "2025-10-28T17:38:20.005Z" }, + { url = "https://files.pythonhosted.org/packages/95/03/28bce0355e4d34a7c034727505a02d19548549e190bedd13a721e35380b7/scipy-1.16.3-cp314-cp314-win_arm64.whl", hash = "sha256:6020470b9d00245926f2d5bb93b119ca0340f0d564eb6fbaad843eaebf9d690f", size = 26135027, upload-time = "2025-10-28T17:38:24.966Z" }, + { url = "https://files.pythonhosted.org/packages/b2/6f/69f1e2b682efe9de8fe9f91040f0cd32f13cfccba690512ba4c582b0bc29/scipy-1.16.3-cp314-cp314t-macosx_10_14_x86_64.whl", hash = "sha256:e1d27cbcb4602680a49d787d90664fa4974063ac9d4134813332a8c53dbe667c", size = 37028379, upload-time = "2025-10-28T17:37:14.061Z" }, + { url = "https://files.pythonhosted.org/packages/7c/2d/e826f31624a5ebbab1cd93d30fd74349914753076ed0593e1d56a98c4fb4/scipy-1.16.3-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:9b9c9c07b6d56a35777a1b4cc8966118fb16cfd8daf6743867d17d36cfad2d40", size = 29400052, upload-time = "2025-10-28T17:37:21.709Z" }, + { url = "https://files.pythonhosted.org/packages/69/27/d24feb80155f41fd1f156bf144e7e049b4e2b9dd06261a242905e3bc7a03/scipy-1.16.3-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:3a4c460301fb2cffb7f88528f30b3127742cff583603aa7dc964a52c463b385d", size = 21391183, upload-time = "2025-10-28T17:37:29.559Z" }, + { url = "https://files.pythonhosted.org/packages/f8/d3/1b229e433074c5738a24277eca520a2319aac7465eea7310ea6ae0e98ae2/scipy-1.16.3-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:f667a4542cc8917af1db06366d3f78a5c8e83badd56409f94d1eac8d8d9133fa", size = 23930174, upload-time = "2025-10-28T17:37:36.306Z" }, + { url = "https://files.pythonhosted.org/packages/16/9d/d9e148b0ec680c0f042581a2be79a28a7ab66c0c4946697f9e7553ead337/scipy-1.16.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:f379b54b77a597aa7ee5e697df0d66903e41b9c85a6dd7946159e356319158e8", size = 33497852, upload-time = "2025-10-28T17:37:42.228Z" }, + { url = "https://files.pythonhosted.org/packages/2f/22/4e5f7561e4f98b7bea63cf3fd7934bff1e3182e9f1626b089a679914d5c8/scipy-1.16.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4aff59800a3b7f786b70bfd6ab551001cb553244988d7d6b8299cb1ea653b353", size = 35798595, upload-time = "2025-10-28T17:37:48.102Z" }, + { url = "https://files.pythonhosted.org/packages/83/42/6644d714c179429fc7196857866f219fef25238319b650bb32dde7bf7a48/scipy-1.16.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:da7763f55885045036fabcebd80144b757d3db06ab0861415d1c3b7c69042146", size = 36186269, upload-time = "2025-10-28T17:37:53.72Z" }, + { url = "https://files.pythonhosted.org/packages/ac/70/64b4d7ca92f9cf2e6fc6aaa2eecf80bb9b6b985043a9583f32f8177ea122/scipy-1.16.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ffa6eea95283b2b8079b821dc11f50a17d0571c92b43e2b5b12764dc5f9b285d", size = 38802779, upload-time = "2025-10-28T17:37:59.393Z" }, + { url = "https://files.pythonhosted.org/packages/61/82/8d0e39f62764cce5ffd5284131e109f07cf8955aef9ab8ed4e3aa5e30539/scipy-1.16.3-cp314-cp314t-win_amd64.whl", hash = "sha256:d9f48cafc7ce94cf9b15c6bffdc443a81a27bf7075cf2dcd5c8b40f85d10c4e7", size = 39471128, upload-time = "2025-10-28T17:38:05.259Z" }, + { url = "https://files.pythonhosted.org/packages/64/47/a494741db7280eae6dc033510c319e34d42dd41b7ac0c7ead39354d1a2b5/scipy-1.16.3-cp314-cp314t-win_arm64.whl", hash = "sha256:21d9d6b197227a12dcbf9633320a4e34c6b0e51c57268df255a0942983bac562", size = 26464127, upload-time = "2025-10-28T17:38:11.34Z" }, +] + [[package]] name = "secretstorage" version = "3.4.0" @@ -2865,6 +3757,70 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/ff/2e2eed29e02c14a5cb6c57f09b2d5b40e65d6cc71f45b52e0be295ccbc2f/secretstorage-3.4.0-py3-none-any.whl", hash = "sha256:0e3b6265c2c63509fb7415717607e4b2c9ab767b7f344a57473b779ca13bd02e", size = 15272, upload-time = "2025-09-09T16:42:12.744Z" }, ] +[[package]] +name = "sentencepiece" +version = "0.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/31/5b7cccb307b485db1a2372d6d2980b0a65d067f8be5ca943a103b4acd5b3/sentencepiece-0.2.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e10fa50bdbaa5e2445dbd387979980d391760faf0ec99a09bd7780ff37eaec44", size = 1942557, upload-time = "2025-08-12T06:59:12.379Z" }, + { url = "https://files.pythonhosted.org/packages/1f/41/0ac923a8e685ad290c5afc8ae55c5844977b8d75076fcc04302b9a324274/sentencepiece-0.2.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2f27ae6deea72efdb6f361750c92f6c21fd0ad087445082770cc34015213c526", size = 1325384, upload-time = "2025-08-12T06:59:14.334Z" }, + { url = "https://files.pythonhosted.org/packages/fc/ef/3751555d67daf9003384978f169d31c775cb5c7baf28633caaf1eb2b2b4d/sentencepiece-0.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:60937c959e6f44159fdd9f56fbdd302501f96114a5ba436829496d5f32d8de3f", size = 1253317, upload-time = "2025-08-12T06:59:16.247Z" }, + { url = "https://files.pythonhosted.org/packages/46/a5/742c69b7bd144eb32b6e5fd50dbd8abbbc7a95fce2fe16e50156fa400e3b/sentencepiece-0.2.1-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d8b1d91545578852f128650b8cce4ec20f93d39b378ff554ebe66290f2dabb92", size = 1316379, upload-time = "2025-08-12T06:59:17.825Z" }, + { url = "https://files.pythonhosted.org/packages/c8/89/8deeafbba2871e8fa10f20f17447786f4ac38085925335728d360eaf4cae/sentencepiece-0.2.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:27e38eee653abc3d387862e67bc5c8b6f428cd604e688b85d29170b7e725c26c", size = 1387926, upload-time = "2025-08-12T06:59:19.395Z" }, + { url = "https://files.pythonhosted.org/packages/c3/ca/67fe73005f0ab617c6a970b199754e28e524b6873aa7025224fad3cda252/sentencepiece-0.2.1-cp310-cp310-win32.whl", hash = "sha256:251874d720ac7f28024a168501f3c7bb15d1802245f6e66de565f18bbb9b5eaa", size = 999550, upload-time = "2025-08-12T06:59:20.844Z" }, + { url = "https://files.pythonhosted.org/packages/6d/33/dc5b54042050d2dda4229c3ce1f862541c99966390b6aa20f54d520d2dc2/sentencepiece-0.2.1-cp310-cp310-win_amd64.whl", hash = "sha256:e52144670738b4b477fade6c2a9b6af71a8d0094514c9853ac9f6fc1fcfabae7", size = 1054613, upload-time = "2025-08-12T06:59:22.255Z" }, + { url = "https://files.pythonhosted.org/packages/fa/19/1ea47f46ff97fe04422b78997da1a37cd632f414aae042d27a9009c5b733/sentencepiece-0.2.1-cp310-cp310-win_arm64.whl", hash = "sha256:9076430ac25dfa7147d9d05751dbc66a04bc1aaac371c07f84952979ea59f0d0", size = 1033884, upload-time = "2025-08-12T06:59:24.194Z" }, + { url = "https://files.pythonhosted.org/packages/d8/15/46afbab00733d81788b64be430ca1b93011bb9388527958e26cc31832de5/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:6356d0986b8b8dc351b943150fcd81a1c6e6e4d439772e8584c64230e58ca987", size = 1942560, upload-time = "2025-08-12T06:59:25.82Z" }, + { url = "https://files.pythonhosted.org/packages/fa/79/7c01b8ef98a0567e9d84a4e7a910f8e7074fcbf398a5cd76f93f4b9316f9/sentencepiece-0.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:8f8ba89a3acb3dc1ae90f65ec1894b0b9596fdb98ab003ff38e058f898b39bc7", size = 1325385, upload-time = "2025-08-12T06:59:27.722Z" }, + { url = "https://files.pythonhosted.org/packages/bb/88/2b41e07bd24f33dcf2f18ec3b74247aa4af3526bad8907b8727ea3caba03/sentencepiece-0.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:02593eca45440ef39247cee8c47322a34bdcc1d8ae83ad28ba5a899a2cf8d79a", size = 1253319, upload-time = "2025-08-12T06:59:29.306Z" }, + { url = "https://files.pythonhosted.org/packages/a0/54/38a1af0c6210a3c6f95aa46d23d6640636d020fba7135cd0d9a84ada05a7/sentencepiece-0.2.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0a0d15781a171d188b661ae4bde1d998c303f6bd8621498c50c671bd45a4798e", size = 1316162, upload-time = "2025-08-12T06:59:30.914Z" }, + { url = "https://files.pythonhosted.org/packages/ef/66/fb191403ade791ad2c3c1e72fe8413e63781b08cfa3aa4c9dfc536d6e795/sentencepiece-0.2.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4f5a3e0d9f445ed9d66c0fec47d4b23d12cfc858b407a03c194c1b26c2ac2a63", size = 1387785, upload-time = "2025-08-12T06:59:32.491Z" }, + { url = "https://files.pythonhosted.org/packages/a9/2d/3bd9b08e70067b2124518b308db6a84a4f8901cc8a4317e2e4288cdd9b4d/sentencepiece-0.2.1-cp311-cp311-win32.whl", hash = "sha256:6d297a1748d429ba8534eebe5535448d78b8acc32d00a29b49acf28102eeb094", size = 999555, upload-time = "2025-08-12T06:59:34.475Z" }, + { url = "https://files.pythonhosted.org/packages/32/b8/f709977f5fda195ae1ea24f24e7c581163b6f142b1005bc3d0bbfe4d7082/sentencepiece-0.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:82d9ead6591015f009cb1be1cb1c015d5e6f04046dbb8c9588b931e869a29728", size = 1054617, upload-time = "2025-08-12T06:59:36.461Z" }, + { url = "https://files.pythonhosted.org/packages/7a/40/a1fc23be23067da0f703709797b464e8a30a1c78cc8a687120cd58d4d509/sentencepiece-0.2.1-cp311-cp311-win_arm64.whl", hash = "sha256:39f8651bd10974eafb9834ce30d9bcf5b73e1fc798a7f7d2528f9820ca86e119", size = 1033877, upload-time = "2025-08-12T06:59:38.391Z" }, + { url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" }, + { url = "https://files.pythonhosted.org/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" }, + { url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" }, + { url = "https://files.pythonhosted.org/packages/89/fa/d3d5ebcba3cb9e6d3775a096251860c41a6bc53a1b9461151df83fe93255/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99f955df238021bf11f0fc37cdb54fd5e5b5f7fd30ecc3d93fb48b6815437167", size = 1316273, upload-time = "2025-08-12T06:59:44.476Z" }, + { url = "https://files.pythonhosted.org/packages/04/88/14f2f4a2b922d8b39be45bf63d79e6cd3a9b2f248b2fcb98a69b12af12f5/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdfecef430d985f1c2bcbfff3defd1d95dae876fbd0173376012d2d7d24044b", size = 1387881, upload-time = "2025-08-12T06:59:46.09Z" }, + { url = "https://files.pythonhosted.org/packages/fd/b8/903e5ccb77b4ef140605d5d71b4f9e0ad95d456d6184688073ed11712809/sentencepiece-0.2.1-cp312-cp312-win32.whl", hash = "sha256:a483fd29a34c3e34c39ac5556b0a90942bec253d260235729e50976f5dba1068", size = 999540, upload-time = "2025-08-12T06:59:48.023Z" }, + { url = "https://files.pythonhosted.org/packages/2d/81/92df5673c067148c2545b1bfe49adfd775bcc3a169a047f5a0e6575ddaca/sentencepiece-0.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:4cdc7c36234fda305e85c32949c5211faaf8dd886096c7cea289ddc12a2d02de", size = 1054671, upload-time = "2025-08-12T06:59:49.895Z" }, + { url = "https://files.pythonhosted.org/packages/fe/02/c5e3bc518655d714622bec87d83db9cdba1cd0619a4a04e2109751c4f47f/sentencepiece-0.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:daeb5e9e9fcad012324807856113708614d534f596d5008638eb9b40112cd9e4", size = 1033923, upload-time = "2025-08-12T06:59:51.952Z" }, + { url = "https://files.pythonhosted.org/packages/ba/4a/85fbe1706d4d04a7e826b53f327c4b80f849cf1c7b7c5e31a20a97d8f28b/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dcd8161eee7b41aae57ded06272905dbd680a0a04b91edd0f64790c796b2f706", size = 1943150, upload-time = "2025-08-12T06:59:53.588Z" }, + { url = "https://files.pythonhosted.org/packages/c2/83/4cfb393e287509fc2155480b9d184706ef8d9fa8cbf5505d02a5792bf220/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c6c8f42949f419ff8c7e9960dbadcfbc982d7b5efc2f6748210d3dd53a7de062", size = 1325651, upload-time = "2025-08-12T06:59:55.073Z" }, + { url = "https://files.pythonhosted.org/packages/8d/de/5a007fb53b1ab0aafc69d11a5a3dd72a289d5a3e78dcf2c3a3d9b14ffe93/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:097f3394e99456e9e4efba1737c3749d7e23563dd1588ce71a3d007f25475fff", size = 1253641, upload-time = "2025-08-12T06:59:56.562Z" }, + { url = "https://files.pythonhosted.org/packages/2c/d2/f552be5928105588f4f4d66ee37dd4c61460d8097e62d0e2e0eec41bc61d/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d7b670879c370d350557edabadbad1f6561a9e6968126e6debca4029e5547820", size = 1316271, upload-time = "2025-08-12T06:59:58.109Z" }, + { url = "https://files.pythonhosted.org/packages/96/df/0cfe748ace5485be740fed9476dee7877f109da32ed0d280312c94ec259f/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c7f0fd2f2693309e6628aeeb2e2faf6edd221134dfccac3308ca0de01f8dab47", size = 1387882, upload-time = "2025-08-12T07:00:00.701Z" }, + { url = "https://files.pythonhosted.org/packages/ac/dd/f7774d42a881ced8e1739f393ab1e82ece39fc9abd4779e28050c2e975b5/sentencepiece-0.2.1-cp313-cp313-win32.whl", hash = "sha256:92b3816aa2339355fda2c8c4e021a5de92180b00aaccaf5e2808972e77a4b22f", size = 999541, upload-time = "2025-08-12T07:00:02.709Z" }, + { url = "https://files.pythonhosted.org/packages/dd/e9/932b9eae6fd7019548321eee1ab8d5e3b3d1294df9d9a0c9ac517c7b636d/sentencepiece-0.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:10ed3dab2044c47f7a2e7b4969b0c430420cdd45735d78c8f853191fa0e3148b", size = 1054669, upload-time = "2025-08-12T07:00:04.915Z" }, + { url = "https://files.pythonhosted.org/packages/c9/3a/76488a00ea7d6931689cda28726a1447d66bf1a4837943489314593d5596/sentencepiece-0.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:ac650534e2251083c5f75dde4ff28896ce7c8904133dc8fef42780f4d5588fcd", size = 1033922, upload-time = "2025-08-12T07:00:06.496Z" }, + { url = "https://files.pythonhosted.org/packages/4a/b6/08fe2ce819e02ccb0296f4843e3f195764ce9829cbda61b7513f29b95718/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8dd4b477a7b069648d19363aad0cab9bad2f4e83b2d179be668efa672500dc94", size = 1946052, upload-time = "2025-08-12T07:00:08.136Z" }, + { url = "https://files.pythonhosted.org/packages/ab/d9/1ea0e740591ff4c6fc2b6eb1d7510d02f3fb885093f19b2f3abd1363b402/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0c0f672da370cc490e4c59d89e12289778310a0e71d176c541e4834759e1ae07", size = 1327408, upload-time = "2025-08-12T07:00:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/99/7e/1fb26e8a21613f6200e1ab88824d5d203714162cf2883248b517deb500b7/sentencepiece-0.2.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad8493bea8432dae8d6830365352350f3b4144415a1d09c4c8cb8d30cf3b6c3c", size = 1254857, upload-time = "2025-08-12T07:00:11.021Z" }, + { url = "https://files.pythonhosted.org/packages/bc/85/c72fd1f3c7a6010544d6ae07f8ddb38b5e2a7e33bd4318f87266c0bbafbf/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b81a24733726e3678d2db63619acc5a8dccd074f7aa7a54ecd5ca33ca6d2d596", size = 1315722, upload-time = "2025-08-12T07:00:12.989Z" }, + { url = "https://files.pythonhosted.org/packages/4a/e8/661e5bd82a8aa641fd6c1020bd0e890ef73230a2b7215ddf9c8cd8e941c2/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0a81799d0a68d618e89063fb423c3001a034c893069135ffe51fee439ae474d6", size = 1387452, upload-time = "2025-08-12T07:00:15.088Z" }, + { url = "https://files.pythonhosted.org/packages/99/5e/ae66c361023a470afcbc1fbb8da722c72ea678a2fcd9a18f1a12598c7501/sentencepiece-0.2.1-cp313-cp313t-win32.whl", hash = "sha256:89a3ea015517c42c0341d0d962f3e6aaf2cf10d71b1932d475c44ba48d00aa2b", size = 1002501, upload-time = "2025-08-12T07:00:16.966Z" }, + { url = "https://files.pythonhosted.org/packages/c1/03/d332828c4ff764e16c1b56c2c8f9a33488bbe796b53fb6b9c4205ddbf167/sentencepiece-0.2.1-cp313-cp313t-win_amd64.whl", hash = "sha256:33f068c9382dc2e7c228eedfd8163b52baa86bb92f50d0488bf2b7da7032e484", size = 1057555, upload-time = "2025-08-12T07:00:18.573Z" }, + { url = "https://files.pythonhosted.org/packages/88/14/5aee0bf0864df9bd82bd59e7711362908e4935e3f9cdc1f57246b5d5c9b9/sentencepiece-0.2.1-cp313-cp313t-win_arm64.whl", hash = "sha256:b3616ad246f360e52c85781e47682d31abfb6554c779e42b65333d4b5f44ecc0", size = 1036042, upload-time = "2025-08-12T07:00:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/24/9c/89eb8b2052f720a612478baf11c8227dcf1dc28cd4ea4c0c19506b5af2a2/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:5d0350b686c320068702116276cfb26c066dc7e65cfef173980b11bb4d606719", size = 1943147, upload-time = "2025-08-12T07:00:21.809Z" }, + { url = "https://files.pythonhosted.org/packages/82/0b/a1432bc87f97c2ace36386ca23e8bd3b91fb40581b5e6148d24b24186419/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c7f54a31cde6fa5cb030370566f68152a742f433f8d2be458463d06c208aef33", size = 1325624, upload-time = "2025-08-12T07:00:23.289Z" }, + { url = "https://files.pythonhosted.org/packages/ea/99/bbe054ebb5a5039457c590e0a4156ed073fb0fe9ce4f7523404dd5b37463/sentencepiece-0.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c83b85ab2d6576607f31df77ff86f28182be4a8de6d175d2c33ca609925f5da1", size = 1253670, upload-time = "2025-08-12T07:00:24.69Z" }, + { url = "https://files.pythonhosted.org/packages/19/ad/d5c7075f701bd97971d7c2ac2904f227566f51ef0838dfbdfdccb58cd212/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1855f57db07b51fb51ed6c9c452f570624d2b169b36f0f79ef71a6e6c618cd8b", size = 1316247, upload-time = "2025-08-12T07:00:26.435Z" }, + { url = "https://files.pythonhosted.org/packages/fb/03/35fbe5f3d9a7435eebd0b473e09584bd3cc354ce118b960445b060d33781/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01e6912125cb45d3792f530a4d38f8e21bf884d6b4d4ade1b2de5cf7a8d2a52b", size = 1387894, upload-time = "2025-08-12T07:00:28.339Z" }, + { url = "https://files.pythonhosted.org/packages/dc/aa/956ef729aafb6c8f9c443104c9636489093bb5c61d6b90fc27aa1a865574/sentencepiece-0.2.1-cp314-cp314-win32.whl", hash = "sha256:c415c9de1447e0a74ae3fdb2e52f967cb544113a3a5ce3a194df185cbc1f962f", size = 1096698, upload-time = "2025-08-12T07:00:29.764Z" }, + { url = "https://files.pythonhosted.org/packages/b8/cb/fe400d8836952cc535c81a0ce47dc6875160e5fedb71d2d9ff0e9894c2a6/sentencepiece-0.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:881b2e44b14fc19feade3cbed314be37de639fc415375cefaa5bc81a4be137fd", size = 1155115, upload-time = "2025-08-12T07:00:32.865Z" }, + { url = "https://files.pythonhosted.org/packages/32/89/047921cf70f36c7b6b6390876b2399b3633ab73b8d0cb857e5a964238941/sentencepiece-0.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:2005242a16d2dc3ac5fe18aa7667549134d37854823df4c4db244752453b78a8", size = 1133890, upload-time = "2025-08-12T07:00:34.763Z" }, + { url = "https://files.pythonhosted.org/packages/a1/11/5b414b9fae6255b5fb1e22e2ed3dc3a72d3a694e5703910e640ac78346bb/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a19adcec27c524cb7069a1c741060add95f942d1cbf7ad0d104dffa0a7d28a2b", size = 1946081, upload-time = "2025-08-12T07:00:36.97Z" }, + { url = "https://files.pythonhosted.org/packages/77/eb/7a5682bb25824db8545f8e5662e7f3e32d72a508fdce086029d89695106b/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:e37e4b4c4a11662b5db521def4e44d4d30ae69a1743241412a93ae40fdcab4bb", size = 1327406, upload-time = "2025-08-12T07:00:38.669Z" }, + { url = "https://files.pythonhosted.org/packages/03/b0/811dae8fb9f2784e138785d481469788f2e0d0c109c5737372454415f55f/sentencepiece-0.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:477c81505db072b3ab627e7eab972ea1025331bd3a92bacbf798df2b75ea86ec", size = 1254846, upload-time = "2025-08-12T07:00:40.611Z" }, + { url = "https://files.pythonhosted.org/packages/ef/23/195b2e7ec85ebb6a547969f60b723c7aca5a75800ece6cc3f41da872d14e/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:010f025a544ef770bb395091d57cb94deb9652d8972e0d09f71d85d5a0816c8c", size = 1315721, upload-time = "2025-08-12T07:00:42.914Z" }, + { url = "https://files.pythonhosted.org/packages/7e/aa/553dbe4178b5f23eb28e59393dddd64186178b56b81d9b8d5c3ff1c28395/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:733e59ff1794d26db706cd41fc2d7ca5f6c64a820709cb801dc0ea31780d64ab", size = 1387458, upload-time = "2025-08-12T07:00:44.56Z" }, + { url = "https://files.pythonhosted.org/packages/66/7c/08ff0012507297a4dd74a5420fdc0eb9e3e80f4e88cab1538d7f28db303d/sentencepiece-0.2.1-cp314-cp314t-win32.whl", hash = "sha256:d3233770f78e637dc8b1fda2cd7c3b99ec77e7505041934188a4e7fe751de3b0", size = 1099765, upload-time = "2025-08-12T07:00:46.058Z" }, + { url = "https://files.pythonhosted.org/packages/91/d5/2a69e1ce15881beb9ddfc7e3f998322f5cedcd5e4d244cb74dade9441663/sentencepiece-0.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:5e4366c97b68218fd30ea72d70c525e6e78a6c0a88650f57ac4c43c63b234a9d", size = 1157807, upload-time = "2025-08-12T07:00:47.673Z" }, + { url = "https://files.pythonhosted.org/packages/f3/16/54f611fcfc2d1c46cbe3ec4169780b2cfa7cf63708ef2b71611136db7513/sentencepiece-0.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:105e36e75cbac1292642045458e8da677b2342dcd33df503e640f0b457cb6751", size = 1136264, upload-time = "2025-08-12T07:00:49.485Z" }, +] + [[package]] name = "setuptools" version = "80.9.0" @@ -2892,6 +3848,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl", hash = "sha256:2f6da418d1f1e0fddd844478f41680e794e6051915791a034ff65e5f100525a2", size = 10235, upload-time = "2024-02-25T23:20:01.196Z" }, ] +[[package]] +name = "sounddevice" +version = "0.5.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cffi" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/4e/4f/28e734898b870db15b6474453f19813d3c81b91c806d9e6f867bd6e4dd03/sounddevice-0.5.3.tar.gz", hash = "sha256:cbac2b60198fbab84533697e7c4904cc895ec69d5fb3973556c9eb74a4629b2c", size = 53465, upload-time = "2025-10-19T13:23:57.922Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/e7/9020e9f0f3df00432728f4c4044387468a743e3d9a4f91123d77be10010e/sounddevice-0.5.3-py3-none-any.whl", hash = "sha256:ea7738baa0a9f9fef7390f649e41c9f2c8ada776180e56c2ffd217133c92a806", size = 32670, upload-time = "2025-10-19T13:23:51.779Z" }, + { url = "https://files.pythonhosted.org/packages/2f/39/714118f8413e0e353436914f2b976665161f1be2b6483ac15a8f61484c14/sounddevice-0.5.3-py3-none-macosx_10_6_x86_64.macosx_10_6_universal2.whl", hash = "sha256:278dc4451fff70934a176df048b77d80d7ce1623a6ec9db8b34b806f3112f9c2", size = 108306, upload-time = "2025-10-19T13:23:53.277Z" }, + { url = "https://files.pythonhosted.org/packages/f5/74/52186e3e5c833d00273f7949a9383adff93692c6e02406bf359cb4d3e921/sounddevice-0.5.3-py3-none-win32.whl", hash = "sha256:845d6927bcf14e84be5292a61ab3359cf8e6b9145819ec6f3ac2619ff089a69c", size = 312882, upload-time = "2025-10-19T13:23:54.829Z" }, + { url = "https://files.pythonhosted.org/packages/66/c7/16123d054aef6d445176c9122bfbe73c11087589b2413cab22aff5a7839a/sounddevice-0.5.3-py3-none-win_amd64.whl", hash = "sha256:f55ad20082efc2bdec06928e974fbcae07bc6c405409ae1334cefe7d377eb687", size = 364025, upload-time = "2025-10-19T13:23:56.362Z" }, +] + [[package]] name = "spoutgl" version = "0.1.1" @@ -2934,6 +3905,54 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "tifffile" +version = "2025.5.10" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and sys_platform == 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/44/d0/18fed0fc0916578a4463f775b0fbd9c5fed2392152d039df2fb533bfdd5d/tifffile-2025.5.10.tar.gz", hash = "sha256:018335d34283aa3fd8c263bae5c3c2b661ebc45548fde31504016fcae7bf1103", size = 365290, upload-time = "2025-05-10T19:22:34.386Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/06/bd0a6097da704a7a7c34a94cfd771c3ea3c2f405dd214e790d22c93f6be1/tifffile-2025.5.10-py3-none-any.whl", hash = "sha256:e37147123c0542d67bc37ba5cdd67e12ea6fbe6e86c52bee037a9eb6a064e5ad", size = 226533, upload-time = "2025-05-10T19:22:27.279Z" }, +] + +[[package]] +name = "tifffile" +version = "2025.12.12" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +] +dependencies = [ + { name = "numpy", marker = "python_full_version >= '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/31/b9/4253513a66f0a836ec3a5104266cf73f7812bfbbcda9d87d8c0e93b28293/tifffile-2025.12.12.tar.gz", hash = "sha256:97e11fd6b1d8dc971896a098c841d9cd4e6eb958ac040dd6fb8b332c3f7288b6", size = 373597, upload-time = "2025-12-13T03:42:53.765Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/5c/e444e1b024a519e488326525f0c154396c6b16baff17e00623f2c21dfc42/tifffile-2025.12.12-py3-none-any.whl", hash = "sha256:e3e3f1290ec6741ca248a5b5a997125209b5c2962f6bd9aef01ea9352c25d0ee", size = 232132, upload-time = "2025-12-13T03:42:52.072Z" }, +] + [[package]] name = "tokenizers" version = "0.22.1" @@ -3013,9 +4032,14 @@ name = "torch" version = "2.8.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.11.*' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ { name = "filelock", marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, @@ -3040,10 +4064,16 @@ name = "torch" version = "2.8.0+cu128" source = { registry = "https://download.pytorch.org/whl/cu128" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version < '3.11' and sys_platform == 'win32'", ] @@ -3120,13 +4150,17 @@ name = "torchvision" version = "0.23.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.11.*' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version >= '3.13' and sys_platform == 'darwin'", + "python_full_version == '3.12.*' and sys_platform == 'darwin'", + "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version == '3.11.*' and sys_platform == 'darwin'", + "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", + "python_full_version < '3.11' and sys_platform == 'darwin'", + "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' and sys_platform != 'linux' and sys_platform != 'win32'" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "numpy", marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "pillow", marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, ] @@ -3143,16 +4177,21 @@ name = "torchvision" version = "0.23.0+cu128" source = { registry = "https://download.pytorch.org/whl/cu128" } resolution-markers = [ - "python_full_version >= '3.12' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'linux'", - "python_full_version >= '3.12' and sys_platform == 'win32'", + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", + "python_full_version >= '3.13' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and sys_platform == 'win32'", "python_full_version == '3.11.*' and sys_platform == 'win32'", "python_full_version < '3.11' and sys_platform == 'win32'", ] dependencies = [ - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform == 'linux') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, + { name = "numpy", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pillow", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] @@ -3188,8 +4227,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, { name = "huggingface-hub" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "numpy", version = "2.3.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "numpy" }, { name = "packaging" }, { name = "pyyaml" }, { name = "regex" }, @@ -3335,6 +4373,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/b5/123f13c975e9f27ab9c0770f514345bd406d0e8d3b7a0723af9d43f710af/wcwidth-0.2.14-py2.py3-none-any.whl", hash = "sha256:a7bb560c8aee30f9957e5f9895805edd20602f2d7f720186dfd906e82b4982e1", size = 37286, upload-time = "2025-09-22T16:29:51.641Z" }, ] +[[package]] +name = "xformers" +version = "0.0.32.post2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, + { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/47/99/68d2e7a17b65a08180d6b120d1d293d5855a0c5999cb30a79fb466b49642/xformers-0.0.32.post2.tar.gz", hash = "sha256:9538be803969c6e1ca16a3ece921e472c24f79970b10be1087a389dcb66e412a", size = 12105990, upload-time = "2025-08-15T11:48:56.584Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/16/4ed2fbf7c20ba20565d9a1fc544c2bef39d409c8eca1949654c6ff9f8452/xformers-0.0.32.post2-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3d1fefe0006eb00a730b3b197ff41aab085f1209b50419baab6152445f09e508", size = 117196673, upload-time = "2025-08-15T11:48:46.726Z" }, + { url = "https://files.pythonhosted.org/packages/50/19/ea64609a535f87a44cb177dc37c6d6a58f6d540adfff819d565fb0657cf7/xformers-0.0.32.post2-cp39-abi3-win_amd64.whl", hash = "sha256:87fa13f4b406ed640554cea6687d0020b496bd6b446ca17d24f9438a079e6f09", size = 100221079, upload-time = "2025-08-15T11:48:53.234Z" }, +] + [[package]] name = "yarl" version = "1.22.0" From 0584a9df0960594dff29a90f6f94c267d2d68fcc Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Tue, 16 Dec 2025 13:45:17 -0500 Subject: [PATCH 2/7] TRT support for personalive. Signed-off-by: BuffMcBigHuge --- pyproject.toml | 12 +- .../core/pipelines/personalive/model.yaml | 7 + .../core/pipelines/personalive/pipeline.py | 250 ++++++++++++- .../personalive/tensorrt/__init__.py | 58 ++++ .../pipelines/personalive/tensorrt/builder.py | 243 +++++++++++++ .../pipelines/personalive/tensorrt/convert.py | 328 ++++++++++++++++++ .../pipelines/personalive/tensorrt/export.py | 129 +++++++ .../personalive/tensorrt/framed_model.py | 305 ++++++++++++++++ .../pipelines/personalive/tensorrt/runner.py | 211 +++++++++++ uv.lock | 270 ++++++++------ 10 files changed, 1707 insertions(+), 106 deletions(-) create mode 100644 src/scope/core/pipelines/personalive/tensorrt/__init__.py create mode 100644 src/scope/core/pipelines/personalive/tensorrt/builder.py create mode 100644 src/scope/core/pipelines/personalive/tensorrt/convert.py create mode 100644 src/scope/core/pipelines/personalive/tensorrt/export.py create mode 100644 src/scope/core/pipelines/personalive/tensorrt/framed_model.py create mode 100644 src/scope/core/pipelines/personalive/tensorrt/runner.py diff --git a/pyproject.toml b/pyproject.toml index c3f53f57..329a881f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,9 +57,16 @@ dependencies = [ "PyOpenGL>=3.1.10; sys_platform == 'win32'", # PersonaLive dependencies "mediapipe>=0.10.11", - "decord>=0.6.0", "xformers>=0.0.32.post2", - "scikit-image>=0.22.0", +] + +[project.optional-dependencies] +tensorrt = [ + "tensorrt>=10.0.0; sys_platform == 'linux' or sys_platform == 'win32'", + "tensorrt-cu12>=10.0.0; sys_platform == 'linux' or sys_platform == 'win32'", + "polygraphy>=0.49.0", + "onnx>=1.17.0", + "onnx-graphsurgeon>=0.5.0", ] [project.scripts] @@ -67,6 +74,7 @@ daydream-scope = "scope.server.app:main" build = "scope.server.build:main" download_models = "scope.server.download_models:main" patch-xformers = "scope.server.patch_xformers:main" +convert-personalive-trt = "scope.core.pipelines.personalive.tensorrt.convert:main" [project.urls] Homepage = "https://github.com/daydreamlive/scope" diff --git a/src/scope/core/pipelines/personalive/model.yaml b/src/scope/core/pipelines/personalive/model.yaml index 520642fb..b26b4f36 100644 --- a/src/scope/core/pipelines/personalive/model.yaml +++ b/src/scope/core/pipelines/personalive/model.yaml @@ -9,6 +9,13 @@ vae_model_path: "sd-vae-ft-mse" # PersonaLive custom weights personalive_weights_path: "PersonaLive/pretrained_weights/personalive" +# TensorRT acceleration (optional) +# Run `convert-personalive-trt` to generate the engine file +# Engine path format: tensorrt/unet_work_{height}x{width}.engine +tensorrt: + onnx_path: "PersonaLive/pretrained_weights/onnx/unet_opt/unet_opt.onnx" + engine_dir: "PersonaLive/pretrained_weights/tensorrt" + # Inference configuration dtype: "fp16" batch_size: 1 diff --git a/src/scope/core/pipelines/personalive/pipeline.py b/src/scope/core/pipelines/personalive/pipeline.py index 1a3da79c..2062ff86 100644 --- a/src/scope/core/pipelines/personalive/pipeline.py +++ b/src/scope/core/pipelines/personalive/pipeline.py @@ -2,6 +2,10 @@ Animates a reference portrait image using driving video frames. Based on: https://github.com/GVCLab/PersonaLive + +Supports optional TensorRT acceleration when installed: + pip install daydream-scope[tensorrt] + convert-personalive-trt --model-dir ./models --height 512 --width 512 """ import logging @@ -37,6 +41,12 @@ get_boxes, ) +# TensorRT support (optional) +try: + from .tensorrt import TRT_AVAILABLE, TRTRunner, get_engine_path +except ImportError: + TRT_AVAILABLE = False + if TYPE_CHECKING: from ..schema import BasePipelineConfig @@ -259,8 +269,16 @@ def __init__( # Enable memory efficient attention self._enable_xformers() + # TensorRT support (optional) + self.use_tensorrt = False + self.trt_runner = None + self._model_dir = model_dir # Store for TensorRT path lookup + + if TRT_AVAILABLE: + self._init_tensorrt(model_dir) + torch.cuda.empty_cache() - logger.info("PersonaLive pipeline initialized") + logger.info(f"PersonaLive pipeline initialized (TensorRT: {self.use_tensorrt})") def _enable_xformers(self): """Enable xformers memory efficient attention if available.""" @@ -271,6 +289,41 @@ def _enable_xformers(self): except Exception as e: logger.warning(f"Could not enable xformers: {e}") + def _init_tensorrt(self, model_dir: Path): + """Initialize TensorRT engine if available. + + Args: + model_dir: Base model directory. + """ + engine_path = get_engine_path(model_dir, self.height, self.width) + + if not engine_path.exists(): + logger.info( + f"TensorRT engine not found at {engine_path}. " + f"Run 'convert-personalive-trt' to create it for faster inference." + ) + return + + try: + device_index = self.device.index if self.device.index is not None else 0 + self.trt_runner = TRTRunner(engine_path, device=self.device) + + # Setup output-to-input bindings for recurrent state + self.trt_runner.bind({ + "motion_hidden_states_out": "motion_hidden_states", + "pose_cond_fea_out": "pose_cond_fea", + "latents": "sample", + }) + + self.use_tensorrt = True + logger.info(f"TensorRT engine loaded from {engine_path}") + + except Exception as e: + logger.warning(f"Failed to load TensorRT engine: {e}") + logger.warning("Falling back to PyTorch inference") + self.trt_runner = None + self.use_tensorrt = False + def prepare(self, **kwargs) -> Requirements | None: """Return input requirements. @@ -381,7 +434,29 @@ def fuse_reference(self, ref_image: Image.Image): encoder_hidden_states=self.encoder_hidden_states, return_dict=False, ) - self.reference_control_reader.update(self.reference_control_writer) + + # Handle TensorRT vs PyTorch mode differently + if self.use_tensorrt: + # For TensorRT, extract reference hidden states and prefill + reference_hidden_states = self.reference_control_writer.output() + self._reference_hidden_states = reference_hidden_states + + # Prefill TensorRT runner with constant inputs + self.trt_runner.clear_prefill() + self.trt_runner.prefill(encoder_hidden_states=self.encoder_hidden_states) + + # Prefill reference hidden states + ref_hidden_names = [ + "d00", "d01", "d10", "d11", "d20", "d21", "m", + "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" + ] + for name in ref_hidden_names: + if name in reference_hidden_states: + self.trt_runner.prefill(**{name: reference_hidden_states[name]}) + else: + # For PyTorch, update the reader + self.reference_control_reader.update(self.reference_control_writer) + self.encoder_hidden_states = self.encoder_hidden_states.to(self.device) # Prepare conditioning tensor for pose encoder @@ -412,6 +487,18 @@ def fuse_reference(self, ref_image: Image.Image): r = (i + 1) * self.temporal_window_size self.latents_pile.append(noisy_latents_first[:, :, l:r]) + # For TensorRT, also prefill initial latents + if self.use_tensorrt: + sample = torch.cat(list(self.latents_pile), dim=2) + # Add placeholder for new latents + new_latents = self.ref_image_latents.unsqueeze(2).repeat( + 1, 1, self.temporal_window_size, 1, 1 + ) + noise = torch.randn_like(new_latents) + new_latents = self.scheduler.add_noise(new_latents, noise, self.timesteps[-1:]) + sample = torch.cat([sample, new_latents], dim=2) + self.trt_runner.prefill(latents=sample) + self.reference_fused = True logger.info("Reference image fused successfully") @@ -470,6 +557,165 @@ def __call__( def _process_frames(self, images: torch.Tensor) -> torch.Tensor: """Process driving frames through the pipeline. + Args: + images: Input tensor of shape (B, C, T, H, W) in [-1, 1] or [0, 1] range. + + Returns: + Output frames tensor of shape (T, H, W, C) in [0, 1] range. + """ + # Route to TensorRT or PyTorch implementation + if self.use_tensorrt: + return self._process_frames_trt(images) + else: + return self._process_frames_pytorch(images) + + @torch.no_grad() + def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: + """Process frames using TensorRT engine. + + Args: + images: Input tensor of shape (B, C, T, H, W) in [-1, 1] or [0, 1] range. + + Returns: + Output frames tensor of shape (T, H, W, C) in [0, 1] range. + """ + batch_size = self.batch_size + device = self.device + temporal_window_size = self.temporal_window_size + temporal_adaptive_step = self.temporal_adaptive_step + + # Reshape from (B, C, T, H, W) to (T, C, H, W) for processing + if images.dim() == 5: + images = images.squeeze(0).permute(1, 0, 2, 3) + + num_frames = images.shape[0] + if num_frames != temporal_window_size: + images = images[-temporal_window_size:] + + if images.min() < 0: + images = images / 2 + 0.5 + images = images.clamp(0, 1) + + # Get keypoints using pose encoder (still PyTorch) + tgt_cond_tensor = self._fast_resize(images, 256, 256).clamp(0, 1) + + if self.first_frame: + mot_bbox_param, kps_ref, kps_frame1, kps_dri = self.pose_encoder.interpolate_kps_online( + self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + ) + self.kps_ref = kps_ref + self.kps_frame1 = kps_frame1 + else: + mot_bbox_param, kps_dri = self.pose_encoder.get_kps( + self.kps_ref, self.kps_frame1, tgt_cond_tensor + ) + + keypoints = draw_keypoints(mot_bbox_param, device=device) + boxes = get_boxes(kps_dri) + keypoints = rearrange(keypoints.unsqueeze(2), 'f c b h w -> b c f h w') + keypoints = keypoints.to(device=device, dtype=self.dtype) + + # Prepare motion input (face crops) + if self.first_frame: + ref_box = get_boxes(mot_bbox_param[:1]) + ref_face = self._crop_face_tensor(self.ref_image_tensor, ref_box[0]) + motion_face = [ref_face] + for i, frame in enumerate(images): + motion_face.append(self._crop_face_tensor(frame, boxes[i])) + pose_cond_tensor = torch.cat(motion_face, dim=0).transpose(0, 1) + pose_cond_tensor = pose_cond_tensor.unsqueeze(0) + + # For first frame, compute initial states using PyTorch motion encoder + motion_hidden_states = self.motion_encoder(pose_cond_tensor) + ref_motion = motion_hidden_states[:, :1] + dri_motion = motion_hidden_states[:, 1:] + + init_motion_hidden_states = self._interpolate_tensors( + ref_motion, dri_motion[:, :1], num=12 + 1 + )[:, :-1] + + # Compute initial pose features using PyTorch pose guider + pose_fea = self.pose_guider(keypoints) + + # Prefill TensorRT with initial states + self.trt_runner.prefill( + motion_hidden_states_out=init_motion_hidden_states, + pose_cond_fea_out=pose_fea[:, :, :temporal_window_size * (temporal_adaptive_step - 1)], + ) + + self.motion_bank = ref_motion + self.first_frame = False + + # Prepare TensorRT inputs + pose = keypoints[:, :, -temporal_window_size:] + motion = dri_motion.transpose(1, 2).unsqueeze(0) # Reshape for TRT + else: + motion_face = [] + for i, frame in enumerate(images): + motion_face.append(self._crop_face_tensor(frame, boxes[i])) + motion = torch.cat(motion_face, dim=0).transpose(0, 1).unsqueeze(0) + pose = keypoints + + # Prepare new noise for next iteration + new_noise = torch.randn( + batch_size, 4, temporal_window_size, + self.height // 8, self.width // 8, + device=device, dtype=self.dtype + ) + + # Run TensorRT inference + outputs = self.trt_runner( + output_names=["pred_video", "motion_out", "latent_first"], + return_torch=True, + pose=pose, + motion=motion.to(self.dtype), + new_noise=new_noise, + ) + + video = outputs["pred_video"] + motion_out = outputs.get("motion_out") + + # Keyframe tracking + if motion_out is not None and self.count > 8: + idx_to_add, self.motion_bank, _ = self._calculate_dis( + self.motion_bank, motion_out.unsqueeze(0), threshold=17.0 + ) + + if len(idx_to_add) > 0 and self.num_khf < 3: + # Update reference hidden states for keyframe + latent_first = outputs.get("latent_first") + if latent_first is not None: + self.reference_control_writer.clear() + self.reference_unet( + latent_first.to(self.reference_unet.dtype), + torch.zeros((batch_size,), dtype=self.dtype, device=device), + encoder_hidden_states=self.encoder_hidden_states, + return_dict=False, + ) + new_ref_hidden = self.reference_control_writer.output() + + # Concatenate new keyframe features + ref_hidden_names = [ + "d00", "d01", "d10", "d11", "d20", "d21", "m", + "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" + ] + for name in ref_hidden_names: + if name in self._reference_hidden_states and name in new_ref_hidden: + self._reference_hidden_states[name] = torch.cat( + [self._reference_hidden_states[name], new_ref_hidden[name]], dim=1 + ) + self.trt_runner.prefill(**{name: self._reference_hidden_states[name]}) + + logger.debug("Added history keyframe (TensorRT)") + self.num_khf += 1 + + self.count += 1 + return video.cpu() + + @torch.no_grad() + def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: + """Process frames using PyTorch (original implementation). + Args: images: Input tensor of shape (B, C, T, H, W) in [-1, 1] or [0, 1] range. diff --git a/src/scope/core/pipelines/personalive/tensorrt/__init__.py b/src/scope/core/pipelines/personalive/tensorrt/__init__.py new file mode 100644 index 00000000..4ec3e734 --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/__init__.py @@ -0,0 +1,58 @@ +"""TensorRT acceleration for PersonaLive pipeline. + +This module provides TensorRT acceleration for the PersonaLive pipeline using +NVIDIA's official TensorRT and polygraphy packages. + +Installation: + pip install daydream-scope[tensorrt] + +Usage: + # Convert models to TensorRT (run once) + convert-personalive-trt --model-dir ./models --height 512 --width 512 + + # The pipeline will automatically use TensorRT when engine is available +""" + +# Check for TensorRT availability before importing +TRT_AVAILABLE = False +try: + import tensorrt # noqa: F401 + from polygraphy.backend.trt import TrtRunner as _TrtRunner # noqa: F401 + + TRT_AVAILABLE = True +except ImportError: + pass + +if TRT_AVAILABLE: + from .builder import build_engine, get_engine_path, is_engine_available + from .runner import TRTRunner + + __all__ = [ + "build_engine", + "get_engine_path", + "is_engine_available", + "TRTRunner", + "TRT_AVAILABLE", + ] +else: + # Provide stub functions when TRT is not available + def get_engine_path(*args, **kwargs): + raise RuntimeError("TensorRT not available. Install with: pip install daydream-scope[tensorrt]") + + def build_engine(*args, **kwargs): + raise RuntimeError("TensorRT not available. Install with: pip install daydream-scope[tensorrt]") + + def is_engine_available(*args, **kwargs): + return False + + class TRTRunner: + def __init__(self, *args, **kwargs): + raise RuntimeError("TensorRT not available. Install with: pip install daydream-scope[tensorrt]") + + __all__ = [ + "build_engine", + "get_engine_path", + "is_engine_available", + "TRTRunner", + "TRT_AVAILABLE", + ] diff --git a/src/scope/core/pipelines/personalive/tensorrt/builder.py b/src/scope/core/pipelines/personalive/tensorrt/builder.py new file mode 100644 index 00000000..12e63f16 --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/builder.py @@ -0,0 +1,243 @@ +"""TensorRT engine builder using polygraphy. + +This module provides utilities to build TensorRT engines from ONNX models +using NVIDIA's polygraphy library. +""" + +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from polygraphy.backend.trt import Profile + +logger = logging.getLogger(__name__) + +# Check for TensorRT availability +TRT_AVAILABLE = False +try: + import tensorrt as trt + from polygraphy.backend.trt import ( + CreateConfig, + Profile, + engine_from_network, + network_from_onnx_path, + save_engine, + ) + from polygraphy.logger import G_LOGGER + + TRT_AVAILABLE = True + # Set polygraphy logger to WARNING to reduce noise + G_LOGGER.severity = G_LOGGER.WARNING +except ImportError: + pass + + +def get_engine_path(model_dir: Path, height: int = 512, width: int = 512) -> Path: + """Get the expected TensorRT engine path for given resolution. + + Args: + model_dir: Base model directory. + height: Output height. + width: Output width. + + Returns: + Path to the TensorRT engine file. + """ + tensorrt_dir = model_dir / "PersonaLive" / "pretrained_weights" / "tensorrt" + return tensorrt_dir / f"unet_work_{height}x{width}.engine" + + +def get_onnx_path(model_dir: Path) -> Path: + """Get the ONNX model path. + + Args: + model_dir: Base model directory. + + Returns: + Path to the optimized ONNX model. + """ + onnx_dir = model_dir / "PersonaLive" / "pretrained_weights" / "onnx" + return onnx_dir / "unet_opt" / "unet_opt.onnx" + + +def create_optimization_profile( + batch_size: int = 1, + height: int = 512, + width: int = 512, + temporal_window_size: int = 4, + temporal_adaptive_step: int = 4, +) -> "Profile": + """Create TensorRT optimization profile for PersonaLive. + + This defines the input shapes for the TensorRT engine, including + min/opt/max shapes for dynamic dimensions. + + Args: + batch_size: Batch size (usually 1). + height: Output height. + width: Output width. + temporal_window_size: Number of frames per temporal window. + temporal_adaptive_step: Number of adaptive steps. + + Returns: + Polygraphy Profile object. + """ + if not TRT_AVAILABLE: + raise RuntimeError("TensorRT is not available. Install with: pip install daydream-scope[tensorrt]") + + # Derived dimensions + tb = temporal_window_size * temporal_adaptive_step # temporal batch size (16) + lh, lw = height // 8, width // 8 # latent height/width + ml, mc = 32, 16 # motion latent size, motion channels + mh, mw = 224, 224 # motion input size + emb = 768 # CLIP embedding dim + lc, ic = 4, 3 # latent channels, image channels + + # UNet channel dimensions + cd0, cd1, cd2, cm = 320, 640, 1280, 1280 + cu1, cu2, cu3 = 1280, 640, 320 + + profile = Profile() + + # Fixed shape inputs + fixed_inputs = { + "sample": (batch_size, lc, tb, lh, lw), + "encoder_hidden_states": (batch_size, 1, emb), + "motion_hidden_states": (batch_size, temporal_window_size * (temporal_adaptive_step - 1), ml, mc), + "motion": (batch_size, ic, temporal_window_size, mh, mw), + "pose_cond_fea": (batch_size, cd0, temporal_window_size * (temporal_adaptive_step - 1), lh, lw), + "pose": (batch_size, ic, temporal_window_size, height, width), + "new_noise": (batch_size, lc, temporal_window_size, lh, lw), + } + + for name, shape in fixed_inputs.items(): + profile.add(name, min=shape, opt=shape, max=shape) + + # Dynamic shape inputs (for reference hidden states that can grow with keyframes) + # Base shapes at 1x, max at 4x for keyframe accumulation + dynamic_inputs = { + "d00": (batch_size, lh * lw, cd0), + "d01": (batch_size, lh * lw, cd0), + "d10": (batch_size, lh * lw // 4, cd1), + "d11": (batch_size, lh * lw // 4, cd1), + "d20": (batch_size, lh * lw // 16, cd2), + "d21": (batch_size, lh * lw // 16, cd2), + "m": (batch_size, lh * lw // 64, cm), + "u10": (batch_size, lh * lw // 16, cu1), + "u11": (batch_size, lh * lw // 16, cu1), + "u12": (batch_size, lh * lw // 16, cu1), + "u20": (batch_size, lh * lw // 4, cu2), + "u21": (batch_size, lh * lw // 4, cu2), + "u22": (batch_size, lh * lw // 4, cu2), + "u30": (batch_size, lh * lw, cu3), + "u31": (batch_size, lh * lw, cu3), + "u32": (batch_size, lh * lw, cu3), + } + + for name, base_shape in dynamic_inputs.items(): + dim0, dim1_base, dim2 = base_shape + # Allow 1x to 4x scaling for keyframe accumulation + min_shape = (dim0, dim1_base, dim2) + opt_shape = (dim0, dim1_base * 2, dim2) + max_shape = (dim0, dim1_base * 4, dim2) + profile.add(name, min=min_shape, opt=opt_shape, max=max_shape) + + return profile + + +def build_engine( + onnx_path: Path | str, + engine_path: Path | str, + height: int = 512, + width: int = 512, + fp16: bool = True, + batch_size: int = 1, + temporal_window_size: int = 4, + temporal_adaptive_step: int = 4, +) -> bool: + """Build TensorRT engine from ONNX model. + + Args: + onnx_path: Path to the ONNX model. + engine_path: Path to save the TensorRT engine. + height: Output height. + width: Output width. + fp16: Use FP16 precision. + batch_size: Batch size. + temporal_window_size: Temporal window size. + temporal_adaptive_step: Temporal adaptive steps. + + Returns: + True if engine was built successfully. + + Raises: + RuntimeError: If TensorRT is not available. + FileNotFoundError: If ONNX model doesn't exist. + """ + if not TRT_AVAILABLE: + raise RuntimeError("TensorRT is not available. Install with: pip install daydream-scope[tensorrt]") + + onnx_path = Path(onnx_path) + engine_path = Path(engine_path) + + if not onnx_path.exists(): + raise FileNotFoundError(f"ONNX model not found: {onnx_path}") + + # Create output directory + engine_path.parent.mkdir(parents=True, exist_ok=True) + + logger.info(f"Building TensorRT engine from {onnx_path}") + logger.info(f"Resolution: {height}x{width}, FP16: {fp16}") + + # Create optimization profile + profile = create_optimization_profile( + batch_size=batch_size, + height=height, + width=width, + temporal_window_size=temporal_window_size, + temporal_adaptive_step=temporal_adaptive_step, + ) + + # Build engine using polygraphy + logger.info("Loading ONNX model and building TensorRT engine...") + logger.info("This may take 10-30 minutes depending on your GPU.") + + try: + engine = engine_from_network( + network_from_onnx_path( + str(onnx_path), + flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM], + ), + config=CreateConfig( + fp16=fp16, + refittable=False, + profiles=[profile], + ), + ) + + logger.info(f"Saving engine to {engine_path}") + save_engine(engine, str(engine_path)) + + logger.info("TensorRT engine built successfully!") + return True + + except Exception as e: + logger.error(f"Failed to build TensorRT engine: {e}") + raise + + +def is_engine_available(model_dir: Path, height: int = 512, width: int = 512) -> bool: + """Check if a TensorRT engine is available for the given configuration. + + Args: + model_dir: Base model directory. + height: Output height. + width: Output width. + + Returns: + True if engine file exists. + """ + engine_path = get_engine_path(model_dir, height, width) + return engine_path.exists() diff --git a/src/scope/core/pipelines/personalive/tensorrt/convert.py b/src/scope/core/pipelines/personalive/tensorrt/convert.py new file mode 100644 index 00000000..99adf4b4 --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/convert.py @@ -0,0 +1,328 @@ +"""CLI script for converting PersonaLive models to TensorRT. + +This script handles the full conversion pipeline: +1. Load PyTorch models +2. Create bundled UNetWork model +3. Export to ONNX +4. Optimize ONNX +5. Build TensorRT engine + +Usage: + convert-personalive-trt --model-dir ./models --height 512 --width 512 +""" + +import argparse +import gc +import logging +import sys +import time +from pathlib import Path + +import torch +from diffusers import AutoencoderKL +from diffusers.models.attention_processor import AttnProcessor +from omegaconf import OmegaConf + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", +) +logger = logging.getLogger(__name__) + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Convert PersonaLive models to TensorRT engine", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--model-dir", + type=Path, + required=True, + help="Path to the models directory containing PersonaLive weights", + ) + parser.add_argument( + "--height", + type=int, + default=512, + help="Output video height", + ) + parser.add_argument( + "--width", + type=int, + default=512, + help="Output video width", + ) + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Batch size", + ) + parser.add_argument( + "--fp32", + action="store_true", + help="Use FP32 instead of FP16 precision", + ) + parser.add_argument( + "--skip-onnx", + action="store_true", + help="Skip ONNX export (use existing ONNX files)", + ) + parser.add_argument( + "--device", + type=str, + default="cuda:0", + help="CUDA device to use", + ) + return parser.parse_args() + + +def load_models(model_dir: Path, device: torch.device, dtype: torch.dtype): + """Load all PersonaLive models. + + Args: + model_dir: Base model directory. + device: Target device. + dtype: Data type for models. + + Returns: + Dictionary containing all loaded models and config. + """ + from ..modules import ( + DDIMScheduler, + MotEncoder, + PoseGuider, + UNet3DConditionModel, + ) + + logger.info("Loading PersonaLive models...") + + # Paths + personalive_dir = model_dir / "PersonaLive" / "pretrained_weights" + pretrained_base_path = personalive_dir / "sd-image-variations-diffusers" + vae_path = personalive_dir / "sd-vae-ft-mse" + weights_path = personalive_dir / "personalive" + + # Load model config for UNet kwargs + config_path = Path(__file__).parent.parent / "model.yaml" + model_config = OmegaConf.load(config_path) + unet_kwargs = OmegaConf.to_container(model_config.unet_additional_kwargs) + sched_kwargs = OmegaConf.to_container(model_config.noise_scheduler_kwargs) + + # Load pose guider + logger.info("Loading PoseGuider...") + pose_guider = PoseGuider().to(device=device, dtype=dtype) + pose_guider_path = weights_path / "pose_guider.pth" + if pose_guider_path.exists(): + state_dict = torch.load(pose_guider_path, map_location="cpu") + pose_guider.load_state_dict(state_dict) + del state_dict + + # Load motion encoder + logger.info("Loading MotEncoder...") + motion_encoder = MotEncoder().to(dtype=dtype, device=device).eval() + motion_encoder_path = weights_path / "motion_encoder.pth" + if motion_encoder_path.exists(): + state_dict = torch.load(motion_encoder_path, map_location="cpu") + motion_encoder.load_state_dict(state_dict) + del state_dict + motion_encoder.set_attn_processor(AttnProcessor()) + + # Load denoising UNet + logger.info("Loading UNet3DConditionModel...") + denoising_unet = UNet3DConditionModel.from_pretrained_2d( + str(pretrained_base_path), + "", + subfolder="unet", + unet_additional_kwargs=unet_kwargs, + ).to(dtype=dtype, device=device) + + denoising_unet_path = weights_path / "denoising_unet.pth" + if denoising_unet_path.exists(): + state_dict = torch.load(denoising_unet_path, map_location="cpu") + denoising_unet.load_state_dict(state_dict, strict=False) + del state_dict + + temporal_module_path = weights_path / "temporal_module.pth" + if temporal_module_path.exists(): + state_dict = torch.load(temporal_module_path, map_location="cpu") + denoising_unet.load_state_dict(state_dict, strict=False) + del state_dict + + denoising_unet.set_attn_processor(AttnProcessor()) + + # Load VAE + logger.info("Loading VAE...") + vae = AutoencoderKL.from_pretrained(str(vae_path)).to(device=device, dtype=dtype) + vae.set_default_attn_processor() + + # Setup scheduler + scheduler = DDIMScheduler(**sched_kwargs) + scheduler.to(device) + + # Create timesteps tensor + # For TensorRT, we use a fixed timestep pattern + timesteps = torch.tensor( + [0, 0, 0, 0, 333, 333, 333, 333, 666, 666, 666, 666, 999, 999, 999, 999], + device=device, + ).long() + scheduler.set_step_length(333) + + return { + "pose_guider": pose_guider, + "motion_encoder": motion_encoder, + "denoising_unet": denoising_unet, + "vae": vae, + "scheduler": scheduler, + "timesteps": timesteps, + } + + +def export_to_onnx( + models: dict, + onnx_path: Path, + height: int, + width: int, + batch_size: int, + device: torch.device, + dtype: torch.dtype, +): + """Export bundled model to ONNX. + + Args: + models: Dictionary of loaded models. + onnx_path: Path to save ONNX model. + height: Output height. + width: Output width. + batch_size: Batch size. + device: Target device. + dtype: Data type. + """ + from .export import export_onnx, optimize_onnx + from .framed_model import UNetWork + + logger.info("Creating bundled UNetWork model...") + + # Create bundled model + unet_work = UNetWork( + pose_guider=models["pose_guider"], + motion_encoder=models["motion_encoder"], + denoising_unet=models["denoising_unet"], + vae=models["vae"], + scheduler=models["scheduler"], + timesteps=models["timesteps"], + ) + + # Generate sample inputs + logger.info("Generating sample inputs...") + sample_inputs = unet_work.get_sample_input(batch_size, height, width, dtype, device) + + # Export to ONNX + raw_onnx_path = onnx_path.parent / "unet" / "unet.onnx" + export_onnx( + model=unet_work, + onnx_path=raw_onnx_path, + sample_inputs=sample_inputs, + input_names=UNetWork.get_input_names(), + output_names=UNetWork.get_output_names(), + dynamic_axes=UNetWork.get_dynamic_axes(), + opset_version=17, + ) + + # Optimize ONNX + logger.info("Optimizing ONNX model...") + optimize_onnx(raw_onnx_path, onnx_path) + + # Clean up + del unet_work, sample_inputs + gc.collect() + torch.cuda.empty_cache() + + +def main(): + """Main entry point for TensorRT conversion.""" + args = parse_args() + + # Validate inputs + if not args.model_dir.exists(): + logger.error(f"Model directory not found: {args.model_dir}") + sys.exit(1) + + personalive_weights = args.model_dir / "PersonaLive" / "pretrained_weights" / "personalive" + if not personalive_weights.exists(): + logger.error(f"PersonaLive weights not found at: {personalive_weights}") + logger.error("Please download PersonaLive models first using: download_models") + sys.exit(1) + + # Check TensorRT availability + try: + from .builder import TRT_AVAILABLE, build_engine, get_engine_path, get_onnx_path + + if not TRT_AVAILABLE: + raise ImportError() + except ImportError: + logger.error("TensorRT is not available.") + logger.error("Install with: pip install daydream-scope[tensorrt]") + sys.exit(1) + + # Setup + device = torch.device(args.device) + dtype = torch.float32 if args.fp32 else torch.float16 + + logger.info(f"Converting PersonaLive models to TensorRT") + logger.info(f" Model directory: {args.model_dir}") + logger.info(f" Resolution: {args.height}x{args.width}") + logger.info(f" Precision: {'FP32' if args.fp32 else 'FP16'}") + logger.info(f" Device: {device}") + + start_time = time.time() + + # Get paths + onnx_path = get_onnx_path(args.model_dir) + engine_path = get_engine_path(args.model_dir, args.height, args.width) + + # Step 1: Export to ONNX (if needed) + if not args.skip_onnx or not onnx_path.exists(): + logger.info("Step 1: Exporting to ONNX...") + models = load_models(args.model_dir, device, dtype) + export_to_onnx( + models=models, + onnx_path=onnx_path, + height=args.height, + width=args.width, + batch_size=args.batch_size, + device=device, + dtype=dtype, + ) + # Clean up models + del models + gc.collect() + torch.cuda.empty_cache() + else: + logger.info("Step 1: Skipping ONNX export (using existing files)") + + # Step 2: Build TensorRT engine + logger.info("Step 2: Building TensorRT engine...") + logger.info("This may take 10-30 minutes depending on your GPU.") + + build_engine( + onnx_path=onnx_path, + engine_path=engine_path, + height=args.height, + width=args.width, + fp16=not args.fp32, + batch_size=args.batch_size, + ) + + elapsed = time.time() - start_time + logger.info(f"Conversion completed in {elapsed / 60:.1f} minutes") + logger.info(f"TensorRT engine saved to: {engine_path}") + logger.info("") + logger.info("The PersonaLive pipeline will automatically use TensorRT") + logger.info("when the engine file is present.") + + +if __name__ == "__main__": + main() diff --git a/src/scope/core/pipelines/personalive/tensorrt/export.py b/src/scope/core/pipelines/personalive/tensorrt/export.py new file mode 100644 index 00000000..3d7bebcc --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/export.py @@ -0,0 +1,129 @@ +"""ONNX export utilities for PersonaLive models. + +This module provides utilities to export PersonaLive models to ONNX format +for subsequent TensorRT conversion. +""" + +import gc +import logging +import os +from pathlib import Path + +import torch + +logger = logging.getLogger(__name__) + + +def export_onnx( + model: torch.nn.Module, + onnx_path: Path | str, + sample_inputs: dict[str, torch.Tensor], + input_names: list[str], + output_names: list[str], + dynamic_axes: dict[str, dict[int, str]] | None = None, + opset_version: int = 17, +) -> None: + """Export a PyTorch model to ONNX format. + + Args: + model: PyTorch model to export. + onnx_path: Path to save the ONNX model. + sample_inputs: Sample inputs for tracing. + input_names: Names of input tensors. + output_names: Names of output tensors. + dynamic_axes: Dynamic axis specifications. + opset_version: ONNX opset version. + """ + onnx_path = Path(onnx_path) + onnx_path.parent.mkdir(parents=True, exist_ok=True) + + logger.info(f"Exporting model to ONNX: {onnx_path}") + + # Prepare inputs as tuple in correct order + inputs = tuple(sample_inputs[name] for name in input_names) + + with torch.inference_mode(), torch.autocast("cuda"): + torch.onnx.export( + model, + inputs, + str(onnx_path), + export_params=True, + opset_version=opset_version, + do_constant_folding=True, + input_names=input_names, + output_names=output_names, + dynamic_axes=dynamic_axes or {}, + ) + + logger.info(f"ONNX model exported to {onnx_path}") + + # Clean up + gc.collect() + torch.cuda.empty_cache() + + +def optimize_onnx( + onnx_path: Path | str, + onnx_opt_path: Path | str, +) -> None: + """Optimize ONNX model and save with external data. + + For large models, this saves weights as external data to avoid + protobuf size limits. + + Args: + onnx_path: Path to input ONNX model. + onnx_opt_path: Path to save optimized ONNX model. + """ + try: + import onnx + except ImportError: + raise RuntimeError("onnx package required. Install with: pip install onnx") + + onnx_path = Path(onnx_path) + onnx_opt_path = Path(onnx_opt_path) + onnx_opt_path.parent.mkdir(parents=True, exist_ok=True) + + logger.info(f"Optimizing ONNX model: {onnx_path} -> {onnx_opt_path}") + + model = onnx.load(str(onnx_path)) + name = onnx_opt_path.stem + + # Save with external data for large models + onnx.save( + model, + str(onnx_opt_path), + save_as_external_data=True, + all_tensors_to_one_file=True, + location=f"{name}.onnx.data", + size_threshold=1024, + ) + + logger.info(f"Optimized ONNX model saved to {onnx_opt_path}") + + +def verify_onnx(onnx_path: Path | str) -> bool: + """Verify ONNX model is valid. + + Args: + onnx_path: Path to ONNX model. + + Returns: + True if model is valid. + """ + try: + import onnx + except ImportError: + logger.warning("onnx package not available, skipping verification") + return True + + onnx_path = Path(onnx_path) + + try: + model = onnx.load(str(onnx_path)) + onnx.checker.check_model(model) + logger.info(f"ONNX model {onnx_path} is valid") + return True + except Exception as e: + logger.error(f"ONNX model verification failed: {e}") + return False diff --git a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py new file mode 100644 index 00000000..0bec0373 --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py @@ -0,0 +1,305 @@ +"""Framed model for ONNX export. + +This module provides the UNetWork class that bundles multiple models +into a single module for ONNX export and TensorRT conversion. + +The bundled models include: +- PoseGuider: Extracts pose features from keypoint images +- MotEncoder: Encodes facial motion features +- UNet3DConditionModel: Main denoising UNet +- VAE decoder: Decodes latents to images +- DDIMScheduler: Denoising scheduler step +""" + +import torch +import torch.nn as nn +from einops import rearrange + + +class UNetWork(nn.Module): + """Bundled model for TensorRT export. + + This class wraps the pose_guider, motion_encoder, denoising_unet, + vae decoder, and scheduler into a single module that can be exported + to ONNX for TensorRT conversion. + + The forward pass handles a complete denoising step including: + 1. Encoding new pose features + 2. Encoding new motion features + 3. Running the UNet denoiser + 4. Scheduler step + 5. VAE decoding for the output frames + """ + + def __init__( + self, + pose_guider: nn.Module, + motion_encoder: nn.Module, + denoising_unet: nn.Module, + vae: nn.Module, + scheduler, + timesteps: torch.Tensor, + ): + """Initialize the bundled model. + + Args: + pose_guider: PoseGuider model. + motion_encoder: MotEncoder model. + denoising_unet: UNet3DConditionModel for denoising. + vae: AutoencoderKL for decoding. + scheduler: DDIMScheduler instance. + timesteps: Fixed timesteps tensor for denoising. + """ + super().__init__() + self.pose_guider = pose_guider + self.motion_encoder = motion_encoder + self.unet = denoising_unet + self.vae = vae + self.scheduler = scheduler + self.timesteps = timesteps + + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: + """Decode latents to images using VAE. + + Args: + latents: Latent tensor of shape (B, C, H, W). + + Returns: + Image tensor of shape (B, H, W, C) in [0, 1] range. + """ + latents = latents / 0.18215 + images = self.vae.decode(latents).sample + images = rearrange(images, "b c h w -> b h w c") + images = (images / 2 + 0.5).clamp(0, 1) + return images + + def forward( + self, + sample: torch.Tensor, + encoder_hidden_states: torch.Tensor, + motion_hidden_states: torch.Tensor, + motion: torch.Tensor, + pose_cond_fea: torch.Tensor, + pose: torch.Tensor, + new_noise: torch.Tensor, + d00: torch.Tensor, + d01: torch.Tensor, + d10: torch.Tensor, + d11: torch.Tensor, + d20: torch.Tensor, + d21: torch.Tensor, + m: torch.Tensor, + u10: torch.Tensor, + u11: torch.Tensor, + u12: torch.Tensor, + u20: torch.Tensor, + u21: torch.Tensor, + u22: torch.Tensor, + u30: torch.Tensor, + u31: torch.Tensor, + u32: torch.Tensor, + ): + """Forward pass for the bundled model. + + Args: + sample: Noisy latent tensor (B, C, T, H, W). + encoder_hidden_states: CLIP embeddings (B, 1, D). + motion_hidden_states: Accumulated motion features (B, T, D1, D2). + motion: New face crop tensor (B, C, T, H, W). + pose_cond_fea: Accumulated pose features (B, C, T, H, W). + pose: New pose keypoints (B, C, T, H, W). + new_noise: Noise for next iteration (B, C, T, H, W). + d00-u32: Reference hidden states from reference UNet. + + Returns: + Tuple of: + - pred_video: Decoded video frames (T, H, W, C). + - latents: Updated latent tensor for next iteration. + - pose_cond_fea_out: Updated pose features. + - motion_hidden_states_out: Updated motion features. + - motion_out: Motion features for keyframe tracking. + - latent_first: First latent for potential keyframe update. + """ + # Encode new pose features + new_pose_cond_fea = self.pose_guider(pose) + pose_cond_fea = torch.cat([pose_cond_fea, new_pose_cond_fea], dim=2) + + # Encode new motion features + new_motion_hidden_states = self.motion_encoder(motion) + motion_hidden_states = torch.cat([motion_hidden_states, new_motion_hidden_states], dim=1) + + # Prepare encoder hidden states for UNet + encoder_hidden_states_combined = [encoder_hidden_states, motion_hidden_states] + + # Run UNet with explicit reference hidden states + score = self.unet( + sample, + self.timesteps, + encoder_hidden_states_combined, + pose_cond_fea, + d00, d01, d10, d11, d20, d21, m, + u10, u11, u12, u20, u21, u22, u30, u31, u32, + ) + + # Scheduler step + score = rearrange(score, 'b c f h w -> (b f) c h w') + sample_flat = rearrange(sample, 'b c f h w -> (b f) c h w') + + latents_model_input, pred_original_sample = self.scheduler.step( + score, self.timesteps, sample_flat, return_dict=False + ) + + latents_model_input = latents_model_input.to(sample.dtype) + pred_original_sample = pred_original_sample.to(sample.dtype) + + # Reshape back to 5D + latents_model_input = rearrange(latents_model_input, '(b f) c h w -> b c f h w', f=16) + + # Decode first 4 frames (temporal_window_size) + pred_video = self._decode_latents(pred_original_sample[:4]) + + # Prepare outputs for next iteration + # Shift latents: drop first temporal_window_size, add new_noise + latents = torch.cat([latents_model_input[:, :, 4:, :, :], new_noise], dim=2) + + # Shift pose and motion features + pose_cond_fea_out = pose_cond_fea[:, :, 4:, :, :] + motion_hidden_states_out = motion_hidden_states[:, 4:, :, :] + + # Motion for keyframe tracking (first frame's motion) + motion_out = motion_hidden_states[:, :1, :, :] + + # First predicted latent for potential keyframe reference update + latent_first = pred_original_sample[:1] + + return ( + pred_video, + latents, + pose_cond_fea_out, + motion_hidden_states_out, + motion_out, + latent_first, + ) + + def get_sample_input( + self, + batch_size: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + ) -> dict[str, torch.Tensor]: + """Generate sample inputs for ONNX export. + + Args: + batch_size: Batch size (usually 1). + height: Output height. + width: Output width. + dtype: Data type (usually float16). + device: Target device. + + Returns: + Dictionary of sample input tensors. + """ + # Constants + tw = 4 # temporal_window_size + ts = 4 # temporal_adaptive_step + tb = tw * ts # temporal batch size (16) + ml, mc = 32, 16 # motion latent size, channels + mh, mw = 224, 224 # motion input size + + lh, lw = height // 8, width // 8 # latent height/width + + # UNet channels + cd0, cd1, cd2, cm = 320, 640, 1280, 1280 + cu1, cu2, cu3 = 1280, 640, 320 + + emb = 768 # CLIP embedding dim + lc, ic = 4, 3 # latent channels, image channels + + shapes = { + "sample": (batch_size, lc, tb, lh, lw), + "encoder_hidden_states": (batch_size, 1, emb), + "motion_hidden_states": (batch_size, tw * (ts - 1), ml, mc), + "motion": (batch_size, ic, tw, mh, mw), + "pose_cond_fea": (batch_size, cd0, tw * (ts - 1), lh, lw), + "pose": (batch_size, ic, tw, height, width), + "new_noise": (batch_size, lc, tw, lh, lw), + "d00": (batch_size, lh * lw, cd0), + "d01": (batch_size, lh * lw, cd0), + "d10": (batch_size, lh * lw // 4, cd1), + "d11": (batch_size, lh * lw // 4, cd1), + "d20": (batch_size, lh * lw // 16, cd2), + "d21": (batch_size, lh * lw // 16, cd2), + "m": (batch_size, lh * lw // 64, cm), + "u10": (batch_size, lh * lw // 16, cu1), + "u11": (batch_size, lh * lw // 16, cu1), + "u12": (batch_size, lh * lw // 16, cu1), + "u20": (batch_size, lh * lw // 4, cu2), + "u21": (batch_size, lh * lw // 4, cu2), + "u22": (batch_size, lh * lw // 4, cu2), + "u30": (batch_size, lh * lw, cu3), + "u31": (batch_size, lh * lw, cu3), + "u32": (batch_size, lh * lw, cu3), + } + + return {name: torch.randn(shape, dtype=dtype, device=device) for name, shape in shapes.items()} + + @staticmethod + def get_input_names() -> list[str]: + """Get list of input tensor names for ONNX export.""" + return [ + "sample", + "encoder_hidden_states", + "motion_hidden_states", + "motion", + "pose_cond_fea", + "pose", + "new_noise", + "d00", "d01", "d10", "d11", "d20", "d21", "m", + "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32", + ] + + @staticmethod + def get_output_names() -> list[str]: + """Get list of output tensor names for ONNX export.""" + return [ + "pred_video", + "latents", + "pose_cond_fea_out", + "motion_hidden_states_out", + "motion_out", + "latent_first", + ] + + @staticmethod + def get_dynamic_axes() -> dict[str, dict[int, str]]: + """Get dynamic axis specifications for ONNX export. + + Returns: + Dictionary mapping tensor names to their dynamic axes. + """ + return { + # Resolution-dependent axes + "sample": {3: "h_64", 4: "w_64"}, + "pose_cond_fea": {3: "h_64", 4: "w_64"}, + "pose": {3: "h_512", 4: "w_512"}, + "new_noise": {3: "h_64", 4: "w_64"}, + # Dynamic reference hidden states (for keyframe accumulation) + "d00": {1: "len_4096"}, + "d01": {1: "len_4096"}, + "u30": {1: "len_4096"}, + "u31": {1: "len_4096"}, + "u32": {1: "len_4096"}, + "d10": {1: "len_1024"}, + "d11": {1: "len_1024"}, + "u20": {1: "len_1024"}, + "u21": {1: "len_1024"}, + "u22": {1: "len_1024"}, + "d20": {1: "len_256"}, + "d21": {1: "len_256"}, + "u10": {1: "len_256"}, + "u11": {1: "len_256"}, + "u12": {1: "len_256"}, + "m": {1: "len_64"}, + } diff --git a/src/scope/core/pipelines/personalive/tensorrt/runner.py b/src/scope/core/pipelines/personalive/tensorrt/runner.py new file mode 100644 index 00000000..70f9df8c --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/runner.py @@ -0,0 +1,211 @@ +"""TensorRT engine runner using polygraphy. + +This module provides a TensorRT engine runner that wraps the engine +and handles inference with PyTorch tensor I/O. +""" + +import logging +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + +# Check for TensorRT availability +TRT_AVAILABLE = False +try: + import tensorrt as trt + from polygraphy.backend.common import BytesFromPath + from polygraphy.backend.trt import EngineFromBytes, TrtRunner + + TRT_AVAILABLE = True +except ImportError: + pass + + +class TRTRunner: + """TensorRT engine runner with PyTorch tensor support. + + This class wraps a TensorRT engine and provides methods for inference + with PyTorch tensors as input/output. + + Attributes: + engine_path: Path to the TensorRT engine file. + device: CUDA device for inference. + runner: Polygraphy TrtRunner instance. + """ + + def __init__( + self, + engine_path: Path | str, + device: torch.device | None = None, + ): + """Initialize TensorRT runner. + + Args: + engine_path: Path to the TensorRT engine file. + device: CUDA device for inference. + + Raises: + RuntimeError: If TensorRT is not available. + FileNotFoundError: If engine file doesn't exist. + """ + if not TRT_AVAILABLE: + raise RuntimeError("TensorRT is not available. Install with: pip install daydream-scope[tensorrt]") + + engine_path = Path(engine_path) + if not engine_path.exists(): + raise FileNotFoundError(f"TensorRT engine not found: {engine_path}") + + self.engine_path = engine_path + self.device = device or torch.device("cuda") + self._runner = None + self._engine = None + self._context = None + + # Pre-filled inputs (constants that don't change per frame) + self._prefilled: dict[str, np.ndarray] = {} + + # Output bindings (for reusing memory between outputs and inputs) + self._output_bindings: dict[str, str] = {} + + logger.info(f"Loading TensorRT engine from {engine_path}") + self._load_engine() + + def _load_engine(self): + """Load the TensorRT engine.""" + # Load engine bytes + engine_bytes = BytesFromPath(str(self.engine_path)) + + # Create engine from bytes + self._engine = EngineFromBytes(engine_bytes) + + # Create runner (context manager handles activation) + self._runner = TrtRunner(self._engine) + self._runner.activate() + + logger.info("TensorRT engine loaded successfully") + + def prefill(self, **inputs: torch.Tensor | np.ndarray): + """Pre-fill constant inputs that don't change per frame. + + These inputs will be automatically included in every inference call. + + Args: + **inputs: Named inputs as PyTorch tensors or numpy arrays. + """ + for name, tensor in inputs.items(): + if isinstance(tensor, torch.Tensor): + # Convert to numpy, ensuring contiguous memory + tensor = tensor.detach().cpu().numpy() + self._prefilled[name] = np.ascontiguousarray(tensor) + + def bind(self, output_to_input: dict[str, str]): + """Bind output tensors to input tensors for memory reuse. + + This allows outputs from one inference to be used as inputs + to the next inference without memory copies. + + Args: + output_to_input: Mapping from output names to input names. + """ + self._output_bindings = output_to_input.copy() + + def clear_prefill(self): + """Clear all pre-filled inputs.""" + self._prefilled.clear() + + def __call__( + self, + output_names: list[str] | None = None, + return_torch: bool = True, + **inputs: torch.Tensor | np.ndarray, + ) -> dict[str, torch.Tensor | np.ndarray]: + """Run inference on the TensorRT engine. + + Args: + output_names: List of output names to return. If None, returns all outputs. + return_torch: If True, return PyTorch tensors. Otherwise, return numpy arrays. + **inputs: Named inputs as PyTorch tensors or numpy arrays. + + Returns: + Dictionary mapping output names to tensors/arrays. + """ + if self._runner is None: + raise RuntimeError("TensorRT runner not initialized") + + # Prepare feed dict with prefilled values + feed_dict = dict(self._prefilled) + + # Add current inputs + for name, tensor in inputs.items(): + if isinstance(tensor, torch.Tensor): + tensor = tensor.detach().cpu().numpy() + feed_dict[name] = np.ascontiguousarray(tensor) + + # Run inference + outputs = self._runner.infer(feed_dict=feed_dict) + + # Filter outputs if requested + if output_names is not None: + outputs = {k: v for k, v in outputs.items() if k in output_names} + + # Update prefilled values with bound outputs + for output_name, input_name in self._output_bindings.items(): + if output_name in outputs: + self._prefilled[input_name] = outputs[output_name] + + # Convert to PyTorch if requested + if return_torch: + outputs = { + name: torch.from_numpy(arr).to(self.device) + for name, arr in outputs.items() + } + + return outputs + + def get_input_names(self) -> list[str]: + """Get list of input tensor names.""" + if self._runner is None: + return [] + return list(self._runner.get_input_metadata().keys()) + + def get_output_names(self) -> list[str]: + """Get list of output tensor names.""" + if self._runner is None: + return [] + return list(self._runner.get_output_metadata().keys()) + + def get_input_shapes(self) -> dict[str, tuple]: + """Get input tensor shapes.""" + if self._runner is None: + return {} + return {name: tuple(meta.shape) for name, meta in self._runner.get_input_metadata().items()} + + def get_output_shapes(self) -> dict[str, tuple]: + """Get output tensor shapes.""" + if self._runner is None: + return {} + return {name: tuple(meta.shape) for name, meta in self._runner.get_output_metadata().items()} + + def __del__(self): + """Clean up resources.""" + if self._runner is not None: + try: + self._runner.deactivate() + except Exception: + pass + + def __repr__(self) -> str: + """String representation.""" + inputs = self.get_input_shapes() + outputs = self.get_output_shapes() + return ( + f"TRTRunner(\n" + f" engine={self.engine_path.name},\n" + f" inputs={list(inputs.keys())},\n" + f" outputs={list(outputs.keys())}\n" + f")" + ) diff --git a/uv.lock b/uv.lock index 50e54fd6..f10ec6f5 100644 --- a/uv.lock +++ b/uv.lock @@ -808,6 +808,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/0d/c3/e90f4a4feae6410f914f8ebac129b9ae7a8c92eb60a638012dde42030a9d/cryptography-46.0.3-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:6b5063083824e5509fdba180721d55909ffacccc8adbec85268b48439423d78c", size = 3438528, upload-time = "2025-10-15T23:18:26.227Z" }, ] +[[package]] +name = "cuda-toolkit" +version = "13.0.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/35/b6/0dd6813dfaf67e1136f6fd94cac44df2c1071de3da05a68cdb4f2402d166/cuda_toolkit-13.0.1-py2.py3-none-any.whl", hash = "sha256:fc763424c5f5dde5c865358777928d23aa18a0165d380980bd3a792176188ffa", size = 2359, upload-time = "2025-09-08T21:06:25.243Z" }, +] + +[package.optional-dependencies] +cudart = [ + { name = "nvidia-cuda-runtime", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] + [[package]] name = "cycler" version = "0.12.1" @@ -824,7 +837,6 @@ source = { editable = "." } dependencies = [ { name = "accelerate" }, { name = "aiortc" }, - { name = "decord" }, { name = "diffusers" }, { name = "easydict" }, { name = "einops" }, @@ -843,7 +855,6 @@ dependencies = [ { name = "safetensors" }, { name = "sageattention", version = "2.2.0", source = { url = "https://github.com/daydreamlive/SageAttention/releases/download/v2.2.0-linux/sageattention-2.2.0-cp310-cp310-linux_x86_64.whl" }, marker = "sys_platform == 'linux'" }, { name = "sageattention", version = "2.2.0+cu128torch2.8.0.post3", source = { url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl" }, marker = "sys_platform == 'win32'" }, - { name = "scikit-image" }, { name = "spoutgl", marker = "sys_platform == 'win32'" }, { name = "torch", version = "2.8.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform != 'linux' and sys_platform != 'win32'" }, { name = "torch", version = "2.8.0+cu128", source = { registry = "https://download.pytorch.org/whl/cu128" }, marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -859,6 +870,15 @@ dependencies = [ { name = "xformers" }, ] +[package.optional-dependencies] +tensorrt = [ + { name = "onnx" }, + { name = "onnx-graphsurgeon" }, + { name = "polygraphy" }, + { name = "tensorrt", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tensorrt-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, +] + [package.dev-dependencies] dev = [ { name = "freezegun" }, @@ -874,7 +894,6 @@ dev = [ requires-dist = [ { name = "accelerate", specifier = ">=1.1.1" }, { name = "aiortc", specifier = ">=1.13.0" }, - { name = "decord", specifier = ">=0.6.0" }, { name = "diffusers", specifier = ">=0.31.0" }, { name = "easydict", specifier = ">=1.13" }, { name = "einops", specifier = ">=0.8.1" }, @@ -888,13 +907,17 @@ requires-dist = [ { name = "lmdb", specifier = ">=1.7.3" }, { name = "mediapipe", specifier = ">=0.10.11" }, { name = "omegaconf", specifier = ">=2.3.0" }, + { name = "onnx", marker = "extra == 'tensorrt'", specifier = ">=1.17.0" }, + { name = "onnx-graphsurgeon", marker = "extra == 'tensorrt'", specifier = ">=0.5.0" }, { name = "peft", specifier = ">=0.17.1" }, + { name = "polygraphy", marker = "extra == 'tensorrt'", specifier = ">=0.49.0" }, { name = "pyopengl", marker = "sys_platform == 'win32'", specifier = ">=3.1.10" }, { name = "safetensors", specifier = ">=0.6.2" }, { name = "sageattention", marker = "sys_platform == 'linux'", url = "https://github.com/daydreamlive/SageAttention/releases/download/v2.2.0-linux/sageattention-2.2.0-cp310-cp310-linux_x86_64.whl" }, { name = "sageattention", marker = "sys_platform == 'win32'", url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl" }, - { name = "scikit-image", specifier = ">=0.22.0" }, { name = "spoutgl", marker = "sys_platform == 'win32'", specifier = ">=0.1.1" }, + { name = "tensorrt", marker = "(sys_platform == 'linux' and extra == 'tensorrt') or (sys_platform == 'win32' and extra == 'tensorrt')", specifier = ">=10.0.0" }, + { name = "tensorrt-cu12", marker = "(sys_platform == 'linux' and extra == 'tensorrt') or (sys_platform == 'win32' and extra == 'tensorrt')", specifier = ">=10.0.0" }, { name = "torch", marker = "sys_platform != 'linux' and sys_platform != 'win32'", specifier = "==2.8.0" }, { name = "torch", marker = "sys_platform == 'linux' or sys_platform == 'win32'", specifier = "==2.8.0", index = "https://download.pytorch.org/whl/cu128" }, { name = "torchao", specifier = "==0.13.0" }, @@ -908,6 +931,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.35.0" }, { name = "xformers", specifier = ">=0.0.32.post2" }, ] +provides-extras = ["tensorrt"] [package.metadata.requires-dev] dev = [ @@ -920,18 +944,6 @@ dev = [ { name = "twine", specifier = ">=5.0.0" }, ] -[[package]] -name = "decord" -version = "0.6.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" }, - { url = "https://files.pythonhosted.org/packages/6c/be/e15b5b866da452e62635a7b27513f31cb581fa2ea9cc9b768b535d62a955/decord-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02665d7c4f1193a330205a791bc128f7e108eb6ae5b67144437a02f700943bad", size = 24733380, upload-time = "2021-06-14T21:30:57.766Z" }, -] - [[package]] name = "diffusers" version = "0.35.2" @@ -1856,18 +1868,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/e9/0d4add7873a73e462aeb45c036a2dead2562b825aa46ba326727b3f31016/kiwisolver-1.4.9-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:fb940820c63a9590d31d88b815e7a3aa5915cad3ce735ab45f0c730b39547de1", size = 73929, upload-time = "2025-08-10T21:27:48.236Z" }, ] -[[package]] -name = "lazy-loader" -version = "0.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/6f/6b/c875b30a1ba490860c93da4cabf479e03f584eba06fe5963f6f6644653d8/lazy_loader-0.4.tar.gz", hash = "sha256:47c75182589b91a4e1a85a136c074285a5ad4d9f39c63e0d7fb76391c4574cd1", size = 15431, upload-time = "2024-04-05T13:03:12.261Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/83/60/d497a310bde3f01cb805196ac61b7ad6dc5dcf8dce66634dc34364b20b4f/lazy_loader-0.4-py3-none-any.whl", hash = "sha256:342aa8e14d543a154047afb4ba8ef17f5563baad3fc610d7b15b213b0f119efc", size = 12097, upload-time = "2024-04-05T13:03:10.514Z" }, -] - [[package]] name = "lmdb" version = "1.7.5" @@ -2515,12 +2515,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, ] +[[package]] +name = "nvidia-cuda-runtime" +version = "13.0.88" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/66/7020819766edab221110be12f1f691e62efba29f667992693aea6c10649f/nvidia_cuda_runtime-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:8ba16a33998ee3ee2bf11c8f13e36e3ff86a1d6ad165ffeee863d6eabd0b5f67", size = 2259852, upload-time = "2025-09-04T08:25:58.873Z" }, + { url = "https://files.pythonhosted.org/packages/b1/1d/9ea7b0d420efcae3c546feb3ca14334c0d710ac19c156ac48fce07c3d6ae/nvidia_cuda_runtime-13.0.88-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:422a38f7c739ecbdd3a4b6ec7ed0a1d74db43d8e9bc26685f3d845638922b178", size = 2243181, upload-time = "2025-09-04T08:26:06.897Z" }, + { url = "https://files.pythonhosted.org/packages/b3/76/d71097672718defe63af4702b9b0d9c32f9f8df4582ce1a1b891bc65b93b/nvidia_cuda_runtime-13.0.88-py3-none-win_amd64.whl", hash = "sha256:0156cf8f3453025938aab80bcca909506f4e82f1584a17736df7edeba882cc02", size = 2926776, upload-time = "2025-09-04T08:37:58.121Z" }, +] + [[package]] name = "nvidia-cuda-runtime-cu12" version = "12.8.90" source = { registry = "https://pypi.org/simple" } wheels = [ + { url = "https://files.pythonhosted.org/packages/7c/75/f865a3b236e4647605ea34cc450900854ba123834a5f1598e160b9530c3a/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:52bf7bbee900262ffefe5e9d5a2a69a30d97e2bc5bb6cc866688caa976966e3d", size = 965265, upload-time = "2025-03-07T01:39:43.533Z" }, { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, + { url = "https://files.pythonhosted.org/packages/30/a5/a515b7600ad361ea14bfa13fb4d6687abf500adc270f19e89849c0590492/nvidia_cuda_runtime_cu12-12.8.90-py3-none-win_amd64.whl", hash = "sha256:c0c6027f01505bfed6c3b21ec546f69c687689aad5f1a377554bc6ca4aa993a8", size = 944318, upload-time = "2025-03-07T01:51:01.794Z" }, ] [[package]] @@ -2630,6 +2642,61 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, ] +[[package]] +name = "onnx" +version = "1.19.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" }, + { name = "numpy" }, + { name = "protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5b/bf/b0a63ee9f3759dcd177b28c6f2cb22f2aecc6d9b3efecaabc298883caa5f/onnx-1.19.0.tar.gz", hash = "sha256:aa3f70b60f54a29015e41639298ace06adf1dd6b023b9b30f1bca91bb0db9473", size = 11949859, upload-time = "2025-08-27T02:34:27.107Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/b3/8a6f3b05d18dffdc7c18839bd829587c826c8513f4bdbe21ddf37dacce50/onnx-1.19.0-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:e927d745939d590f164e43c5aec7338c5a75855a15130ee795f492fc3a0fa565", size = 18310869, upload-time = "2025-08-27T02:32:47.346Z" }, + { url = "https://files.pythonhosted.org/packages/b9/92/550d6155ab3f2c00e95add1726397c95b4b79d6eb4928d049ff591ad4c84/onnx-1.19.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:c6cdcb237c5c4202463bac50417c5a7f7092997a8469e8b7ffcd09f51de0f4a9", size = 18028144, upload-time = "2025-08-27T02:32:50.306Z" }, + { url = "https://files.pythonhosted.org/packages/79/21/9bcc715ea6d9aab3f6c583bfc59504a14777e39e0591030e7345f4e40315/onnx-1.19.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ed0b85a33deacb65baffe6ca4ce91adf2bb906fa2dee3856c3c94e163d2eb563", size = 18200923, upload-time = "2025-08-27T02:32:54.325Z" }, + { url = "https://files.pythonhosted.org/packages/c8/90/3a6f0741ff22270e2f4b741f440ab68ba5525ebc94775cd6f2c01f531374/onnx-1.19.0-cp310-cp310-win32.whl", hash = "sha256:89a9cefe75547aec14a796352c2243e36793bbbcb642d8897118595ab0c2395b", size = 16332097, upload-time = "2025-08-27T02:32:56.997Z" }, + { url = "https://files.pythonhosted.org/packages/4c/4c/ef61d359865712803d488672607023d36bfcd21fa008d8dc1d6ee8e8b23c/onnx-1.19.0-cp310-cp310-win_amd64.whl", hash = "sha256:a16a82bfdf4738691c0a6eda5293928645ab8b180ab033df84080817660b5e66", size = 16451402, upload-time = "2025-08-27T02:33:00.534Z" }, + { url = "https://files.pythonhosted.org/packages/db/5c/b959b17608cfb6ccf6359b39fe56a5b0b7d965b3d6e6a3c0add90812c36e/onnx-1.19.0-cp311-cp311-macosx_12_0_universal2.whl", hash = "sha256:206f00c47b85b5c7af79671e3307147407991a17994c26974565aadc9e96e4e4", size = 18312580, upload-time = "2025-08-27T02:33:03.081Z" }, + { url = "https://files.pythonhosted.org/packages/2c/ee/ac052bbbc832abe0debb784c2c57f9582444fb5f51d63c2967fd04432444/onnx-1.19.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:4d7bee94abaac28988b50da675ae99ef8dd3ce16210d591fbd0b214a5930beb3", size = 18029165, upload-time = "2025-08-27T02:33:05.771Z" }, + { url = "https://files.pythonhosted.org/packages/5c/c9/8687ba0948d46fd61b04e3952af9237883bbf8f16d716e7ed27e688d73b8/onnx-1.19.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7730b96b68c0c354bbc7857961bb4909b9aaa171360a8e3708d0a4c749aaadeb", size = 18202125, upload-time = "2025-08-27T02:33:09.325Z" }, + { url = "https://files.pythonhosted.org/packages/e2/16/6249c013e81bd689f46f96c7236d7677f1af5dd9ef22746716b48f10e506/onnx-1.19.0-cp311-cp311-win32.whl", hash = "sha256:7cb7a3ad8059d1a0dfdc5e0a98f71837d82002e441f112825403b137227c2c97", size = 16332738, upload-time = "2025-08-27T02:33:12.448Z" }, + { url = "https://files.pythonhosted.org/packages/6a/28/34a1e2166e418c6a78e5c82e66f409d9da9317832f11c647f7d4e23846a6/onnx-1.19.0-cp311-cp311-win_amd64.whl", hash = "sha256:d75452a9be868bd30c3ef6aa5991df89bbfe53d0d90b2325c5e730fbd91fff85", size = 16452303, upload-time = "2025-08-27T02:33:15.176Z" }, + { url = "https://files.pythonhosted.org/packages/e6/b7/639664626e5ba8027860c4d2a639ee02b37e9c322215c921e9222513c3aa/onnx-1.19.0-cp311-cp311-win_arm64.whl", hash = "sha256:23c7959370d7b3236f821e609b0af7763cff7672a758e6c1fc877bac099e786b", size = 16425340, upload-time = "2025-08-27T02:33:17.78Z" }, + { url = "https://files.pythonhosted.org/packages/0d/94/f56f6ca5e2f921b28c0f0476705eab56486b279f04e1d568ed64c14e7764/onnx-1.19.0-cp312-cp312-macosx_12_0_universal2.whl", hash = "sha256:61d94e6498ca636756f8f4ee2135708434601b2892b7c09536befb19bc8ca007", size = 18322331, upload-time = "2025-08-27T02:33:20.373Z" }, + { url = "https://files.pythonhosted.org/packages/c8/00/8cc3f3c40b54b28f96923380f57c9176872e475face726f7d7a78bd74098/onnx-1.19.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:224473354462f005bae985c72028aaa5c85ab11de1b71d55b06fdadd64a667dd", size = 18027513, upload-time = "2025-08-27T02:33:23.44Z" }, + { url = "https://files.pythonhosted.org/packages/61/90/17c4d2566fd0117a5e412688c9525f8950d467f477fbd574e6b32bc9cb8d/onnx-1.19.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ae475c85c89bc4d1f16571006fd21a3e7c0e258dd2c091f6e8aafb083d1ed9b", size = 18202278, upload-time = "2025-08-27T02:33:26.103Z" }, + { url = "https://files.pythonhosted.org/packages/bc/6e/a9383d9cf6db4ac761a129b081e9fa5d0cd89aad43cf1e3fc6285b915c7d/onnx-1.19.0-cp312-cp312-win32.whl", hash = "sha256:323f6a96383a9cdb3960396cffea0a922593d221f3929b17312781e9f9b7fb9f", size = 16333080, upload-time = "2025-08-27T02:33:28.559Z" }, + { url = "https://files.pythonhosted.org/packages/a7/2e/3ff480a8c1fa7939662bdc973e41914add2d4a1f2b8572a3c39c2e4982e5/onnx-1.19.0-cp312-cp312-win_amd64.whl", hash = "sha256:50220f3499a499b1a15e19451a678a58e22ad21b34edf2c844c6ef1d9febddc2", size = 16453927, upload-time = "2025-08-27T02:33:31.177Z" }, + { url = "https://files.pythonhosted.org/packages/57/37/ad500945b1b5c154fe9d7b826b30816ebd629d10211ea82071b5bcc30aa4/onnx-1.19.0-cp312-cp312-win_arm64.whl", hash = "sha256:efb768299580b786e21abe504e1652ae6189f0beed02ab087cd841cb4bb37e43", size = 16426022, upload-time = "2025-08-27T02:33:33.515Z" }, + { url = "https://files.pythonhosted.org/packages/be/29/d7b731f63d243f815d9256dce0dca3c151dcaa1ac59f73e6ee06c9afbe91/onnx-1.19.0-cp313-cp313-macosx_12_0_universal2.whl", hash = "sha256:9aed51a4b01acc9ea4e0fe522f34b2220d59e9b2a47f105ac8787c2e13ec5111", size = 18322412, upload-time = "2025-08-27T02:33:36.723Z" }, + { url = "https://files.pythonhosted.org/packages/58/f5/d3106becb42cb374f0e17ff4c9933a97f1ee1d6a798c9452067f7d3ff61b/onnx-1.19.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce2cdc3eb518bb832668c4ea9aeeda01fbaa59d3e8e5dfaf7aa00f3d37119404", size = 18026565, upload-time = "2025-08-27T02:33:39.493Z" }, + { url = "https://files.pythonhosted.org/packages/83/fa/b086d17bab3900754c7ffbabfb244f8e5e5da54a34dda2a27022aa2b373b/onnx-1.19.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8b546bd7958734b6abcd40cfede3d025e9c274fd96334053a288ab11106bd0aa", size = 18202077, upload-time = "2025-08-27T02:33:42.115Z" }, + { url = "https://files.pythonhosted.org/packages/35/f2/5e2dfb9d4cf873f091c3f3c6d151f071da4295f9893fbf880f107efe3447/onnx-1.19.0-cp313-cp313-win32.whl", hash = "sha256:03086bffa1cf5837430cf92f892ca0cd28c72758d8905578c2bf8ffaf86c6743", size = 16333198, upload-time = "2025-08-27T02:33:45.172Z" }, + { url = "https://files.pythonhosted.org/packages/79/67/b3751a35c2522f62f313156959575619b8fa66aa883db3adda9d897d8eb2/onnx-1.19.0-cp313-cp313-win_amd64.whl", hash = "sha256:1715b51eb0ab65272e34ef51cb34696160204b003566cd8aced2ad20a8f95cb8", size = 16453836, upload-time = "2025-08-27T02:33:47.779Z" }, + { url = "https://files.pythonhosted.org/packages/14/b9/1df85effc960fbbb90bb7bc36eb3907c676b104bc2f88bce022bcfdaef63/onnx-1.19.0-cp313-cp313-win_arm64.whl", hash = "sha256:6bf5acdb97a3ddd6e70747d50b371846c313952016d0c41133cbd8f61b71a8d5", size = 16425877, upload-time = "2025-08-27T02:33:50.357Z" }, + { url = "https://files.pythonhosted.org/packages/23/2b/089174a1427be9149f37450f8959a558ba20f79fca506ba461d59379d3a1/onnx-1.19.0-cp313-cp313t-macosx_12_0_universal2.whl", hash = "sha256:46cf29adea63e68be0403c68de45ba1b6acc9bb9592c5ddc8c13675a7c71f2cb", size = 18348546, upload-time = "2025-08-27T02:33:56.132Z" }, + { url = "https://files.pythonhosted.org/packages/c0/d6/3458f0e3a9dc7677675d45d7d6528cb84ad321c8670cc10c69b32c3e03da/onnx-1.19.0-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:246f0de1345498d990a443d55a5b5af5101a3e25a05a2c3a5fe8b7bd7a7d0707", size = 18033067, upload-time = "2025-08-27T02:33:58.661Z" }, + { url = "https://files.pythonhosted.org/packages/e4/16/6e4130e1b4b29465ee1fb07d04e8d6f382227615c28df8f607ba50909e2a/onnx-1.19.0-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ae0d163ffbc250007d984b8dd692a4e2e4506151236b50ca6e3560b612ccf9ff", size = 18205741, upload-time = "2025-08-27T02:34:01.538Z" }, + { url = "https://files.pythonhosted.org/packages/fe/d8/f64d010fd024b2a2b11ce0c4ee179e4f8f6d4ccc95f8184961c894c22af1/onnx-1.19.0-cp313-cp313t-win_amd64.whl", hash = "sha256:7c151604c7cca6ae26161c55923a7b9b559df3344938f93ea0074d2d49e7fe78", size = 16453839, upload-time = "2025-08-27T02:34:06.515Z" }, + { url = "https://files.pythonhosted.org/packages/67/ec/8761048eabef4dad55af4c002c672d139b9bd47c3616abaed642a1710063/onnx-1.19.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:236bc0e60d7c0f4159300da639953dd2564df1c195bce01caba172a712e75af4", size = 18027605, upload-time = "2025-08-27T02:34:08.962Z" }, +] + +[[package]] +name = "onnx-graphsurgeon" +version = "0.5.8" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "onnx" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/53/98334c4f64a9e289a8cb48f5e7966b8ff015414d0bf26587cf46d764f1d8/onnx_graphsurgeon-0.5.8-py2.py3-none-any.whl", hash = "sha256:6f611ea29a8e4740fbab1aae52bf4c40b8b9918f8459058d20b99acc79fce121", size = 57923, upload-time = "2025-04-10T18:49:24.483Z" }, +] + [[package]] name = "opencv-contrib-python" version = "4.11.0.86" @@ -2803,6 +2870,14 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "polygraphy" +version = "0.49.26" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a1/25/e92af58cca8f8f8833d3346506a186938148b7156671c8154d94b8985d35/polygraphy-0.49.26-py2.py3-none-any.whl", hash = "sha256:5787d218a133163b42c92800134afaba6b266127646efb77416a9530137d1a45", size = 372849, upload-time = "2025-07-16T18:26:21.446Z" }, +] + [[package]] name = "pre-commit" version = "4.3.0" @@ -3551,48 +3626,6 @@ wheels = [ { url = "https://github.com/woct0rdho/SageAttention/releases/download/v2.2.0-windows.post3/sageattention-2.2.0+cu128torch2.8.0.post3-cp39-abi3-win_amd64.whl", hash = "sha256:7dabcd00e63229b28f046c5a69ec37cf4756afb375dbadd1975dadec045ae21c" }, ] -[[package]] -name = "scikit-image" -version = "0.25.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "imageio" }, - { name = "lazy-loader" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "networkx", version = "3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "scipy", version = "1.16.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "tifffile", version = "2025.5.10", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "tifffile", version = "2025.12.12", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/c7/a8/3c0f256012b93dd2cb6fda9245e9f4bff7dc0486880b248005f15ea2255e/scikit_image-0.25.2.tar.gz", hash = "sha256:e5a37e6cd4d0c018a7a55b9d601357e3382826d3888c10d0213fc63bff977dde", size = 22693594, upload-time = "2025-02-18T18:05:24.538Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/11/cb/016c63f16065c2d333c8ed0337e18a5cdf9bc32d402e4f26b0db362eb0e2/scikit_image-0.25.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d3278f586793176599df6a4cf48cb6beadae35c31e58dc01a98023af3dc31c78", size = 13988922, upload-time = "2025-02-18T18:04:11.069Z" }, - { url = "https://files.pythonhosted.org/packages/30/ca/ff4731289cbed63c94a0c9a5b672976603118de78ed21910d9060c82e859/scikit_image-0.25.2-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:5c311069899ce757d7dbf1d03e32acb38bb06153236ae77fcd820fd62044c063", size = 13192698, upload-time = "2025-02-18T18:04:15.362Z" }, - { url = "https://files.pythonhosted.org/packages/39/6d/a2aadb1be6d8e149199bb9b540ccde9e9622826e1ab42fe01de4c35ab918/scikit_image-0.25.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:be455aa7039a6afa54e84f9e38293733a2622b8c2fb3362b822d459cc5605e99", size = 14153634, upload-time = "2025-02-18T18:04:18.496Z" }, - { url = "https://files.pythonhosted.org/packages/96/08/916e7d9ee4721031b2f625db54b11d8379bd51707afaa3e5a29aecf10bc4/scikit_image-0.25.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a4c464b90e978d137330be433df4e76d92ad3c5f46a22f159520ce0fdbea8a09", size = 14767545, upload-time = "2025-02-18T18:04:22.556Z" }, - { url = "https://files.pythonhosted.org/packages/5f/ee/c53a009e3997dda9d285402f19226fbd17b5b3cb215da391c4ed084a1424/scikit_image-0.25.2-cp310-cp310-win_amd64.whl", hash = "sha256:60516257c5a2d2f74387c502aa2f15a0ef3498fbeaa749f730ab18f0a40fd054", size = 12812908, upload-time = "2025-02-18T18:04:26.364Z" }, - { url = "https://files.pythonhosted.org/packages/c4/97/3051c68b782ee3f1fb7f8f5bb7d535cf8cb92e8aae18fa9c1cdf7e15150d/scikit_image-0.25.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f4bac9196fb80d37567316581c6060763b0f4893d3aca34a9ede3825bc035b17", size = 14003057, upload-time = "2025-02-18T18:04:30.395Z" }, - { url = "https://files.pythonhosted.org/packages/19/23/257fc696c562639826065514d551b7b9b969520bd902c3a8e2fcff5b9e17/scikit_image-0.25.2-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d989d64ff92e0c6c0f2018c7495a5b20e2451839299a018e0e5108b2680f71e0", size = 13180335, upload-time = "2025-02-18T18:04:33.449Z" }, - { url = "https://files.pythonhosted.org/packages/ef/14/0c4a02cb27ca8b1e836886b9ec7c9149de03053650e9e2ed0625f248dd92/scikit_image-0.25.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b2cfc96b27afe9a05bc92f8c6235321d3a66499995675b27415e0d0c76625173", size = 14144783, upload-time = "2025-02-18T18:04:36.594Z" }, - { url = "https://files.pythonhosted.org/packages/dd/9b/9fb556463a34d9842491d72a421942c8baff4281025859c84fcdb5e7e602/scikit_image-0.25.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:24cc986e1f4187a12aa319f777b36008764e856e5013666a4a83f8df083c2641", size = 14785376, upload-time = "2025-02-18T18:04:39.856Z" }, - { url = "https://files.pythonhosted.org/packages/de/ec/b57c500ee85885df5f2188f8bb70398481393a69de44a00d6f1d055f103c/scikit_image-0.25.2-cp311-cp311-win_amd64.whl", hash = "sha256:b4f6b61fc2db6340696afe3db6b26e0356911529f5f6aee8c322aa5157490c9b", size = 12791698, upload-time = "2025-02-18T18:04:42.868Z" }, - { url = "https://files.pythonhosted.org/packages/35/8c/5df82881284459f6eec796a5ac2a0a304bb3384eec2e73f35cfdfcfbf20c/scikit_image-0.25.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8db8dd03663112783221bf01ccfc9512d1cc50ac9b5b0fe8f4023967564719fb", size = 13986000, upload-time = "2025-02-18T18:04:47.156Z" }, - { url = "https://files.pythonhosted.org/packages/ce/e6/93bebe1abcdce9513ffec01d8af02528b4c41fb3c1e46336d70b9ed4ef0d/scikit_image-0.25.2-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:483bd8cc10c3d8a7a37fae36dfa5b21e239bd4ee121d91cad1f81bba10cfb0ed", size = 13235893, upload-time = "2025-02-18T18:04:51.049Z" }, - { url = "https://files.pythonhosted.org/packages/53/4b/eda616e33f67129e5979a9eb33c710013caa3aa8a921991e6cc0b22cea33/scikit_image-0.25.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9d1e80107bcf2bf1291acfc0bf0425dceb8890abe9f38d8e94e23497cbf7ee0d", size = 14178389, upload-time = "2025-02-18T18:04:54.245Z" }, - { url = "https://files.pythonhosted.org/packages/6b/b5/b75527c0f9532dd8a93e8e7cd8e62e547b9f207d4c11e24f0006e8646b36/scikit_image-0.25.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a17e17eb8562660cc0d31bb55643a4da996a81944b82c54805c91b3fe66f4824", size = 15003435, upload-time = "2025-02-18T18:04:57.586Z" }, - { url = "https://files.pythonhosted.org/packages/34/e3/49beb08ebccda3c21e871b607c1cb2f258c3fa0d2f609fed0a5ba741b92d/scikit_image-0.25.2-cp312-cp312-win_amd64.whl", hash = "sha256:bdd2b8c1de0849964dbc54037f36b4e9420157e67e45a8709a80d727f52c7da2", size = 12899474, upload-time = "2025-02-18T18:05:01.166Z" }, - { url = "https://files.pythonhosted.org/packages/e6/7c/9814dd1c637f7a0e44342985a76f95a55dd04be60154247679fd96c7169f/scikit_image-0.25.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:7efa888130f6c548ec0439b1a7ed7295bc10105458a421e9bf739b457730b6da", size = 13921841, upload-time = "2025-02-18T18:05:03.963Z" }, - { url = "https://files.pythonhosted.org/packages/84/06/66a2e7661d6f526740c309e9717d3bd07b473661d5cdddef4dd978edab25/scikit_image-0.25.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:dd8011efe69c3641920614d550f5505f83658fe33581e49bed86feab43a180fc", size = 13196862, upload-time = "2025-02-18T18:05:06.986Z" }, - { url = "https://files.pythonhosted.org/packages/4e/63/3368902ed79305f74c2ca8c297dfeb4307269cbe6402412668e322837143/scikit_image-0.25.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:28182a9d3e2ce3c2e251383bdda68f8d88d9fff1a3ebe1eb61206595c9773341", size = 14117785, upload-time = "2025-02-18T18:05:10.69Z" }, - { url = "https://files.pythonhosted.org/packages/cd/9b/c3da56a145f52cd61a68b8465d6a29d9503bc45bc993bb45e84371c97d94/scikit_image-0.25.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b8abd3c805ce6944b941cfed0406d88faeb19bab3ed3d4b50187af55cf24d147", size = 14977119, upload-time = "2025-02-18T18:05:13.871Z" }, - { url = "https://files.pythonhosted.org/packages/8a/97/5fcf332e1753831abb99a2525180d3fb0d70918d461ebda9873f66dcc12f/scikit_image-0.25.2-cp313-cp313-win_amd64.whl", hash = "sha256:64785a8acefee460ec49a354706db0b09d1f325674107d7fa3eadb663fb56d6f", size = 12885116, upload-time = "2025-02-18T18:05:17.844Z" }, - { url = "https://files.pythonhosted.org/packages/10/cc/75e9f17e3670b5ed93c32456fda823333c6279b144cd93e2c03aa06aa472/scikit_image-0.25.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:330d061bd107d12f8d68f1d611ae27b3b813b8cdb0300a71d07b1379178dd4cd", size = 13862801, upload-time = "2025-02-18T18:05:20.783Z" }, -] - [[package]] name = "scipy" version = "1.15.3" @@ -3906,52 +3939,85 @@ wheels = [ ] [[package]] -name = "tifffile" -version = "2025.5.10" +name = "tensorrt" +version = "10.14.1.48.post1" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'linux'", - "python_full_version < '3.11' and platform_machine != 'aarch64' and sys_platform == 'linux'", - "python_full_version < '3.11' and sys_platform == 'win32'", - "python_full_version < '3.11' and sys_platform == 'darwin'", - "python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +dependencies = [ + { name = "tensorrt-cu13", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/cd/d7/1455a247ecc3a85929cd3136e1ebcb0c2afbef3fb893cb47efbebec95034/tensorrt-10.14.1.48.post1.tar.gz", hash = "sha256:5aed87c541fcef881f874fb5724eadbe291851fe2152c962b9cba57cafd4872b", size = 16725, upload-time = "2025-11-07T17:14:17.028Z" } + +[[package]] +name = "tensorrt-cu12" +version = "10.13.3.9" +source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "python_full_version < '3.11'" }, + { name = "tensorrt-cu12-bindings", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tensorrt-cu12-libs", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/44/d0/18fed0fc0916578a4463f775b0fbd9c5fed2392152d039df2fb533bfdd5d/tifffile-2025.5.10.tar.gz", hash = "sha256:018335d34283aa3fd8c263bae5c3c2b661ebc45548fde31504016fcae7bf1103", size = 365290, upload-time = "2025-05-10T19:22:34.386Z" } +sdist = { url = "https://files.pythonhosted.org/packages/1a/aa/5a630ac02d05f7bbc774bd11b2bc9a3d31aa0bd7d38b4cfeef06733544b3/tensorrt_cu12-10.13.3.9.tar.gz", hash = "sha256:b507abc536ceeaac3b2e6527bd65100bef19000d204e7072db6fe85f9f859a84", size = 18175, upload-time = "2025-09-05T19:03:16.524Z" } + +[[package]] +name = "tensorrt-cu12-bindings" +version = "10.13.3.9" +source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/06/bd0a6097da704a7a7c34a94cfd771c3ea3c2f405dd214e790d22c93f6be1/tifffile-2025.5.10-py3-none-any.whl", hash = "sha256:e37147123c0542d67bc37ba5cdd67e12ea6fbe6e86c52bee037a9eb6a064e5ad", size = 226533, upload-time = "2025-05-10T19:22:27.279Z" }, + { url = "https://files.pythonhosted.org/packages/03/b4/768d8e1d4cd949422ede434397d5f68b2281a2a269602257ab2922674f36/tensorrt_cu12_bindings-10.13.3.9-cp310-none-manylinux_2_28_x86_64.whl", hash = "sha256:3ca2635394789dd62346ff5ee3f02def53cee2cd21247cc342461dc721893a1e", size = 1181470, upload-time = "2025-09-05T19:03:39.63Z" }, + { url = "https://files.pythonhosted.org/packages/0e/3c/ff7b3cea9181b4d5c01f135af3cf66b547f9988d7b04ac0657abbdc2b300/tensorrt_cu12_bindings-10.13.3.9-cp310-none-win_amd64.whl", hash = "sha256:cb69c187370c3cfd1e3176acf5dcc8c3e7e9b526046e4646a7372b9cde375f70", size = 874901, upload-time = "2025-09-05T18:37:36.624Z" }, + { url = "https://files.pythonhosted.org/packages/03/b4/9040435df7db4491b6a9dbac68dfd39ba8de9b79e06cc877c064fe97c2a6/tensorrt_cu12_bindings-10.13.3.9-cp311-none-manylinux_2_28_x86_64.whl", hash = "sha256:33aa9d104c04b008871645a97bd8a1a02279cf498d251df5e767d9d5183e0646", size = 1181763, upload-time = "2025-09-05T19:04:50.567Z" }, + { url = "https://files.pythonhosted.org/packages/67/56/66b1b9e75cbb13e706cdcc14b4b7801b0940fcb0ae10759b67067e8208f4/tensorrt_cu12_bindings-10.13.3.9-cp311-none-win_amd64.whl", hash = "sha256:f87962f5ea2ae67437a78a5aa08b00345bfb1f19ec815e2566da26e4223c72ad", size = 877325, upload-time = "2025-09-05T18:45:53.776Z" }, + { url = "https://files.pythonhosted.org/packages/58/cc/8eddae7c3b0bd6d51290624c59f05b51ec3ded917b3689f1975f8dcfc601/tensorrt_cu12_bindings-10.13.3.9-cp312-none-manylinux_2_28_x86_64.whl", hash = "sha256:193174cd0aa5ec7d63943d8d3f84ca9ec525477a9db3a9100bf2f77071b3fae8", size = 1184541, upload-time = "2025-09-05T19:04:03.708Z" }, + { url = "https://files.pythonhosted.org/packages/f5/e4/19aff120be055418bac8f012cf38c682e540f983b55a5037d98c00b8b65c/tensorrt_cu12_bindings-10.13.3.9-cp312-none-win_amd64.whl", hash = "sha256:8dd3ad04874563f28a7ab5e6db394082b50351006da6c834253d864d777591f3", size = 877013, upload-time = "2025-09-05T18:30:47.278Z" }, + { url = "https://files.pythonhosted.org/packages/71/1c/07281273e6db001f726e365b70d7d648df7da0af2da68eabe83f58622f7c/tensorrt_cu12_bindings-10.13.3.9-cp313-none-manylinux_2_28_x86_64.whl", hash = "sha256:90c7d502682d868a918665f1035b4acb07254bcd76cc551febea9a7009ee6269", size = 1184451, upload-time = "2025-09-05T19:04:26.969Z" }, + { url = "https://files.pythonhosted.org/packages/76/ba/5e20932b15ed85bb402203e5afd9859eb1bff494024918c660776d9ec737/tensorrt_cu12_bindings-10.13.3.9-cp313-none-win_amd64.whl", hash = "sha256:5b4bdc2bab9f3109bb44425e4789dcb540f0cba22af4cb25a36aec523d65b508", size = 876995, upload-time = "2025-09-05T18:46:17.103Z" }, ] [[package]] -name = "tifffile" -version = "2025.12.12" +name = "tensorrt-cu12-libs" +version = "10.13.3.9" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'linux'", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", - "python_full_version >= '3.13' and platform_machine != 'aarch64' and sys_platform == 'linux'", - "python_full_version == '3.12.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'linux'", - "python_full_version == '3.11.*' and platform_machine != 'aarch64' and sys_platform == 'linux'", - "python_full_version >= '3.13' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'win32'", - "python_full_version >= '3.13' and sys_platform == 'darwin'", - "python_full_version == '3.12.*' and sys_platform == 'darwin'", - "python_full_version >= '3.13' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.12.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", - "python_full_version == '3.11.*' and sys_platform == 'darwin'", - "python_full_version == '3.11.*' and sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32'", +dependencies = [ + { name = "nvidia-cuda-runtime-cu12", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/e1/d3/82b9a603cac94975dbb25d33c7ffb1dba6851980f2dc630f2bf9caf0bffb/tensorrt_cu12_libs-10.13.3.9.tar.gz", hash = "sha256:99d11dd42a27705bcb2ea4e2e0f76817a2a9403d892ec1272c502bd2722ebbd9", size = 703, upload-time = "2025-09-05T18:03:45.78Z" } + +[[package]] +name = "tensorrt-cu13" +version = "10.14.1.48.post1" +source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "python_full_version >= '3.11'" }, + { name = "tensorrt-cu13-bindings", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "tensorrt-cu13-libs", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/31/b9/4253513a66f0a836ec3a5104266cf73f7812bfbbcda9d87d8c0e93b28293/tifffile-2025.12.12.tar.gz", hash = "sha256:97e11fd6b1d8dc971896a098c841d9cd4e6eb958ac040dd6fb8b332c3f7288b6", size = 373597, upload-time = "2025-12-13T03:42:53.765Z" } +sdist = { url = "https://files.pythonhosted.org/packages/9f/c3/2905ffaf789b53d68c84930d081e14cb96b7516a5c13f8e66c5545a47464/tensorrt_cu13-10.14.1.48.post1.tar.gz", hash = "sha256:f667aee4b4b908520820b824ada7b7168f29e2659aacee2d5ae42139c6cb884e", size = 18328, upload-time = "2025-11-07T17:14:39.291Z" } + +[[package]] +name = "tensorrt-cu13-bindings" +version = "10.14.1.48.post1" +source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d5/5c/e444e1b024a519e488326525f0c154396c6b16baff17e00623f2c21dfc42/tifffile-2025.12.12-py3-none-any.whl", hash = "sha256:e3e3f1290ec6741ca248a5b5a997125209b5c2962f6bd9aef01ea9352c25d0ee", size = 232132, upload-time = "2025-12-13T03:42:52.072Z" }, + { url = "https://files.pythonhosted.org/packages/2e/dd/be2121fabc3a3bc5bca2d91b76630a74fc32fa43ae186584a865dc97d2d9/tensorrt_cu13_bindings-10.14.1.48.post1-cp310-none-manylinux_2_28_x86_64.whl", hash = "sha256:d8f13ac102187c4a45d11cc671f7f7e59a01ec4e4f1c707fb748d30042a49416", size = 879656, upload-time = "2025-11-07T17:16:03.153Z" }, + { url = "https://files.pythonhosted.org/packages/25/98/b42abb0a4ddeeb9e77278c5276db9136bb34e95f80e41171b730c87b8752/tensorrt_cu13_bindings-10.14.1.48.post1-cp310-none-manylinux_2_35_aarch64.whl", hash = "sha256:35e34a14411a24fde813471e4ef1882615992ee1e093c5c254f4fca070b98973", size = 958832, upload-time = "2025-11-07T18:28:04.98Z" }, + { url = "https://files.pythonhosted.org/packages/d8/d5/903808b18cc08ba0d8e4a6a469b08244c7c2b36291d95fc09058ec0bb39f/tensorrt_cu13_bindings-10.14.1.48.post1-cp310-none-win_amd64.whl", hash = "sha256:8f8ee0d11ee6c3f53d80c768dd9a18b75f7325babb3022b549236bb524fc807a", size = 766043, upload-time = "2025-11-07T18:57:50.67Z" }, + { url = "https://files.pythonhosted.org/packages/1d/f6/37a3a59fc925b9d9cdd44375468b934adb4906482a42b2228a1967d6146c/tensorrt_cu13_bindings-10.14.1.48.post1-cp311-none-manylinux_2_28_x86_64.whl", hash = "sha256:30e220ec5cd049709270e9c835d2d69ae5b05d5c269fc067e27155d2b1ba8ff5", size = 879697, upload-time = "2025-11-07T17:15:40.107Z" }, + { url = "https://files.pythonhosted.org/packages/3c/ce/0df4587d1690d615c0fa2a22f7e9c32a5547f66f370c4a0cd789c7ebf965/tensorrt_cu13_bindings-10.14.1.48.post1-cp311-none-manylinux_2_35_aarch64.whl", hash = "sha256:31a6f3d8217fa52533aecbe15073114fa26c2f7835d8d3c8b5371b22899c912a", size = 959312, upload-time = "2025-11-07T18:26:17.172Z" }, + { url = "https://files.pythonhosted.org/packages/de/2f/a591eb4c6251705896a579c5b0b5b4210d7a547b0fbf867da0ae9221ae1e/tensorrt_cu13_bindings-10.14.1.48.post1-cp311-none-win_amd64.whl", hash = "sha256:83c58af9c32e527567cf4927957126a487d8b89604822c2be2522bf3e348228b", size = 766024, upload-time = "2025-11-07T18:59:17.105Z" }, + { url = "https://files.pythonhosted.org/packages/65/5a/4101e26fffadc08630dd4850045fdf4ec3d58e0cda6fa67935c0daede5c4/tensorrt_cu13_bindings-10.14.1.48.post1-cp312-none-manylinux_2_28_x86_64.whl", hash = "sha256:1a098123ddf9e69629d66320666efbb2e237db76cd4faf715b7377471b9f4275", size = 879890, upload-time = "2025-11-07T17:17:27.641Z" }, + { url = "https://files.pythonhosted.org/packages/a1/b1/d8098e73df5e43610408b2726280d5cf36f3b9b2a45250aa1545c8970606/tensorrt_cu13_bindings-10.14.1.48.post1-cp312-none-manylinux_2_35_aarch64.whl", hash = "sha256:e8db2ddd753097b23e13bd29b61e24431db0724f6f143e917ef107a4f58dff43", size = 959698, upload-time = "2025-11-07T18:28:56.083Z" }, + { url = "https://files.pythonhosted.org/packages/82/75/9032bbef64e0e1bcf35171e6524a4a027e3620ebc0bd8310bfb3b74fd221/tensorrt_cu13_bindings-10.14.1.48.post1-cp312-none-win_amd64.whl", hash = "sha256:c762c813bcaddf88e0aae43f8fad2b8637d1088818e9684de2f5e17f773592b5", size = 768710, upload-time = "2025-11-07T18:59:40.181Z" }, + { url = "https://files.pythonhosted.org/packages/2d/fc/c2aa0191c579f0ce79a782bc6859673a6baa2fd1fc889ba4f6a21660480f/tensorrt_cu13_bindings-10.14.1.48.post1-cp313-none-manylinux_2_28_x86_64.whl", hash = "sha256:51928aa2c11c08d4aa5425c60c7d866c088cffccd0f20a4c8addca9de930c603", size = 879778, upload-time = "2025-11-07T17:17:50.286Z" }, + { url = "https://files.pythonhosted.org/packages/90/5d/9c811ab2f8b9f856d2b9947467763819da4a51fa7d851aed7a356490ccdc/tensorrt_cu13_bindings-10.14.1.48.post1-cp313-none-manylinux_2_35_aarch64.whl", hash = "sha256:4984a8f841843a3b23c743a24c5d517cf2e3b5eb0c7186bd96791f011fd1dab1", size = 959675, upload-time = "2025-11-07T18:29:49.568Z" }, + { url = "https://files.pythonhosted.org/packages/c7/ec/60e294d729515dd67fad47789abe8c41c81ed2f2a72f4d4233317c11ae71/tensorrt_cu13_bindings-10.14.1.48.post1-cp313-none-win_amd64.whl", hash = "sha256:dd8bc2507717a6894d5ef26a4acad003657f8c8781370f45c5be19623d5587ae", size = 768834, upload-time = "2025-11-07T18:58:15.436Z" }, +] + +[[package]] +name = "tensorrt-cu13-libs" +version = "10.14.1.48.post1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-toolkit", extra = ["cudart"], marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, ] +sdist = { url = "https://files.pythonhosted.org/packages/38/7f/98c80315da5b2cf19bd2a47b4b676d66a0cebc95d1b5d277349147b082e4/tensorrt_cu13_libs-10.14.1.48.post1.tar.gz", hash = "sha256:8ffa01f315276bff7c7dd01bf11415d0c8c21c45012abd42f2878153ba312d76", size = 726, upload-time = "2025-11-07T17:21:49.844Z" } [[package]] name = "tokenizers" From c1ee453219860bcab964e3eb730192b67d8ee288 Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Wed, 17 Dec 2025 00:54:01 -0500 Subject: [PATCH 3/7] PersonaLive TensorRT testing, developement and profiling. Signed-off-by: BuffMcBigHuge --- .../core/pipelines/personalive/pipeline.py | 355 +++++++---- .../personalive/tensorrt/__init__.py | 32 + .../pipelines/personalive/tensorrt/convert.py | 124 ++-- .../personalive/tensorrt/engine_model.py | 394 ++++++++++++ .../pipelines/personalive/tensorrt/export.py | 106 +-- .../personalive/tensorrt/framed_model.py | 301 +++++---- .../pipelines/personalive/tensorrt/runner.py | 185 +++++- .../personalive/tensorrt/unet_explicit.py | 601 ++++++++++++++++++ 8 files changed, 1743 insertions(+), 355 deletions(-) create mode 100644 src/scope/core/pipelines/personalive/tensorrt/engine_model.py create mode 100644 src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py diff --git a/src/scope/core/pipelines/personalive/pipeline.py b/src/scope/core/pipelines/personalive/pipeline.py index 2062ff86..9481e6d9 100644 --- a/src/scope/core/pipelines/personalive/pipeline.py +++ b/src/scope/core/pipelines/personalive/pipeline.py @@ -10,7 +10,8 @@ import logging import time -from collections import deque +from collections import defaultdict, deque +from contextlib import contextmanager from pathlib import Path from typing import TYPE_CHECKING @@ -41,11 +42,61 @@ get_boxes, ) + +# Profiling utilities +_profiling_enabled = False +_profiling_timings: dict[str, list[float]] = defaultdict(list) + + +def enable_profiling(): + """Enable pipeline profiling.""" + global _profiling_enabled + _profiling_enabled = True + _profiling_timings.clear() + + +def disable_profiling(): + """Disable pipeline profiling.""" + global _profiling_enabled + _profiling_enabled = False + + +def reset_profiling(): + """Reset profiling timings.""" + _profiling_timings.clear() + + +def get_profiling_timings() -> dict[str, list[float]]: + """Get profiling timings dictionary.""" + return dict(_profiling_timings) + + +@contextmanager +def _profile_stage(name: str): + """Context manager to profile a pipeline stage.""" + if not _profiling_enabled: + yield + return + torch.cuda.synchronize() if torch.cuda.is_available() else None + start = time.perf_counter() + yield + torch.cuda.synchronize() if torch.cuda.is_available() else None + elapsed = time.perf_counter() - start + _profiling_timings[name].append(elapsed) + # TensorRT support (optional) try: - from .tensorrt import TRT_AVAILABLE, TRTRunner, get_engine_path + from .tensorrt import TRT_AVAILABLE, TRTRunner, get_engine_path, PYCUDA_AVAILABLE + + # Import EngineModel if pycuda is available (best performance) + if PYCUDA_AVAILABLE: + from .tensorrt import EngineModel + else: + EngineModel = None except ImportError: TRT_AVAILABLE = False + PYCUDA_AVAILABLE = False + EngineModel = None if TYPE_CHECKING: from ..schema import BasePipelineConfig @@ -272,6 +323,7 @@ def __init__( # TensorRT support (optional) self.use_tensorrt = False self.trt_runner = None + self._use_engine_model = False # True if using pycuda EngineModel self._model_dir = model_dir # Store for TensorRT path lookup if TRT_AVAILABLE: @@ -304,9 +356,33 @@ def _init_tensorrt(self, model_dir: Path): ) return + device_index = self.device.index if self.device.index is not None else 0 + + # Prefer EngineModel (pycuda-based) for best performance + if EngineModel is not None: + try: + self.trt_runner = EngineModel(engine_path, device_int=device_index) + self._use_engine_model = True + + # Setup output-to-input bindings for recurrent state (zero-copy) + self.trt_runner.bind({ + "motion_hidden_states_out": "motion_hidden_states", + "pose_cond_fea_out": "pose_cond_fea", + "latents": "sample", + }) + + self.use_tensorrt = True + logger.info(f"TensorRT engine loaded with pycuda (EngineModel) from {engine_path}") + return + + except Exception as e: + logger.warning(f"Failed to load EngineModel: {e}") + logger.info("Falling back to polygraphy TRTRunner") + + # Fall back to TRTRunner (polygraphy-based) try: - device_index = self.device.index if self.device.index is not None else 0 self.trt_runner = TRTRunner(engine_path, device=self.device) + self._use_engine_model = False # Setup output-to-input bindings for recurrent state self.trt_runner.bind({ @@ -316,13 +392,14 @@ def _init_tensorrt(self, model_dir: Path): }) self.use_tensorrt = True - logger.info(f"TensorRT engine loaded from {engine_path}") + logger.info(f"TensorRT engine loaded with polygraphy (TRTRunner) from {engine_path}") except Exception as e: logger.warning(f"Failed to load TensorRT engine: {e}") logger.warning("Falling back to PyTorch inference") self.trt_runner = None self.use_tensorrt = False + self._use_engine_model = False def prepare(self, **kwargs) -> Requirements | None: """Return input requirements. @@ -441,8 +518,9 @@ def fuse_reference(self, ref_image: Image.Image): reference_hidden_states = self.reference_control_writer.output() self._reference_hidden_states = reference_hidden_states - # Prefill TensorRT runner with constant inputs - self.trt_runner.clear_prefill() + # Clear previous prefill values (TRTRunner only - EngineModel persists buffers) + if hasattr(self.trt_runner, 'clear_prefill'): + self.trt_runner.clear_prefill() self.trt_runner.prefill(encoder_hidden_states=self.encoder_hidden_states) # Prefill reference hidden states @@ -487,7 +565,7 @@ def fuse_reference(self, ref_image: Image.Image): r = (i + 1) * self.temporal_window_size self.latents_pile.append(noisy_latents_first[:, :, l:r]) - # For TensorRT, also prefill initial latents + # For TensorRT, also prefill initial latents (using correct input name 'sample') if self.use_tensorrt: sample = torch.cat(list(self.latents_pile), dim=2) # Add placeholder for new latents @@ -497,7 +575,7 @@ def fuse_reference(self, ref_image: Image.Image): noise = torch.randn_like(new_latents) new_latents = self.scheduler.add_noise(new_latents, noise, self.timesteps[-1:]) sample = torch.cat([sample, new_latents], dim=2) - self.trt_runner.prefill(latents=sample) + self.trt_runner.prefill(sample=sample) # Use 'sample' (input name), not 'latents' (output name) self.reference_fused = True logger.info("Reference image fused successfully") @@ -584,133 +662,162 @@ def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: temporal_window_size = self.temporal_window_size temporal_adaptive_step = self.temporal_adaptive_step - # Reshape from (B, C, T, H, W) to (T, C, H, W) for processing - if images.dim() == 5: - images = images.squeeze(0).permute(1, 0, 2, 3) - - num_frames = images.shape[0] - if num_frames != temporal_window_size: - images = images[-temporal_window_size:] - - if images.min() < 0: - images = images / 2 + 0.5 - images = images.clamp(0, 1) - - # Get keypoints using pose encoder (still PyTorch) - tgt_cond_tensor = self._fast_resize(images, 256, 256).clamp(0, 1) - + # Input preprocessing + with _profile_stage("input_preprocessing"): + # Reshape from (B, C, T, H, W) to (T, C, H, W) for processing + if images.dim() == 5: + images = images.squeeze(0).permute(1, 0, 2, 3) + + num_frames = images.shape[0] + if num_frames != temporal_window_size: + images = images[-temporal_window_size:] + + if images.min() < 0: + images = images / 2 + 0.5 + images = images.clamp(0, 1) + + # Pose encoding (keypoint extraction) - this is the FAN/MediaPipe step + with _profile_stage("pose_encoding"): + tgt_cond_tensor = self._fast_resize(images, 256, 256).clamp(0, 1) + + if self.first_frame: + mot_bbox_param, kps_ref, kps_frame1, kps_dri = self.pose_encoder.interpolate_kps_online( + self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + ) + self.kps_ref = kps_ref + self.kps_frame1 = kps_frame1 + else: + mot_bbox_param, kps_dri = self.pose_encoder.get_kps( + self.kps_ref, self.kps_frame1, tgt_cond_tensor + ) + + # Keypoint drawing + with _profile_stage("keypoint_drawing"): + keypoints = draw_keypoints(mot_bbox_param, device=device) + boxes = get_boxes(kps_dri) + keypoints = rearrange(keypoints.unsqueeze(2), 'f c b h w -> b c f h w') + keypoints = keypoints.to(device=device, dtype=self.dtype) + + # Face cropping for motion input + with _profile_stage("face_cropping"): + if self.first_frame: + ref_box = get_boxes(mot_bbox_param[:1]) + ref_face = self._crop_face_tensor(self.ref_image_tensor, ref_box[0]) + motion_face = [ref_face] + for i, frame in enumerate(images): + motion_face.append(self._crop_face_tensor(frame, boxes[i])) + pose_cond_tensor = torch.cat(motion_face, dim=0).transpose(0, 1) + pose_cond_tensor = pose_cond_tensor.unsqueeze(0) + else: + motion_face = [] + for i, frame in enumerate(images): + motion_face.append(self._crop_face_tensor(frame, boxes[i])) + + # First frame initialization (motion encoder + pose guider + TRT prefill) if self.first_frame: - mot_bbox_param, kps_ref, kps_frame1, kps_dri = self.pose_encoder.interpolate_kps_online( - self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 - ) - self.kps_ref = kps_ref - self.kps_frame1 = kps_frame1 + with _profile_stage("first_frame_init"): + # For first frame, compute initial states using PyTorch motion encoder + motion_hidden_states = self.motion_encoder(pose_cond_tensor) + ref_motion = motion_hidden_states[:, :1] + dri_motion = motion_hidden_states[:, 1:] + + init_motion_hidden_states = self._interpolate_tensors( + ref_motion, dri_motion[:, :1], num=12 + 1 + )[:, :-1] + + # Compute initial pose features using PyTorch pose guider + pose_fea = self.pose_guider(keypoints) + + # Prefill TensorRT with initial motion and pose states + self.trt_runner.prefill( + motion_hidden_states=init_motion_hidden_states, + pose_cond_fea=pose_fea[:, :, :temporal_window_size * (temporal_adaptive_step - 1)], + ) + + self.motion_bank = ref_motion + self.first_frame = False + + # Prepare TensorRT inputs + pose = keypoints[:, :, -temporal_window_size:] + motion = torch.cat(motion_face[1:], dim=0).transpose(0, 1).unsqueeze(0) else: - mot_bbox_param, kps_dri = self.pose_encoder.get_kps( - self.kps_ref, self.kps_frame1, tgt_cond_tensor + with _profile_stage("motion_preparation"): + motion = torch.cat(motion_face, dim=0).transpose(0, 1).unsqueeze(0) + pose = keypoints + + # Prepare noise + with _profile_stage("noise_generation"): + new_noise = torch.randn( + batch_size, 4, temporal_window_size, + self.height // 8, self.width // 8, + device=device, dtype=self.dtype ) - keypoints = draw_keypoints(mot_bbox_param, device=device) - boxes = get_boxes(kps_dri) - keypoints = rearrange(keypoints.unsqueeze(2), 'f c b h w -> b c f h w') - keypoints = keypoints.to(device=device, dtype=self.dtype) - - # Prepare motion input (face crops) - if self.first_frame: - ref_box = get_boxes(mot_bbox_param[:1]) - ref_face = self._crop_face_tensor(self.ref_image_tensor, ref_box[0]) - motion_face = [ref_face] - for i, frame in enumerate(images): - motion_face.append(self._crop_face_tensor(frame, boxes[i])) - pose_cond_tensor = torch.cat(motion_face, dim=0).transpose(0, 1) - pose_cond_tensor = pose_cond_tensor.unsqueeze(0) - - # For first frame, compute initial states using PyTorch motion encoder - motion_hidden_states = self.motion_encoder(pose_cond_tensor) - ref_motion = motion_hidden_states[:, :1] - dri_motion = motion_hidden_states[:, 1:] - - init_motion_hidden_states = self._interpolate_tensors( - ref_motion, dri_motion[:, :1], num=12 + 1 - )[:, :-1] - - # Compute initial pose features using PyTorch pose guider - pose_fea = self.pose_guider(keypoints) - - # Prefill TensorRT with initial states - self.trt_runner.prefill( - motion_hidden_states_out=init_motion_hidden_states, - pose_cond_fea_out=pose_fea[:, :, :temporal_window_size * (temporal_adaptive_step - 1)], - ) - - self.motion_bank = ref_motion - self.first_frame = False - - # Prepare TensorRT inputs - pose = keypoints[:, :, -temporal_window_size:] - motion = dri_motion.transpose(1, 2).unsqueeze(0) # Reshape for TRT - else: - motion_face = [] - for i, frame in enumerate(images): - motion_face.append(self._crop_face_tensor(frame, boxes[i])) - motion = torch.cat(motion_face, dim=0).transpose(0, 1).unsqueeze(0) - pose = keypoints - - # Prepare new noise for next iteration - new_noise = torch.randn( - batch_size, 4, temporal_window_size, - self.height // 8, self.width // 8, - device=device, dtype=self.dtype - ) - - # Run TensorRT inference - outputs = self.trt_runner( - output_names=["pred_video", "motion_out", "latent_first"], - return_torch=True, - pose=pose, - motion=motion.to(self.dtype), - new_noise=new_noise, - ) + # TensorRT inference (includes motion_encoder, pose_guider, denoising_unet, vae_decode) + with _profile_stage("tensorrt_inference"): + # EngineModel and TRTRunner have slightly different APIs + if self._use_engine_model: + outputs = self.trt_runner( + output_names=["pred_video", "motion_out", "latent_first"], + return_torch=True, + pose=pose, + motion=motion.to(self.dtype), + new_noise=new_noise, + ) + else: + outputs = self.trt_runner( + output_list=["pred_video", "motion_out", "latent_first"], + return_tensor=True, + pose=pose, + motion=motion.to(self.dtype), + new_noise=new_noise, + ) video = outputs["pred_video"] motion_out = outputs.get("motion_out") # Keyframe tracking - if motion_out is not None and self.count > 8: - idx_to_add, self.motion_bank, _ = self._calculate_dis( - self.motion_bank, motion_out.unsqueeze(0), threshold=17.0 - ) - - if len(idx_to_add) > 0 and self.num_khf < 3: - # Update reference hidden states for keyframe - latent_first = outputs.get("latent_first") - if latent_first is not None: - self.reference_control_writer.clear() - self.reference_unet( - latent_first.to(self.reference_unet.dtype), - torch.zeros((batch_size,), dtype=self.dtype, device=device), - encoder_hidden_states=self.encoder_hidden_states, - return_dict=False, - ) - new_ref_hidden = self.reference_control_writer.output() - - # Concatenate new keyframe features - ref_hidden_names = [ - "d00", "d01", "d10", "d11", "d20", "d21", "m", - "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" - ] - for name in ref_hidden_names: - if name in self._reference_hidden_states and name in new_ref_hidden: - self._reference_hidden_states[name] = torch.cat( - [self._reference_hidden_states[name], new_ref_hidden[name]], dim=1 - ) - self.trt_runner.prefill(**{name: self._reference_hidden_states[name]}) - - logger.debug("Added history keyframe (TensorRT)") - self.num_khf += 1 + with _profile_stage("keyframe_tracking"): + if motion_out is not None and self.count > 8: + # motion_out from TensorRT is already 4D [B, 1, ml, mc], matching motion_bank + idx_to_add, self.motion_bank, _ = self._calculate_dis( + self.motion_bank, motion_out, threshold=17.0 + ) + + if len(idx_to_add) > 0 and self.num_khf < 3: + # Update reference hidden states for keyframe + latent_first = outputs.get("latent_first") + if latent_first is not None: + self.reference_control_writer.clear() + self.reference_unet( + latent_first.to(self.reference_unet.dtype), + torch.zeros((batch_size,), dtype=self.dtype, device=device), + encoder_hidden_states=self.encoder_hidden_states, + return_dict=False, + ) + new_ref_hidden = self.reference_control_writer.output() + + # Concatenate new keyframe features + ref_hidden_names = [ + "d00", "d01", "d10", "d11", "d20", "d21", "m", + "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" + ] + for name in ref_hidden_names: + if name in self._reference_hidden_states and name in new_ref_hidden: + self._reference_hidden_states[name] = torch.cat( + [self._reference_hidden_states[name], new_ref_hidden[name]], dim=1 + ) + self.trt_runner.prefill(**{name: self._reference_hidden_states[name]}) + + logger.debug("Added history keyframe (TensorRT)") + self.num_khf += 1 + + # Output to CPU + with _profile_stage("output_transfer"): + video = video.cpu() self.count += 1 - return video.cpu() + return video @torch.no_grad() def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: diff --git a/src/scope/core/pipelines/personalive/tensorrt/__init__.py b/src/scope/core/pipelines/personalive/tensorrt/__init__.py index 4ec3e734..d1bae587 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/__init__.py +++ b/src/scope/core/pipelines/personalive/tensorrt/__init__.py @@ -6,6 +6,9 @@ Installation: pip install daydream-scope[tensorrt] +For best performance, also install pycuda: + pip install pycuda + Usage: # Convert models to TensorRT (run once) convert-personalive-trt --model-dir ./models --height 512 --width 512 @@ -15,11 +18,30 @@ # Check for TensorRT availability before importing TRT_AVAILABLE = False +CUDA_BUFFERS_AVAILABLE = False +PYCUDA_AVAILABLE = False + try: import tensorrt # noqa: F401 from polygraphy.backend.trt import TrtRunner as _TrtRunner # noqa: F401 TRT_AVAILABLE = True + + # Check for zero-copy CUDA buffer support (polygraphy) + try: + from polygraphy.cuda import DeviceView # noqa: F401 + CUDA_BUFFERS_AVAILABLE = True + except ImportError: + pass + + # Check for pycuda (best performance) + try: + import pycuda.driver # noqa: F401 + import pycuda.autoinit # noqa: F401 + PYCUDA_AVAILABLE = True + except ImportError: + pass + except ImportError: pass @@ -27,12 +49,21 @@ from .builder import build_engine, get_engine_path, is_engine_available from .runner import TRTRunner + # Import EngineModel if pycuda is available + if PYCUDA_AVAILABLE: + from .engine_model import EngineModel + else: + EngineModel = None + __all__ = [ "build_engine", "get_engine_path", "is_engine_available", "TRTRunner", + "EngineModel", "TRT_AVAILABLE", + "CUDA_BUFFERS_AVAILABLE", + "PYCUDA_AVAILABLE", ] else: # Provide stub functions when TRT is not available @@ -55,4 +86,5 @@ def __init__(self, *args, **kwargs): "is_engine_available", "TRTRunner", "TRT_AVAILABLE", + "CUDA_BUFFERS_AVAILABLE", ] diff --git a/src/scope/core/pipelines/personalive/tensorrt/convert.py b/src/scope/core/pipelines/personalive/tensorrt/convert.py index 99adf4b4..390e75f5 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/convert.py +++ b/src/scope/core/pipelines/personalive/tensorrt/convert.py @@ -30,6 +30,14 @@ logger = logging.getLogger(__name__) +def log_gpu_memory(msg: str = ""): + """Log current GPU memory usage.""" + if torch.cuda.is_available(): + allocated = torch.cuda.memory_allocated() / 1024**3 + reserved = torch.cuda.memory_reserved() / 1024**3 + logger.info(f"GPU Memory {msg}: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved") + + def parse_args(): """Parse command line arguments.""" parser = argparse.ArgumentParser( @@ -80,12 +88,12 @@ def parse_args(): def load_models(model_dir: Path, device: torch.device, dtype: torch.dtype): - """Load all PersonaLive models. + """Load all PersonaLive models with aggressive memory management. Args: model_dir: Base model directory. device: Target device. - dtype: Data type for models. + dtype: Data type for models (used for TensorRT, ONNX export uses FP32). Returns: Dictionary containing all loaded models and config. @@ -94,10 +102,19 @@ def load_models(model_dir: Path, device: torch.device, dtype: torch.dtype): DDIMScheduler, MotEncoder, PoseGuider, - UNet3DConditionModel, ) + # Use explicit reference UNet for TensorRT export + from .unet_explicit import UNet3DConditionModelExplicit + + def cleanup(): + """Aggressive memory cleanup.""" + gc.collect() + torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.synchronize() logger.info("Loading PersonaLive models...") + cleanup() # Paths personalive_dir = model_dir / "PersonaLive" / "pretrained_weights" @@ -111,54 +128,65 @@ def load_models(model_dir: Path, device: torch.device, dtype: torch.dtype): unet_kwargs = OmegaConf.to_container(model_config.unet_additional_kwargs) sched_kwargs = OmegaConf.to_container(model_config.noise_scheduler_kwargs) + # Use FP16 like official implementation - autocast handles dtype during export + export_dtype = dtype + # Load pose guider logger.info("Loading PoseGuider...") - pose_guider = PoseGuider().to(device=device, dtype=dtype) + pose_guider = PoseGuider().to(device=device, dtype=export_dtype) pose_guider_path = weights_path / "pose_guider.pth" if pose_guider_path.exists(): state_dict = torch.load(pose_guider_path, map_location="cpu") pose_guider.load_state_dict(state_dict) del state_dict + cleanup() - # Load motion encoder + # Load motion encoder (contains BatchNorm - must be FP32 for ONNX export) logger.info("Loading MotEncoder...") - motion_encoder = MotEncoder().to(dtype=dtype, device=device).eval() + motion_encoder = MotEncoder().to(dtype=export_dtype, device=device).eval() motion_encoder_path = weights_path / "motion_encoder.pth" if motion_encoder_path.exists(): state_dict = torch.load(motion_encoder_path, map_location="cpu") motion_encoder.load_state_dict(state_dict) del state_dict motion_encoder.set_attn_processor(AttnProcessor()) + cleanup() - # Load denoising UNet - logger.info("Loading UNet3DConditionModel...") - denoising_unet = UNet3DConditionModel.from_pretrained_2d( + # Load denoising UNet (explicit reference version for TensorRT) + logger.info("Loading UNet3DConditionModelExplicit for TensorRT...") + denoising_unet = UNet3DConditionModelExplicit.from_pretrained_2d( str(pretrained_base_path), "", subfolder="unet", unet_additional_kwargs=unet_kwargs, - ).to(dtype=dtype, device=device) + ).to(dtype=export_dtype, device=device) denoising_unet_path = weights_path / "denoising_unet.pth" if denoising_unet_path.exists(): state_dict = torch.load(denoising_unet_path, map_location="cpu") denoising_unet.load_state_dict(state_dict, strict=False) del state_dict + cleanup() temporal_module_path = weights_path / "temporal_module.pth" if temporal_module_path.exists(): state_dict = torch.load(temporal_module_path, map_location="cpu") denoising_unet.load_state_dict(state_dict, strict=False) del state_dict + cleanup() + # IMPORTANT: Use standard AttnProcessor for ONNX export + # xformers is NOT compatible with ONNX tracing! denoising_unet.set_attn_processor(AttnProcessor()) + cleanup() # Load VAE logger.info("Loading VAE...") - vae = AutoencoderKL.from_pretrained(str(vae_path)).to(device=device, dtype=dtype) + vae = AutoencoderKL.from_pretrained(str(vae_path)).to(device=device, dtype=export_dtype) vae.set_default_attn_processor() + cleanup() - # Setup scheduler + # Setup scheduler (CPU-only, minimal memory) scheduler = DDIMScheduler(**sched_kwargs) scheduler.to(device) @@ -170,6 +198,10 @@ def load_models(model_dir: Path, device: torch.device, dtype: torch.dtype): ).long() scheduler.set_step_length(333) + logger.info("All models loaded. Final memory cleanup...") + cleanup() + log_gpu_memory("after loading all models") + return { "pose_guider": pose_guider, "motion_encoder": motion_encoder, @@ -189,7 +221,7 @@ def export_to_onnx( device: torch.device, dtype: torch.dtype, ): - """Export bundled model to ONNX. + """Export bundled model to ONNX following official PersonaLive approach. Args: models: Dictionary of loaded models. @@ -198,12 +230,20 @@ def export_to_onnx( width: Output width. batch_size: Batch size. device: Target device. - dtype: Data type. + dtype: Data type for export. """ from .export import export_onnx, optimize_onnx from .framed_model import UNetWork + def cleanup(): + """Aggressive memory cleanup.""" + gc.collect() + torch.cuda.empty_cache() + if torch.cuda.is_available(): + torch.cuda.synchronize() + logger.info("Creating bundled UNetWork model...") + cleanup() # Create bundled model unet_work = UNetWork( @@ -215,30 +255,35 @@ def export_to_onnx( timesteps=models["timesteps"], ) - # Generate sample inputs - logger.info("Generating sample inputs...") - sample_inputs = unet_work.get_sample_input(batch_size, height, width, dtype, device) - # Export to ONNX + # Note: auto_cast=False because scope's MotEncoder has BatchNorm that + # doesn't work well with autocast (dtype mismatch issues) + log_gpu_memory("before ONNX export") raw_onnx_path = onnx_path.parent / "unet" / "unet.onnx" export_onnx( model=unet_work, onnx_path=raw_onnx_path, - sample_inputs=sample_inputs, - input_names=UNetWork.get_input_names(), - output_names=UNetWork.get_output_names(), - dynamic_axes=UNetWork.get_dynamic_axes(), - opset_version=17, + opt_image_height=height, + opt_image_width=width, + opt_batch_size=batch_size, + onnx_opset=17, + dtype=dtype, + device=device, + auto_cast=False, # Disabled due to BatchNorm dtype issues ) + # Note: export_onnx handles cleanup of unet_work internally + cleanup() - # Optimize ONNX + # Clear the models dict + for key in list(models.keys()): + del models[key] + cleanup() + log_gpu_memory("after ONNX export cleanup") + + # Optimize ONNX (CPU-only, doesn't need GPU) logger.info("Optimizing ONNX model...") optimize_onnx(raw_onnx_path, onnx_path) - - # Clean up - del unet_work, sample_inputs - gc.collect() - torch.cuda.empty_cache() + cleanup() def main(): @@ -271,9 +316,14 @@ def main(): device = torch.device(args.device) dtype = torch.float32 if args.fp32 else torch.float16 + # ONNX export uses 256x256 to save memory (official approach) + # TRT engine is built at target resolution with dynamic shapes + onnx_height, onnx_width = 256, 256 + logger.info(f"Converting PersonaLive models to TensorRT") logger.info(f" Model directory: {args.model_dir}") - logger.info(f" Resolution: {args.height}x{args.width}") + logger.info(f" Target resolution: {args.height}x{args.width}") + logger.info(f" ONNX export resolution: {onnx_height}x{onnx_width} (to save memory)") logger.info(f" Precision: {'FP32' if args.fp32 else 'FP16'}") logger.info(f" Device: {device}") @@ -283,15 +333,16 @@ def main(): onnx_path = get_onnx_path(args.model_dir) engine_path = get_engine_path(args.model_dir, args.height, args.width) - # Step 1: Export to ONNX (if needed) + # Step 1: Export to ONNX at 256x256 (if needed) + # This uses less memory - TRT handles target resolution via dynamic shapes if not args.skip_onnx or not onnx_path.exists(): - logger.info("Step 1: Exporting to ONNX...") + logger.info("Step 1: Exporting to ONNX at 256x256...") models = load_models(args.model_dir, device, dtype) export_to_onnx( models=models, onnx_path=onnx_path, - height=args.height, - width=args.width, + height=onnx_height, # 256 to save memory + width=onnx_width, # 256 to save memory batch_size=args.batch_size, device=device, dtype=dtype, @@ -303,9 +354,10 @@ def main(): else: logger.info("Step 1: Skipping ONNX export (using existing files)") - # Step 2: Build TensorRT engine + # Step 2: Build TensorRT engine at target resolution logger.info("Step 2: Building TensorRT engine...") - logger.info("This may take 10-30 minutes depending on your GPU.") + logger.info(f" Building for {args.height}x{args.width} resolution") + logger.info(" This may take 10-30 minutes depending on your GPU.") build_engine( onnx_path=onnx_path, diff --git a/src/scope/core/pipelines/personalive/tensorrt/engine_model.py b/src/scope/core/pipelines/personalive/tensorrt/engine_model.py new file mode 100644 index 00000000..6068e0d1 --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/engine_model.py @@ -0,0 +1,394 @@ +"""TensorRT engine runner using pycuda for direct CUDA memory management. + +Based on PersonaLive official implementation: +PersonaLive/src/modeling/engine_model.py + +This provides significantly better performance than polygraphy's TrtRunner +by using direct CUDA memory allocation and device-to-device copies. +""" + +import logging +import os +import traceback +from pathlib import Path + +import numpy as np +import torch + +logger = logging.getLogger(__name__) + +# Check for TensorRT and pycuda availability +PYCUDA_TRT_AVAILABLE = False +try: + import tensorrt as trt + import pycuda.driver as cuda + import pycuda.autoinit # noqa: F401 - Required for pycuda initialization + + PYCUDA_TRT_AVAILABLE = True + TRT_LOGGER = trt.Logger(trt.Logger.WARNING) +except ImportError as e: + logger.debug(f"pycuda/tensorrt not available: {e}") + + +def _get_engine(engine_file_path: str): + """Load TensorRT engine from file.""" + if os.path.exists(engine_file_path): + logger.info(f"Loading engine from file {engine_file_path}...") + with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime: + return runtime.deserialize_cuda_engine(f.read()) + else: + raise FileNotFoundError(f"Engine file not found: {engine_file_path}") + + +def _numpy_to_torch_dtype(np_dtype): + """Convert numpy dtype to torch dtype.""" + mapping = { + np.float32: torch.float, + np.float64: torch.double, + np.float16: torch.half, + np.int32: torch.int32, + np.int64: torch.int64, + np.int16: torch.int16, + np.int8: torch.int8, + np.uint8: torch.uint8, + np.bool_: torch.bool, + } + return mapping.get(np_dtype, torch.float) + + +class EngineModel: + """TensorRT engine runner with pycuda for direct CUDA memory management. + + This provides significantly better performance than polygraphy by: + 1. Pre-allocating CUDA device memory for all inputs/outputs + 2. Using device-to-device copies (memcpy_dtod_async) when possible + 3. Using async TRT execution with CUDA streams + 4. True zero-copy binding for recurrent state + + Based on PersonaLive official implementation. + """ + + def __init__( + self, + engine_file_path: str | Path, + device_int: int = 0, + ): + """Initialize TensorRT engine with pycuda. + + Args: + engine_file_path: Path to the TensorRT engine file. + device_int: CUDA device index. + """ + if not PYCUDA_TRT_AVAILABLE: + raise RuntimeError( + "pycuda and tensorrt are required. " + "Install with: pip install pycuda tensorrt" + ) + + engine_file_path = str(engine_file_path) + if not os.path.exists(engine_file_path): + raise FileNotFoundError(f"Engine file not found: {engine_file_path}") + + self.device_int = device_int + self._torch_device = torch.device(f"cuda:{device_int}") + + # Create CUDA context for this device + self.ctx = cuda.Device(device_int).make_context() + + try: + # Load TensorRT engine + self.engine = _get_engine(engine_file_path) + + # Get input/output tensor names + self.input_names = [] + self.output_names = [] + + for binding in self.engine: + mode = self.engine.get_tensor_mode(binding) + if mode == trt.TensorIOMode.INPUT: + self.input_names.append(binding) + elif mode == trt.TensorIOMode.OUTPUT: + self.output_names.append(binding) + + # Helper to get safe shape (handle dynamic dimensions) + def get_safe_shape(engine, name): + shape = engine.get_tensor_shape(name) + if -1 in shape: + # Use max profile shape for dynamic dimensions + profile = engine.get_tensor_profile_shape(name, 0) + if profile: + logger.debug(f"Dynamic shape for {name}: {shape} -> Max: {profile[2]}") + return profile[2] + return shape + + # Get shapes and dtypes for all tensors + self.input_shapes = {name: get_safe_shape(self.engine, name) for name in self.input_names} + self.input_dtypes = {name: self.engine.get_tensor_dtype(name) for name in self.input_names} + self.input_nbytes = { + name: trt.volume(self.input_shapes[name]) * trt.nptype(self.input_dtypes[name])().itemsize + for name in self.input_names + } + + self.output_shapes = {name: get_safe_shape(self.engine, name) for name in self.output_names} + self.output_dtypes = {name: self.engine.get_tensor_dtype(name) for name in self.output_names} + self.output_nbytes = { + name: trt.volume(self.output_shapes[name]) * trt.nptype(self.output_dtypes[name])().itemsize + for name in self.output_names + } + + # Allocate CUDA device memory for inputs and outputs + self.dinputs = {name: cuda.mem_alloc(self.input_nbytes[name]) for name in self.input_names} + self.doutputs = {name: cuda.mem_alloc(self.output_nbytes[name]) for name in self.output_names} + + # Create execution context + self.context = self.engine.create_execution_context() + + # Create CUDA stream + self.stream = cuda.Stream() + + # Bind tensor addresses to context + for name in self.input_names: + self.context.set_tensor_address(name, int(self.dinputs[name])) + for name in self.output_names: + self.context.set_tensor_address(name, int(self.doutputs[name])) + + # Allocate page-locked host memory for outputs (for CPU transfer) + self.houtputs = { + name: cuda.pagelocked_empty( + trt.volume(self.output_shapes[name]), + dtype=trt.nptype(self.output_dtypes[name]) + ) + for name in self.output_names + } + + logger.info(f"TensorRT engine loaded successfully ({len(self.input_names)} inputs, {len(self.output_names)} outputs)") + + except Exception as e: + self.ctx.pop() + raise RuntimeError(f"Failed to initialize TensorRT engine: {e}") + + self.ctx.pop() + + def __call__( + self, + output_names: list[str] | None = None, + return_torch: bool = True, + **inputs: torch.Tensor | np.ndarray, + ) -> dict[str, torch.Tensor | np.ndarray]: + """Run inference on the TensorRT engine. + + Args: + output_names: List of output names to return. If None or empty, returns all. + return_torch: If True, return PyTorch tensors on CUDA. + **inputs: Named inputs as PyTorch tensors or numpy arrays. + + Returns: + Dictionary mapping output names to tensors/arrays. + """ + if output_names is None or len(output_names) == 0: + output_names = self.output_names + + self.ctx.push() + result = {} + + try: + # Copy inputs to device memory + for name, hinput in inputs.items(): + if name not in self.input_names: + continue + + if isinstance(hinput, torch.Tensor) and hinput.is_cuda and hinput.device.index == self.device_int: + # GPU tensor on same device - device-to-device copy (fast!) + hinput_con = hinput.contiguous() + cuda.memcpy_dtod_async( + self.dinputs[name], + hinput_con.data_ptr(), + self.input_nbytes[name], + self.stream + ) + else: + # CPU tensor or numpy array - host-to-device copy + if isinstance(hinput, torch.Tensor): + hinput = hinput.detach().cpu().numpy() + hinput_con = np.ascontiguousarray(hinput) + cuda.memcpy_htod_async(self.dinputs[name], hinput_con, self.stream) + + # Set input shapes for any inputs not provided (use prefilled) + for name in self.input_names: + if name not in inputs: + self.context.set_input_shape(name, self.input_shapes[name]) + + # Execute TensorRT inference asynchronously + self.context.execute_async_v3(self.stream.handle) + + # Copy outputs + if return_torch: + # Device-to-device copy to new torch tensor + for name in output_names: + t = torch.zeros( + trt.volume(self.output_shapes[name]), + device=self._torch_device, + dtype=_numpy_to_torch_dtype(trt.nptype(self.output_dtypes[name])) + ) + cuda.memcpy_dtod_async( + t.data_ptr(), + self.doutputs[name], + self.output_nbytes[name], + self.stream + ) + t = t.reshape(tuple(self.output_shapes[name])) + result[name] = t + else: + # Device-to-host copy to numpy + for name in output_names: + cuda.memcpy_dtoh_async(self.houtputs[name], self.doutputs[name], self.stream) + result[name] = self.houtputs[name].reshape(self.output_shapes[name]) + + # Synchronize stream + self.stream.synchronize() + + except Exception as e: + logger.error(f"TensorRT execution failed: {e}") + traceback.print_exc() + self.ctx.pop() + raise + + self.ctx.pop() + return result + + def prefill(self, **inputs: torch.Tensor | np.ndarray) -> bool: + """Pre-fill inputs/outputs with constant values. + + This copies data to the pre-allocated CUDA buffers. + + Args: + **inputs: Named inputs as PyTorch tensors or numpy arrays. + + Returns: + True if successful. + """ + self.ctx.push() + + try: + for name, hinput in inputs.items(): + # Check if it's an input or output buffer + if name in self.input_names: + dst_ptr = self.dinputs[name] + elif name in self.output_names: + dst_ptr = self.doutputs[name] + else: + logger.warning(f"Unknown tensor name for prefill: {name}") + continue + + # Calculate actual bytes to copy + if isinstance(hinput, torch.Tensor): + real_nbytes = hinput.numel() * hinput.element_size() + else: + real_nbytes = hinput.nbytes + + # Copy to device + if isinstance(hinput, torch.Tensor) and hinput.is_cuda and hinput.device.index == self.device_int: + # GPU tensor - device-to-device copy + hinput_con = hinput.contiguous() + cuda.memcpy_dtod_async(dst_ptr, hinput_con.data_ptr(), real_nbytes, self.stream) + else: + # CPU tensor/numpy - host-to-device copy + if isinstance(hinput, torch.Tensor): + hinput = hinput.detach().cpu().numpy() + hinput_con = np.ascontiguousarray(hinput) + cuda.memcpy_htod_async(dst_ptr, hinput_con, self.stream) + + self.stream.synchronize() + + except Exception as e: + logger.error(f"Prefill failed: {e}") + traceback.print_exc() + self.ctx.pop() + return False + + self.ctx.pop() + return True + + def bind(self, output_to_input: dict[str, str]) -> bool: + """Bind output buffers directly to input buffers for zero-copy recurrence. + + This sets the input tensor address to point to the output buffer, + eliminating any memory copies for recurrent state. + + Args: + output_to_input: Mapping from output names to input names. + + Returns: + True if successful. + """ + self.ctx.push() + + try: + for output_name, input_name in output_to_input.items(): + if output_name not in self.output_names: + logger.warning(f"Unknown output for bind: {output_name}") + continue + if input_name not in self.input_names: + logger.warning(f"Unknown input for bind: {input_name}") + continue + + # Point input address to output buffer + self.context.set_tensor_address(input_name, int(self.doutputs[output_name])) + + except Exception as e: + logger.error(f"Bind failed: {e}") + traceback.print_exc() + self.ctx.pop() + return False + + self.ctx.pop() + return True + + def clear_prefill(self): + """Clear is not needed for EngineModel - buffers persist.""" + pass + + def get_input_names(self) -> list[str]: + """Get list of input tensor names.""" + return list(self.input_names) + + def get_output_names(self) -> list[str]: + """Get list of output tensor names.""" + return list(self.output_names) + + def get_input_shapes(self) -> dict[str, tuple]: + """Get input tensor shapes.""" + return {name: tuple(shape) for name, shape in self.input_shapes.items()} + + def get_output_shapes(self) -> dict[str, tuple]: + """Get output tensor shapes.""" + return {name: tuple(shape) for name, shape in self.output_shapes.items()} + + def __del__(self): + """Clean up CUDA resources.""" + try: + if hasattr(self, 'ctx'): + self.ctx.push() + # Free CUDA memory + for ptr in getattr(self, 'dinputs', {}).values(): + ptr.free() + for ptr in getattr(self, 'doutputs', {}).values(): + ptr.free() + self.ctx.pop() + except Exception: + pass + + def __repr__(self) -> str: + """String representation.""" + r = "EngineModel(\n Inputs=[\n" + for name in self.input_names: + dtype = trt.nptype(self.input_dtypes[name]).__name__ + r += f" {name}: {dtype}{tuple(self.input_shapes[name])},\n" + r += " ],\n Outputs=[\n" + for name in self.output_names: + dtype = trt.nptype(self.output_dtypes[name]).__name__ + r += f" {name}: {dtype}{tuple(self.output_shapes[name])},\n" + r += " ]\n)" + return r + + diff --git a/src/scope/core/pipelines/personalive/tensorrt/export.py b/src/scope/core/pipelines/personalive/tensorrt/export.py index 3d7bebcc..8e5917d6 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/export.py +++ b/src/scope/core/pipelines/personalive/tensorrt/export.py @@ -1,12 +1,15 @@ """ONNX export utilities for PersonaLive models. -This module provides utilities to export PersonaLive models to ONNX format -for subsequent TensorRT conversion. +Adapted from PersonaLive official implementation: +PersonaLive/src/modeling/onnx_export.py + +Based on: https://github.com/NVIDIA/TensorRT/blob/main/demo/Diffusion/utilities.py """ import gc import logging import os +from contextlib import contextmanager from pathlib import Path import torch @@ -14,50 +17,78 @@ logger = logging.getLogger(__name__) +@contextmanager +def auto_cast_manager(enabled: bool): + """Context manager for autocast during ONNX export.""" + if enabled: + with torch.inference_mode(), torch.autocast("cuda"): + yield + else: + yield + + +@torch.no_grad() def export_onnx( model: torch.nn.Module, onnx_path: Path | str, - sample_inputs: dict[str, torch.Tensor], - input_names: list[str], - output_names: list[str], - dynamic_axes: dict[str, dict[int, str]] | None = None, - opset_version: int = 17, + opt_image_height: int, + opt_image_width: int, + opt_batch_size: int, + onnx_opset: int, + dtype: torch.dtype, + device: torch.device, + auto_cast: bool = True, ) -> None: """Export a PyTorch model to ONNX format. + This follows the official PersonaLive export approach. + Args: - model: PyTorch model to export. + model: PyTorch model with get_sample_input, get_input_names, + get_output_names, get_dynamic_axes methods. onnx_path: Path to save the ONNX model. - sample_inputs: Sample inputs for tracing. - input_names: Names of input tensors. - output_names: Names of output tensors. - dynamic_axes: Dynamic axis specifications. - opset_version: ONNX opset version. + opt_image_height: Optimization image height. + opt_image_width: Optimization image width. + opt_batch_size: Optimization batch size. + onnx_opset: ONNX opset version. + dtype: Data type for sample inputs. + device: Target device. + auto_cast: Whether to use autocast during export. """ onnx_path = Path(onnx_path) onnx_path.parent.mkdir(parents=True, exist_ok=True) logger.info(f"Exporting model to ONNX: {onnx_path}") + logger.info(f" Resolution: {opt_image_height}x{opt_image_width}") + logger.info(f" Batch size: {opt_batch_size}") + logger.info(f" Opset: {onnx_opset}") + logger.info(f" Auto cast: {auto_cast}") + + with auto_cast_manager(auto_cast): + # Generate sample inputs + inputs = model.get_sample_input( + opt_batch_size, opt_image_height, opt_image_width, dtype, device + ) - # Prepare inputs as tuple in correct order - inputs = tuple(sample_inputs[name] for name in input_names) + logger.info(f"Output names: {model.get_output_names()}") - with torch.inference_mode(), torch.autocast("cuda"): - torch.onnx.export( + # Export to ONNX - use torch.onnx.utils.export like official impl + torch.onnx.utils.export( model, - inputs, + inputs, # Pass dict directly, not as tuple str(onnx_path), export_params=True, - opset_version=opset_version, + opset_version=onnx_opset, do_constant_folding=True, - input_names=input_names, - output_names=output_names, - dynamic_axes=dynamic_axes or {}, + input_names=model.get_input_names(), + output_names=model.get_output_names(), + dynamic_axes=model.get_dynamic_axes(), ) logger.info(f"ONNX model exported to {onnx_path}") # Clean up + del model gc.collect() torch.cuda.empty_cache() @@ -84,7 +115,7 @@ def optimize_onnx( onnx_opt_path = Path(onnx_opt_path) onnx_opt_path.parent.mkdir(parents=True, exist_ok=True) - logger.info(f"Optimizing ONNX model: {onnx_path} -> {onnx_opt_path}") + logger.info(f"Saving ONNX model with external data: {onnx_opt_path}") model = onnx.load(str(onnx_path)) name = onnx_opt_path.stem @@ -99,31 +130,4 @@ def optimize_onnx( size_threshold=1024, ) - logger.info(f"Optimized ONNX model saved to {onnx_opt_path}") - - -def verify_onnx(onnx_path: Path | str) -> bool: - """Verify ONNX model is valid. - - Args: - onnx_path: Path to ONNX model. - - Returns: - True if model is valid. - """ - try: - import onnx - except ImportError: - logger.warning("onnx package not available, skipping verification") - return True - - onnx_path = Path(onnx_path) - - try: - model = onnx.load(str(onnx_path)) - onnx.checker.check_model(model) - logger.info(f"ONNX model {onnx_path} is valid") - return True - except Exception as e: - logger.error(f"ONNX model verification failed: {e}") - return False + logger.info("ONNX optimization done.") diff --git a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py index 0bec0373..85533c1d 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py +++ b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py @@ -1,32 +1,33 @@ -"""Framed model for ONNX export. +"""Bundled model for TensorRT export. -This module provides the UNetWork class that bundles multiple models +This module provides the UNetWork class that bundles multiple PersonaLive models into a single module for ONNX export and TensorRT conversion. -The bundled models include: -- PoseGuider: Extracts pose features from keypoint images -- MotEncoder: Encodes facial motion features -- UNet3DConditionModel: Main denoising UNet -- VAE decoder: Decodes latents to images -- DDIMScheduler: Denoising scheduler step +Based on PersonaLive official implementation: +PersonaLive/src/modeling/framed_models.py """ import torch -import torch.nn as nn +from torch import nn from einops import rearrange +try: + from polygraphy.backend.trt import Profile + POLYGRAPHY_AVAILABLE = True +except ImportError: + POLYGRAPHY_AVAILABLE = False + class UNetWork(nn.Module): """Bundled model for TensorRT export. - This class wraps the pose_guider, motion_encoder, denoising_unet, - vae decoder, and scheduler into a single module that can be exported - to ONNX for TensorRT conversion. + This class wraps the pose_guider, motion_encoder, denoising_unet (explicit reference), + vae decoder, and scheduler into a single module for ONNX/TensorRT export. The forward pass handles a complete denoising step including: 1. Encoding new pose features 2. Encoding new motion features - 3. Running the UNet denoiser + 3. Running the UNet denoiser with explicit reference injection 4. Scheduler step 5. VAE decoding for the output frames """ @@ -45,7 +46,7 @@ def __init__( Args: pose_guider: PoseGuider model. motion_encoder: MotEncoder model. - denoising_unet: UNet3DConditionModel for denoising. + denoising_unet: UNet3DConditionModelExplicit for denoising. vae: AutoencoderKL for decoding. scheduler: DDIMScheduler instance. timesteps: Fixed timesteps tensor for denoising. @@ -58,20 +59,13 @@ def __init__( self.scheduler = scheduler self.timesteps = timesteps - def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: - """Decode latents to images using VAE. - - Args: - latents: Latent tensor of shape (B, C, H, W). - - Returns: - Image tensor of shape (B, H, W, C) in [0, 1] range. - """ - latents = latents / 0.18215 - images = self.vae.decode(latents).sample - images = rearrange(images, "b c h w -> b h w c") - images = (images / 2 + 0.5).clamp(0, 1) - return images + def decode_slice(self, vae, x): + """Decode latents to images using VAE.""" + x = x / 0.18215 + x = vae.decode(x).sample + x = rearrange(x, "b c h w -> b h w c") + x = (x / 2 + 0.5).clamp(0, 1) + return x def forward( self, @@ -82,6 +76,7 @@ def forward( pose_cond_fea: torch.Tensor, pose: torch.Tensor, new_noise: torch.Tensor, + # Explicit reference hidden states d00: torch.Tensor, d01: torch.Tensor, d10: torch.Tensor, @@ -99,36 +94,36 @@ def forward( u31: torch.Tensor, u32: torch.Tensor, ): - """Forward pass for the bundled model. + """Forward pass for TensorRT inference. Args: - sample: Noisy latent tensor (B, C, T, H, W). - encoder_hidden_states: CLIP embeddings (B, 1, D). - motion_hidden_states: Accumulated motion features (B, T, D1, D2). - motion: New face crop tensor (B, C, T, H, W). - pose_cond_fea: Accumulated pose features (B, C, T, H, W). - pose: New pose keypoints (B, C, T, H, W). - new_noise: Noise for next iteration (B, C, T, H, W). + sample: Noisy latent (B, 4, 16, H/8, W/8). + encoder_hidden_states: CLIP embeddings (B, 1, 768). + motion_hidden_states: Accumulated motion features (B, 12, 32, 16). + motion: New motion face crops (B, 3, 4, 224, 224). + pose_cond_fea: Accumulated pose features (B, 320, 12, H/8, W/8). + pose: New pose keypoints (B, 3, 4, H, W). + new_noise: Noise for next iteration (B, 4, 4, H/8, W/8). d00-u32: Reference hidden states from reference UNet. Returns: Tuple of: - - pred_video: Decoded video frames (T, H, W, C). - - latents: Updated latent tensor for next iteration. + - pred_video: Decoded video frames (4, H, W, 3). + - latents: Updated latents for next iteration. - pose_cond_fea_out: Updated pose features. - motion_hidden_states_out: Updated motion features. - - motion_out: Motion features for keyframe tracking. - - latent_first: First latent for potential keyframe update. + - motion_out: First motion frame for keyframe detection. + - latent_first: First frame latent for potential keyframe update. """ - # Encode new pose features + # Encode new pose features and concatenate new_pose_cond_fea = self.pose_guider(pose) pose_cond_fea = torch.cat([pose_cond_fea, new_pose_cond_fea], dim=2) - # Encode new motion features + # Encode new motion features and concatenate new_motion_hidden_states = self.motion_encoder(motion) motion_hidden_states = torch.cat([motion_hidden_states, new_motion_hidden_states], dim=1) - # Prepare encoder hidden states for UNet + # Prepare encoder hidden states for UNet [clip_embeds, motion_features] encoder_hidden_states_combined = [encoder_hidden_states, motion_hidden_states] # Run UNet with explicit reference hidden states @@ -156,7 +151,7 @@ def forward( latents_model_input = rearrange(latents_model_input, '(b f) c h w -> b c f h w', f=16) # Decode first 4 frames (temporal_window_size) - pred_video = self._decode_latents(pred_original_sample[:4]) + pred_video = self.decode_slice(self.vae, pred_original_sample[:4]) # Prepare outputs for next iteration # Shift latents: drop first temporal_window_size, add new_noise @@ -166,126 +161,88 @@ def forward( pose_cond_fea_out = pose_cond_fea[:, :, 4:, :, :] motion_hidden_states_out = motion_hidden_states[:, 4:, :, :] - # Motion for keyframe tracking (first frame's motion) + # First motion frame for keyframe tracking motion_out = motion_hidden_states[:, :1, :, :] - # First predicted latent for potential keyframe reference update + # First frame latent for potential keyframe update latent_first = pred_original_sample[:1] - return ( - pred_video, - latents, - pose_cond_fea_out, - motion_hidden_states_out, - motion_out, - latent_first, - ) + return pred_video, latents, pose_cond_fea_out, motion_hidden_states_out, motion_out, latent_first - def get_sample_input( - self, - batch_size: int, - height: int, - width: int, - dtype: torch.dtype, - device: torch.device, - ) -> dict[str, torch.Tensor]: + def get_sample_input(self, batchsize: int, height: int, width: int, dtype, device): """Generate sample inputs for ONNX export. Args: - batch_size: Batch size (usually 1). - height: Output height. - width: Output width. - dtype: Data type (usually float16). + batchsize: Batch size (typically 1). + height: Output image height. + width: Output image width. + dtype: Tensor dtype. device: Target device. Returns: - Dictionary of sample input tensors. + Dictionary of sample tensors for ONNX export. """ - # Constants - tw = 4 # temporal_window_size - ts = 4 # temporal_adaptive_step - tb = tw * ts # temporal batch size (16) - ml, mc = 32, 16 # motion latent size, channels - mh, mw = 224, 224 # motion input size - + tw, ts, tb = 4, 4, 16 # temporal_window_size, temporal_adaptive_steps, temporal_batch + ml, mc, mh, mw = 32, 16, 224, 224 # motion latent dims + b, h, w = batchsize, height, width lh, lw = height // 8, width // 8 # latent height/width - - # UNet channels - cd0, cd1, cd2, cm = 320, 640, 1280, 1280 - cu1, cu2, cu3 = 1280, 640, 320 - - emb = 768 # CLIP embedding dim - lc, ic = 4, 3 # latent channels, image channels - - shapes = { - "sample": (batch_size, lc, tb, lh, lw), - "encoder_hidden_states": (batch_size, 1, emb), - "motion_hidden_states": (batch_size, tw * (ts - 1), ml, mc), - "motion": (batch_size, ic, tw, mh, mw), - "pose_cond_fea": (batch_size, cd0, tw * (ts - 1), lh, lw), - "pose": (batch_size, ic, tw, height, width), - "new_noise": (batch_size, lc, tw, lh, lw), - "d00": (batch_size, lh * lw, cd0), - "d01": (batch_size, lh * lw, cd0), - "d10": (batch_size, lh * lw // 4, cd1), - "d11": (batch_size, lh * lw // 4, cd1), - "d20": (batch_size, lh * lw // 16, cd2), - "d21": (batch_size, lh * lw // 16, cd2), - "m": (batch_size, lh * lw // 64, cm), - "u10": (batch_size, lh * lw // 16, cu1), - "u11": (batch_size, lh * lw // 16, cu1), - "u12": (batch_size, lh * lw // 16, cu1), - "u20": (batch_size, lh * lw // 4, cu2), - "u21": (batch_size, lh * lw // 4, cu2), - "u22": (batch_size, lh * lw // 4, cu2), - "u30": (batch_size, lh * lw, cu3), - "u31": (batch_size, lh * lw, cu3), - "u32": (batch_size, lh * lw, cu3), + cd0, cd1, cd2, cm, cu1, cu2, cu3 = 320, 640, 1280, 1280, 1280, 640, 320 # unet channels + emb = 768 # CLIP embedding dims + lc, ic = 4, 3 # latent/image channels + + profile = { + "sample": [b, lc, tb, lh, lw], + "encoder_hidden_states": [b, 1, emb], + "motion_hidden_states": [b, tw * (ts - 1), ml, mc], + "motion": [b, ic, tw, mh, mw], + "pose_cond_fea": [b, cd0, tw * (ts - 1), lh, lw], + "pose": [b, ic, tw, h, w], + "new_noise": [b, lc, tw, lh, lw], + "d00": [b, lh * lw, cd0], + "d01": [b, lh * lw, cd0], + "d10": [b, lh * lw // 4, cd1], + "d11": [b, lh * lw // 4, cd1], + "d20": [b, lh * lw // 16, cd2], + "d21": [b, lh * lw // 16, cd2], + "m": [b, lh * lw // 64, cm], + "u10": [b, lh * lw // 16, cu1], + "u11": [b, lh * lw // 16, cu1], + "u12": [b, lh * lw // 16, cu1], + "u20": [b, lh * lw // 4, cu2], + "u21": [b, lh * lw // 4, cu2], + "u22": [b, lh * lw // 4, cu2], + "u30": [b, lh * lw, cu3], + "u31": [b, lh * lw, cu3], + "u32": [b, lh * lw, cu3], } - - return {name: torch.randn(shape, dtype=dtype, device=device) for name, shape in shapes.items()} + return {k: torch.randn(profile[k], dtype=dtype, device=device) for k in profile} @staticmethod - def get_input_names() -> list[str]: - """Get list of input tensor names for ONNX export.""" + def get_input_names(): + """Get ordered input names for ONNX export.""" return [ - "sample", - "encoder_hidden_states", - "motion_hidden_states", - "motion", - "pose_cond_fea", - "pose", - "new_noise", + "sample", "encoder_hidden_states", "motion_hidden_states", + "motion", "pose_cond_fea", "pose", "new_noise", "d00", "d01", "d10", "d11", "d20", "d21", "m", - "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32", + "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" ] @staticmethod - def get_output_names() -> list[str]: - """Get list of output tensor names for ONNX export.""" + def get_output_names(): + """Get ordered output names for ONNX export.""" return [ - "pred_video", - "latents", - "pose_cond_fea_out", - "motion_hidden_states_out", - "motion_out", - "latent_first", + "pred_video", "latents", "pose_cond_fea_out", + "motion_hidden_states_out", "motion_out", "latent_first" ] @staticmethod - def get_dynamic_axes() -> dict[str, dict[int, str]]: - """Get dynamic axis specifications for ONNX export. - - Returns: - Dictionary mapping tensor names to their dynamic axes. - """ + def get_dynamic_axes(): + """Get dynamic axes for ONNX export.""" return { - # Resolution-dependent axes "sample": {3: "h_64", 4: "w_64"}, "pose_cond_fea": {3: "h_64", 4: "w_64"}, "pose": {3: "h_512", 4: "w_512"}, "new_noise": {3: "h_64", 4: "w_64"}, - # Dynamic reference hidden states (for keyframe accumulation) "d00": {1: "len_4096"}, "d01": {1: "len_4096"}, "u30": {1: "len_4096"}, @@ -303,3 +260,79 @@ def get_dynamic_axes() -> dict[str, dict[int, str]]: "u12": {1: "len_256"}, "m": {1: "len_64"}, } + + def get_dynamic_profile(self, batchsize: int, height: int, width: int): + """Get TensorRT optimization profile with dynamic shapes. + + Args: + batchsize: Batch size. + height: Target image height. + width: Target image width. + + Returns: + Polygraphy Profile object with min/opt/max shapes. + """ + if not POLYGRAPHY_AVAILABLE: + raise ImportError("polygraphy is required for TensorRT profiles") + + tw, ts, tb = 4, 4, 16 + ml, mc, mh, mw = 32, 16, 224, 224 + b, h, w = batchsize, height, width + lh, lw = height // 8, width // 8 + cd0, cd1, cd2, cm, cu1, cu2, cu3 = 320, 640, 1280, 1280, 1280, 640, 320 + emb = 768 + lc, ic = 4, 3 + + # Fixed inputs (don't change with keyframes) + fixed_inputs_map = { + "sample": (b, lc, tb, lh, lw), + "encoder_hidden_states": (b, 1, emb), + "motion_hidden_states": (b, tw * (ts - 1), ml, mc), + "motion": (b, ic, tw, mh, mw), + "pose_cond_fea": (b, cd0, tw * (ts - 1), lh, lw), + "pose": (b, ic, tw, h, w), + "new_noise": (b, lc, tw, lh, lw), + } + + # Dynamic inputs (grow with keyframe accumulation: 1x, 2x, 4x) + dynamic_inputs_map = { + "d00": (b, lh * lw, cd0), + "d01": (b, lh * lw, cd0), + "d10": (b, lh * lw // 4, cd1), + "d11": (b, lh * lw // 4, cd1), + "d20": (b, lh * lw // 16, cd2), + "d21": (b, lh * lw // 16, cd2), + "m": (b, lh * lw // 64, cm), + "u10": (b, lh * lw // 16, cu1), + "u11": (b, lh * lw // 16, cu1), + "u12": (b, lh * lw // 16, cu1), + "u20": (b, lh * lw // 4, cu2), + "u21": (b, lh * lw // 4, cu2), + "u22": (b, lh * lw // 4, cu2), + "u30": (b, lh * lw, cu3), + "u31": (b, lh * lw, cu3), + "u32": (b, lh * lw, cu3), + } + + profile = Profile() + + # Fixed inputs have same min/opt/max + for name, shape in fixed_inputs_map.items(): + shape_tuple = tuple(shape) + profile.add(name, min=shape_tuple, opt=shape_tuple, max=shape_tuple) + + # Dynamic inputs can grow 1x-4x in the sequence dimension + for name, base_shape in dynamic_inputs_map.items(): + dim0, dim1_base, dim2 = base_shape + + val_1x = dim1_base * 1 # 1 keyframe + val_2x = dim1_base * 2 # 2 keyframes + val_4x = dim1_base * 4 # 4 keyframes (max) + + min_shape = (dim0, val_1x, dim2) + opt_shape = (dim0, val_2x, dim2) + max_shape = (dim0, val_4x, dim2) + + profile.add(name, min=min_shape, opt=opt_shape, max=max_shape) + + return profile diff --git a/src/scope/core/pipelines/personalive/tensorrt/runner.py b/src/scope/core/pipelines/personalive/tensorrt/runner.py index 70f9df8c..974523fb 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/runner.py +++ b/src/scope/core/pipelines/personalive/tensorrt/runner.py @@ -2,6 +2,8 @@ This module provides a TensorRT engine runner that wraps the engine and handles inference with PyTorch tensor I/O. + +Uses CUDA device buffers (DeviceView) to avoid CPU-GPU memory transfers. """ import logging @@ -15,12 +17,20 @@ # Check for TensorRT availability TRT_AVAILABLE = False +CUDA_BUFFERS_AVAILABLE = False try: import tensorrt as trt from polygraphy.backend.common import BytesFromPath from polygraphy.backend.trt import EngineFromBytes, TrtRunner TRT_AVAILABLE = True + + # Check for CUDA buffer support (polygraphy >= 0.47) + try: + from polygraphy.cuda import DeviceView + CUDA_BUFFERS_AVAILABLE = True + except ImportError: + logger.warning("polygraphy.cuda.DeviceView not available - using numpy path (slower)") except ImportError: pass @@ -74,6 +84,11 @@ def __init__( logger.info(f"Loading TensorRT engine from {engine_path}") self._load_engine() + if CUDA_BUFFERS_AVAILABLE: + logger.info("TensorRT runner using CUDA device buffers (zero-copy)") + else: + logger.warning("TensorRT runner using numpy path (slower - install polygraphy>=0.47 for zero-copy)") + def _load_engine(self): """Load the TensorRT engine.""" # Load engine bytes @@ -92,15 +107,26 @@ def prefill(self, **inputs: torch.Tensor | np.ndarray): """Pre-fill constant inputs that don't change per frame. These inputs will be automatically included in every inference call. + When CUDA buffers are available, keeps tensors on GPU for zero-copy inference. Args: **inputs: Named inputs as PyTorch tensors or numpy arrays. """ for name, tensor in inputs.items(): - if isinstance(tensor, torch.Tensor): - # Convert to numpy, ensuring contiguous memory - tensor = tensor.detach().cpu().numpy() - self._prefilled[name] = np.ascontiguousarray(tensor) + if CUDA_BUFFERS_AVAILABLE: + # Keep as torch tensor on GPU for zero-copy inference + if isinstance(tensor, np.ndarray): + tensor = torch.from_numpy(tensor).to(self.device).contiguous() + elif isinstance(tensor, torch.Tensor): + if not tensor.is_cuda: + tensor = tensor.to(self.device) + tensor = tensor.contiguous() + self._prefilled[name] = tensor + else: + # Fallback to numpy for compatibility + if isinstance(tensor, torch.Tensor): + tensor = tensor.detach().cpu().numpy() + self._prefilled[name] = np.ascontiguousarray(tensor) def bind(self, output_to_input: dict[str, str]): """Bind output tensors to input tensors for memory reuse. @@ -117,6 +143,46 @@ def clear_prefill(self): """Clear all pre-filled inputs.""" self._prefilled.clear() + def _tensor_to_device_view(self, tensor: torch.Tensor): + """Convert PyTorch tensor to polygraphy DeviceView (zero-copy). + + Args: + tensor: PyTorch CUDA tensor (must be contiguous). + + Returns: + DeviceView pointing to the tensor's GPU memory. + """ + if not CUDA_BUFFERS_AVAILABLE: + raise RuntimeError("DeviceView not available") + + # Ensure tensor is contiguous and on CUDA + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + + if not tensor.is_cuda: + tensor = tensor.cuda() + + # Map PyTorch dtype to numpy dtype properly + dtype_map = { + torch.float16: np.float16, + torch.float32: np.float32, + torch.float64: np.float64, + torch.int32: np.int32, + torch.int64: np.int64, + torch.int16: np.int16, + torch.int8: np.int8, + torch.uint8: np.uint8, + torch.bool: np.bool_, + } + np_dtype = dtype_map.get(tensor.dtype, np.float32) + + # Create DeviceView from tensor's data pointer + return DeviceView( + ptr=tensor.data_ptr(), + shape=tuple(tensor.shape), + dtype=np_dtype, + ) + def __call__( self, output_names: list[str] | None = None, @@ -125,6 +191,8 @@ def __call__( ) -> dict[str, torch.Tensor | np.ndarray]: """Run inference on the TensorRT engine. + Uses CUDA device buffers when available to avoid CPU-GPU memory transfers. + Args: output_names: List of output names to return. If None, returns all outputs. return_torch: If True, return PyTorch tensors. Otherwise, return numpy arrays. @@ -136,27 +204,124 @@ def __call__( if self._runner is None: raise RuntimeError("TensorRT runner not initialized") - # Prepare feed dict with prefilled values - feed_dict = dict(self._prefilled) + # Use CUDA device buffers if available (zero-copy path) + if CUDA_BUFFERS_AVAILABLE and return_torch: + return self._call_cuda_buffers(output_names, **inputs) + + # Fallback to numpy path (has CPU-GPU copy overhead) + return self._call_numpy(output_names, return_torch, **inputs) + + def _call_cuda_buffers( + self, + output_names: list[str] | None = None, + **inputs: torch.Tensor | np.ndarray, + ) -> dict[str, torch.Tensor]: + """Run inference using CUDA device buffers (zero-copy). + + Args: + output_names: List of output names to return. + **inputs: Named inputs as PyTorch tensors. + + Returns: + Dictionary mapping output names to PyTorch tensors. + """ + # Prepare feed dict with DeviceView wrappers + feed_dict = {} + + # Add prefilled values (convert from numpy to DeviceView if needed) + for name, arr in self._prefilled.items(): + if isinstance(arr, np.ndarray): + # Convert numpy prefill to torch tensor (one-time cost) + tensor = torch.from_numpy(arr).to(self.device).contiguous() + self._prefilled[name] = tensor + feed_dict[name] = self._tensor_to_device_view(tensor) + elif isinstance(arr, torch.Tensor): + feed_dict[name] = self._tensor_to_device_view(arr) + else: + # Already a DeviceView or compatible + feed_dict[name] = arr # Add current inputs for name, tensor in inputs.items(): + if isinstance(tensor, np.ndarray): + tensor = torch.from_numpy(tensor).to(self.device) if isinstance(tensor, torch.Tensor): - tensor = tensor.detach().cpu().numpy() - feed_dict[name] = np.ascontiguousarray(tensor) + if not tensor.is_cuda: + tensor = tensor.to(self.device) + feed_dict[name] = self._tensor_to_device_view(tensor.contiguous()) # Run inference outputs = self._runner.infer(feed_dict=feed_dict) + # Convert outputs to PyTorch tensors (they come back as DeviceView or numpy) + result = {} + for name, output in outputs.items(): + if hasattr(output, 'numpy'): + # DeviceView - create tensor from same memory + arr = output.numpy() + result[name] = torch.from_numpy(arr).to(self.device) + elif isinstance(output, np.ndarray): + result[name] = torch.from_numpy(output).to(self.device) + elif isinstance(output, torch.Tensor): + result[name] = output + else: + result[name] = output + + # Update prefilled values with bound outputs + for output_name, input_name in self._output_bindings.items(): + if output_name in result: + self._prefilled[input_name] = result[output_name] + # Filter outputs if requested if output_names is not None: - outputs = {k: v for k, v in outputs.items() if k in output_names} + result = {k: v for k, v in result.items() if k in output_names} - # Update prefilled values with bound outputs + return result + + def _call_numpy( + self, + output_names: list[str] | None = None, + return_torch: bool = True, + **inputs: torch.Tensor | np.ndarray, + ) -> dict[str, torch.Tensor | np.ndarray]: + """Run inference using numpy arrays (has CPU-GPU copy overhead). + + Args: + output_names: List of output names to return. + return_torch: If True, return PyTorch tensors. + **inputs: Named inputs as PyTorch tensors or numpy arrays. + + Returns: + Dictionary mapping output names to tensors/arrays. + """ + # Prepare feed dict with prefilled values + feed_dict = {} + + # Add prefilled numpy values + for name, arr in self._prefilled.items(): + if isinstance(arr, torch.Tensor): + arr = arr.detach().cpu().numpy() + feed_dict[name] = np.ascontiguousarray(arr) + + # Add current inputs + for name, tensor in inputs.items(): + if isinstance(tensor, torch.Tensor): + tensor = tensor.detach().cpu().numpy() + feed_dict[name] = np.ascontiguousarray(tensor) + + # Run inference + outputs = self._runner.infer(feed_dict=feed_dict) + + # Update prefilled values with bound outputs BEFORE filtering + # This ensures recurrent states (latents->sample) are properly updated for output_name, input_name in self._output_bindings.items(): if output_name in outputs: self._prefilled[input_name] = outputs[output_name] + # Filter outputs if requested (after bindings are applied) + if output_names is not None: + outputs = {k: v for k, v in outputs.items() if k in output_names} + # Convert to PyTorch if requested if return_torch: outputs = { diff --git a/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py b/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py new file mode 100644 index 00000000..d3d0dd15 --- /dev/null +++ b/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py @@ -0,0 +1,601 @@ +"""UNet3DConditionModel with explicit reference inputs for TensorRT export. + +This module provides a modified UNet3D that accepts reference hidden states +as explicit forward parameters instead of using the attention processor +writer/reader pattern. This is required for ONNX/TensorRT export. + +Ported from PersonaLive official implementation: +PersonaLive/src/models/unet_3d_explicit_reference.py +""" + +from collections import OrderedDict +from dataclasses import dataclass +from os import PathLike +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import torch +import torch.nn as nn +import torch.utils.checkpoint +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.attention_processor import AttentionProcessor +from diffusers.models.embeddings import TimestepEmbedding, Timesteps +from diffusers.models.modeling_utils import ModelMixin +from diffusers.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, BaseOutput, logging +from safetensors.torch import load_file + +from ..modules.resnet import InflatedConv3d, InflatedGroupNorm +from ..modules.unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block + +logger = logging.get_logger(__name__) + + +@dataclass +class UNet3DConditionOutput(BaseOutput): + sample: torch.FloatTensor + + +class UNet3DConditionModelExplicit(ModelMixin, ConfigMixin): + """UNet3D with explicit reference hidden state inputs for TensorRT export. + + This model accepts reference hidden states (d00, d01, ..., u32) as explicit + forward parameters instead of using attention processor hooks. + """ + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + sample_size: Optional[int] = None, + in_channels: int = 4, + out_channels: int = 4, + center_input_sample: bool = False, + flip_sin_to_cos: bool = True, + freq_shift: int = 0, + down_block_types: Tuple[str] = ( + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "DownBlock3D", + ), + mid_block_type: str = "UNetMidBlock3DCrossAttn", + up_block_types: Tuple[str] = ( + "UpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + ), + only_cross_attention: Union[bool, Tuple[bool]] = False, + block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + layers_per_block: int = 2, + downsample_padding: int = 1, + mid_block_scale_factor: float = 1, + act_fn: str = "silu", + norm_num_groups: int = 32, + norm_eps: float = 1e-5, + cross_attention_dim: int = 1280, + attention_head_dim: Union[int, Tuple[int]] = 8, + dual_cross_attention: bool = False, + use_linear_projection: bool = False, + class_embed_type: Optional[str] = None, + num_class_embeds: Optional[int] = None, + upcast_attention: bool = False, + resnet_time_scale_shift: str = "default", + use_inflated_groupnorm=False, + # Additional + use_motion_module=False, + use_temporal_module=False, + motion_module_resolutions=(1, 2, 4, 8), + motion_module_mid_block=False, + motion_module_decoder_only=False, + motion_module_type=None, + temporal_module_type=None, + motion_module_kwargs={}, + temporal_module_kwargs={}, + unet_use_cross_frame_attention=None, + unet_use_temporal_attention=None, + ): + super().__init__() + + self.sample_size = sample_size + time_embed_dim = block_out_channels[0] * 4 + + # input + self.conv_in = InflatedConv3d( + in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1) + ) + + # time + self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift) + timestep_input_dim = block_out_channels[0] + + self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + + # class embedding + if class_embed_type is None and num_class_embeds is not None: + self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim) + elif class_embed_type == "timestep": + self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim) + elif class_embed_type == "identity": + self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim) + else: + self.class_embedding = None + + self.down_blocks = nn.ModuleList([]) + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + if isinstance(only_cross_attention, bool): + only_cross_attention = [only_cross_attention] * len(down_block_types) + + if isinstance(attention_head_dim, int): + attention_head_dim = (attention_head_dim,) * len(down_block_types) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + res = 2**i + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + + down_block = get_down_block( + down_block_type, + num_layers=layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + temb_channels=time_embed_dim, + add_downsample=not is_final_block, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[i], + downsample_padding=downsample_padding, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module + and (res in motion_module_resolutions) + and (not motion_module_decoder_only), + use_temporal_module=use_temporal_module + and (res in motion_module_resolutions) + and (not motion_module_decoder_only), + motion_module_type=motion_module_type, + temporal_module_type=temporal_module_type, + motion_module_kwargs=motion_module_kwargs, + temporal_module_kwargs=temporal_module_kwargs + ) + self.down_blocks.append(down_block) + + # mid + if mid_block_type == "UNetMidBlock3DCrossAttn": + self.mid_block = UNetMidBlock3DCrossAttn( + in_channels=block_out_channels[-1], + temb_channels=time_embed_dim, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + output_scale_factor=mid_block_scale_factor, + resnet_time_scale_shift=resnet_time_scale_shift, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=attention_head_dim[-1], + resnet_groups=norm_num_groups, + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + upcast_attention=upcast_attention, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module and motion_module_mid_block, + use_temporal_module=use_temporal_module and motion_module_mid_block, + motion_module_type=motion_module_type, + temporal_module_type=temporal_module_type, + motion_module_kwargs=motion_module_kwargs, + temporal_module_kwargs=temporal_module_kwargs, + ) + else: + raise ValueError(f"unknown mid_block_type : {mid_block_type}") + + # count how many layers upsample the videos + self.num_upsamplers = 0 + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + reversed_attention_head_dim = list(reversed(attention_head_dim)) + only_cross_attention = list(reversed(only_cross_attention)) + output_channel = reversed_block_out_channels[0] + for i, up_block_type in enumerate(up_block_types): + res = 2 ** (3 - i) + is_final_block = i == len(block_out_channels) - 1 + + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + input_channel = reversed_block_out_channels[ + min(i + 1, len(block_out_channels) - 1) + ] + + # add upsample block for all BUT final layer + if not is_final_block: + add_upsample = True + self.num_upsamplers += 1 + else: + add_upsample = False + + up_block = get_up_block( + up_block_type, + num_layers=layers_per_block + 1, + in_channels=input_channel, + out_channels=output_channel, + prev_output_channel=prev_output_channel, + temb_channels=time_embed_dim, + add_upsample=add_upsample, + resnet_eps=norm_eps, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + cross_attention_dim=cross_attention_dim, + attn_num_head_channels=reversed_attention_head_dim[i], + dual_cross_attention=dual_cross_attention, + use_linear_projection=use_linear_projection, + only_cross_attention=only_cross_attention[i], + upcast_attention=upcast_attention, + resnet_time_scale_shift=resnet_time_scale_shift, + unet_use_cross_frame_attention=unet_use_cross_frame_attention, + unet_use_temporal_attention=unet_use_temporal_attention, + use_inflated_groupnorm=use_inflated_groupnorm, + use_motion_module=use_motion_module + and (res in motion_module_resolutions), + use_temporal_module=use_temporal_module + and (res in motion_module_resolutions), + motion_module_type=motion_module_type, + temporal_module_type=temporal_module_type, + motion_module_kwargs=motion_module_kwargs, + temporal_module_kwargs=temporal_module_kwargs, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if use_inflated_groupnorm: + self.conv_norm_out = InflatedGroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=norm_eps, + ) + else: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], + num_groups=norm_num_groups, + eps=norm_eps, + ) + self.conv_act = nn.SiLU() + self.conv_out = InflatedConv3d( + block_out_channels[0], out_channels, kernel_size=3, padding=1 + ) + + @property + def attn_processors(self) -> Dict[str, AttentionProcessor]: + """Returns attention processors.""" + processors = {} + + def fn_recursive_add_processors( + name: str, + module: torch.nn.Module, + processors: Dict[str, AttentionProcessor], + ): + if hasattr(module, "set_processor"): + processors[f"{name}.processor"] = module.processor + + for sub_name, child in module.named_children(): + if "temporal_transformer" not in sub_name: + fn_recursive_add_processors(f"{name}.{sub_name}", child, processors) + + return processors + + for name, module in self.named_children(): + if "temporal_transformer" not in name: + fn_recursive_add_processors(name, module, processors) + + return processors + + def set_attn_processor( + self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] + ): + """Sets the attention processor.""" + count = len(self.attn_processors.keys()) + + if isinstance(processor, dict) and len(processor) != count: + raise ValueError( + f"A dict of processors was passed, but the number of processors {len(processor)} does not match the" + f" number of attention layers: {count}. Please make sure to pass {count} processor classes." + ) + + def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): + if hasattr(module, "set_processor"): + if not isinstance(processor, dict): + module.set_processor(processor) + else: + module.set_processor(processor.pop(f"{name}.processor")) + + for sub_name, child in module.named_children(): + fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor) + + for name, module in self.named_children(): + fn_recursive_attn_processor(name, module, processor) + + def forward( + self, + sample: torch.FloatTensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + pose_cond_fea: torch.Tensor, + # Explicit reference hidden states for TensorRT + d00: torch.Tensor, + d01: torch.Tensor, + d10: torch.Tensor, + d11: torch.Tensor, + d20: torch.Tensor, + d21: torch.Tensor, + m: torch.Tensor, + u10: torch.Tensor, + u11: torch.Tensor, + u12: torch.Tensor, + u20: torch.Tensor, + u21: torch.Tensor, + u22: torch.Tensor, + u30: torch.Tensor, + u31: torch.Tensor, + u32: torch.Tensor, + class_labels: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, + mid_block_additional_residual: Optional[torch.Tensor] = None, + return_dict: bool = True, + skip_mm: bool = False, + ) -> Union[UNet3DConditionOutput, Tuple]: + """Forward pass with explicit reference hidden states. + + Args: + sample: Noisy latent input (batch, channel, frames, height, width). + timestep: Timesteps tensor. + encoder_hidden_states: [clip_embeds, motion_hidden_states] list. + pose_cond_fea: Pose conditioning features. + d00-d21: Down block reference hidden states. + m: Mid block reference hidden states. + u10-u32: Up block reference hidden states. + + Returns: + UNet output sample tensor. + """ + default_overall_up_factor = 2**self.num_upsamplers + + forward_upsample_size = False + upsample_size = None + + if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]): + forward_upsample_size = True + + if attention_mask is not None: + attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0 + attention_mask = attention_mask.unsqueeze(1) + + if self.config.center_input_sample: + sample = 2 * sample - 1.0 + + # time embedding + timesteps = timestep + if not torch.is_tensor(timesteps): + is_mps = sample.device.type == "mps" + if isinstance(timestep, float): + dtype = torch.float32 if is_mps else torch.float64 + else: + dtype = torch.int32 if is_mps else torch.int64 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif len(timesteps.shape) == 0: + timesteps = timesteps[None].to(sample.device) + + t_emb = self.time_proj(timesteps) + t_emb = t_emb.to(dtype=self.dtype) + emb = self.time_embedding(t_emb) + + if self.class_embedding is not None: + if class_labels is None: + raise ValueError("class_labels should be provided when num_class_embeds > 0") + if self.config.class_embed_type == "timestep": + class_labels = self.time_proj(class_labels) + class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) + emb = emb + class_emb + + # pre-process + sample = self.conv_in(sample) + if pose_cond_fea is not None: + sample = sample + pose_cond_fea + + # down blocks with explicit reference injection + down_block_res_samples = (sample,) + for i, downsample_block in enumerate(self.down_blocks): + # Set reference hidden states for this block + down_reference = None + if i == 0: + down_reference = [d00, d01] + elif i == 1: + down_reference = [d10, d11] + elif i == 2: + down_reference = [d20, d21] + + if ( + hasattr(downsample_block, "has_cross_attention") + and downsample_block.has_cross_attention + ): + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + skip_mm=skip_mm, + down_reference=down_reference, + ) + else: + sample, res_samples = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + skip_mm=skip_mm, + ) + down_block_res_samples += res_samples + + if down_block_additional_residuals is not None: + new_down_block_res_samples = () + for down_block_res_sample, down_block_additional_residual in zip( + down_block_res_samples, down_block_additional_residuals + ): + down_block_res_sample = down_block_res_sample + down_block_additional_residual + new_down_block_res_samples += (down_block_res_sample,) + down_block_res_samples = new_down_block_res_samples + + # mid block with explicit reference + sample = self.mid_block( + sample, + emb, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + skip_mm=skip_mm, + mid_reference=m, + ) + + if mid_block_additional_residual is not None: + sample = sample + mid_block_additional_residual + + # up blocks with explicit reference injection + for i, upsample_block in enumerate(self.up_blocks): + up_reference = None + if i == 1: + up_reference = [u10, u11, u12] + elif i == 2: + up_reference = [u20, u21, u22] + elif i == 3: + up_reference = [u30, u31, u32] + + is_final_block = i == len(self.up_blocks) - 1 + + res_samples = down_block_res_samples[-len(upsample_block.resnets):] + down_block_res_samples = down_block_res_samples[:-len(upsample_block.resnets)] + + if not is_final_block and forward_upsample_size: + upsample_size = down_block_res_samples[-1].shape[2:] + + if ( + hasattr(upsample_block, "has_cross_attention") + and upsample_block.has_cross_attention + ): + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + encoder_hidden_states=encoder_hidden_states, + upsample_size=upsample_size, + attention_mask=attention_mask, + skip_mm=skip_mm, + up_reference=up_reference, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=res_samples, + upsample_size=upsample_size, + encoder_hidden_states=encoder_hidden_states, + skip_mm=skip_mm, + ) + + # post-process + sample = self.conv_norm_out(sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample) + + return sample + + @classmethod + def from_pretrained_2d( + cls, + pretrained_model_path: PathLike, + motion_module_path: PathLike, + subfolder=None, + unet_additional_kwargs=None, + mm_zero_proj_out=False, + ): + """Load from pretrained 2D UNet weights.""" + pretrained_model_path = Path(pretrained_model_path) + motion_module_path = Path(motion_module_path) + if subfolder is not None: + pretrained_model_path = pretrained_model_path.joinpath(subfolder) + logger.info(f"Loading temporal unet from {pretrained_model_path}...") + + config_file = pretrained_model_path / "config.json" + if not (config_file.exists() and config_file.is_file()): + raise RuntimeError(f"{config_file} does not exist or is not a file") + + unet_config = cls.load_config(config_file) + unet_config["_class_name"] = cls.__name__ + unet_config["down_block_types"] = [ + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "CrossAttnDownBlock3D", + "DownBlock3D", + ] + unet_config["up_block_types"] = [ + "UpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + "CrossAttnUpBlock3D", + ] + unet_config["mid_block_type"] = "UNetMidBlock3DCrossAttn" + + model = cls.from_config(unet_config, **unet_additional_kwargs) + + # Load weights + if pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME).exists(): + state_dict = load_file( + pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME), device="cpu" + ) + elif pretrained_model_path.joinpath(WEIGHTS_NAME).exists(): + state_dict = torch.load( + pretrained_model_path.joinpath(WEIGHTS_NAME), + map_location="cpu", + weights_only=True, + ) + else: + raise FileNotFoundError(f"No weights file found in {pretrained_model_path}") + + # Load motion module + if motion_module_path.exists() and motion_module_path.is_file(): + if motion_module_path.suffix.lower() in [".pth", ".pt", ".ckpt"]: + motion_state_dict = torch.load( + motion_module_path, map_location="cpu", weights_only=True + ) + elif motion_module_path.suffix.lower() == ".safetensors": + motion_state_dict = load_file(motion_module_path, device="cpu") + else: + raise RuntimeError(f"Unknown format: {motion_module_path.suffix}") + + motion_state_dict = { + k.replace('motion_modules.', 'temporal_modules.'): v + for k, v in motion_state_dict.items() if "pos_encoder" not in k + } + + if mm_zero_proj_out: + motion_state_dict = { + k: v for k, v in motion_state_dict.items() if "proj_out" not in k + } + + state_dict.update(motion_state_dict) + + m, u = model.load_state_dict(state_dict, strict=False) + logger.debug(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}") + + return model + From 5a455fae669f86fdbf577cceaf092852195f766a Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Wed, 17 Dec 2025 10:38:43 -0500 Subject: [PATCH 4/7] Hide warnings from polygraphy in console, fix to accumulation leak. Signed-off-by: BuffMcBigHuge --- src/scope/core/pipelines/personalive/pipeline.py | 2 ++ src/scope/core/pipelines/personalive/tensorrt/runner.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/src/scope/core/pipelines/personalive/pipeline.py b/src/scope/core/pipelines/personalive/pipeline.py index 9481e6d9..85bfae4e 100644 --- a/src/scope/core/pipelines/personalive/pipeline.py +++ b/src/scope/core/pipelines/personalive/pipeline.py @@ -505,6 +505,8 @@ def fuse_reference(self, ref_image: Image.Image): ref_image_latents = ref_image_latents * 0.18215 # Run reference UNet to cache features + # Clear writer first to prevent accumulation from previous keyframes when re-fusing + self.reference_control_writer.clear() self.reference_unet( ref_image_latents.to(self.reference_unet.device), torch.zeros((self.batch_size,), dtype=self.dtype, device=self.reference_unet.device), diff --git a/src/scope/core/pipelines/personalive/tensorrt/runner.py b/src/scope/core/pipelines/personalive/tensorrt/runner.py index 974523fb..2362478c 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/runner.py +++ b/src/scope/core/pipelines/personalive/tensorrt/runner.py @@ -23,6 +23,11 @@ from polygraphy.backend.common import BytesFromPath from polygraphy.backend.trt import EngineFromBytes, TrtRunner + # Suppress polygraphy deprecation warnings about DeviceView.dtype + # (warns about future API change: DataType.from_dtype(device_view.dtype).numpy()) + from polygraphy.logger import G_LOGGER + G_LOGGER.severity = G_LOGGER.ERROR # Only show errors, not warnings + TRT_AVAILABLE = True # Check for CUDA buffer support (polygraphy >= 0.47) From cdb9f8adbf8ab154a297843ba5cebfec42e3ff69 Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Wed, 17 Dec 2025 17:35:04 -0500 Subject: [PATCH 5/7] Better pipeline conditioning logic for ui. Signed-off-by: BuffMcBigHuge --- .../src/components/InputAndControlsPanel.tsx | 70 ++- frontend/src/components/SettingsPanel.tsx | 559 +++++++++--------- frontend/src/data/pipelines.ts | 139 ++++- frontend/src/pages/StreamPage.tsx | 265 +++++---- frontend/src/types/index.ts | 18 + .../core/pipelines/personalive/__init__.py | 1 + .../pipelines/personalive/blocks/__init__.py | 1 + .../personalive/components/face_detector.py | 1 + .../core/pipelines/personalive/docs/README.md | 1 + .../pipelines/personalive/tensorrt/builder.py | 1 + 10 files changed, 623 insertions(+), 433 deletions(-) diff --git a/frontend/src/components/InputAndControlsPanel.tsx b/frontend/src/components/InputAndControlsPanel.tsx index 176bdef2..7222fa65 100644 --- a/frontend/src/components/InputAndControlsPanel.tsx +++ b/frontend/src/components/InputAndControlsPanel.tsx @@ -12,10 +12,22 @@ import { Badge } from "./ui/badge"; import { Input } from "./ui/input"; import { Upload } from "lucide-react"; import { LabelWithTooltip } from "./ui/label-with-tooltip"; +import { + Tooltip, + TooltipContent, + TooltipProvider, + TooltipTrigger, +} from "./ui/tooltip"; import type { VideoSourceMode } from "../hooks/useVideoSource"; import type { PromptItem, PromptTransition } from "../lib/api"; import type { InputMode } from "../types"; -import { pipelineIsMultiMode, pipelineRequiresReferenceImage } from "../data/pipelines"; +import { + pipelineIsMultiMode, + pipelineRequiresReferenceImage, + pipelineShowsPromptInput, + pipelineCanChangeReferenceWhileStreaming, + getPipelineReferenceImageDescription, +} from "../data/pipelines"; import { PromptInput } from "./PromptInput"; import { TimelinePromptEditor } from "./TimelinePromptEditor"; import type { TimelinePrompt } from "./PromptTimeline"; @@ -192,7 +204,7 @@ export function InputAndControlsPanel({ )} - {/* Reference Image upload - only show for PersonaLive */} + {/* Reference Image upload - only show for pipelines that require it */} {needsReferenceImage && (

Reference Portrait

@@ -215,18 +227,42 @@ export function InputAndControlsPanel({ onChange={handleReferenceImageUpload} className="hidden" id="reference-image-upload" - disabled={isStreaming || isConnecting || isUploadingReference} + disabled={ + isUploadingReference || + ((isStreaming || isConnecting) && + !pipelineCanChangeReferenceWhileStreaming(pipelineId)) + } /> - + {/* Upload button with tooltip when disabled during streaming */} + {(isStreaming || isConnecting) && + !pipelineCanChangeReferenceWhileStreaming(pipelineId) ? ( + + + +
+ +
+
+ +

+ Reference image is processed when the pipeline loads. + Stop the stream to change it. +

+
+
+
+ ) : ( + + )}
{isUploadingReference && (

@@ -235,8 +271,8 @@ export function InputAndControlsPanel({ )} {!referenceImageUrl && (

- This pipeline animates this portrait using your webcam/video as the - driving source. + {getPipelineReferenceImageDescription(pipelineId) || + "This pipeline requires a reference image."}

)} @@ -352,8 +388,8 @@ export function InputAndControlsPanel({ )} - {/* Prompts section - hide for PersonaLive (it uses image conditioning, not text) */} - {!needsReferenceImage && ( + {/* Prompts section - only show for pipelines that support text prompts */} + {pipelineShowsPromptInput(pipelineId) && (
{(() => { // The Input can have two states: Append (default) and Edit (when a prompt is selected and the video is paused) diff --git a/frontend/src/components/SettingsPanel.tsx b/frontend/src/components/SettingsPanel.tsx index d21ca03a..2d7213c5 100644 --- a/frontend/src/components/SettingsPanel.tsx +++ b/frontend/src/components/SettingsPanel.tsx @@ -20,7 +20,17 @@ import { Button } from "./ui/button"; import { Toggle } from "./ui/toggle"; import { SliderWithInput } from "./ui/slider-with-input"; import { Hammer, Info, Minus, Plus, RotateCcw } from "lucide-react"; -import { PIPELINES, pipelineSupportsLoRA } from "../data/pipelines"; +import { + PIPELINES, + pipelineSupportsLoRA, + pipelineShowsResolutionControl, + pipelineShowsSeedControl, + pipelineShowsDenoisingSteps, + pipelineShowsCacheManagement, + pipelineShowsQuantization, + pipelineShowsKvCacheAttentionBias, + pipelineShowsNoiseControls, +} from "../data/pipelines"; import { PARAMETER_METADATA } from "../data/parameterMetadata"; import { DenoisingStepsSlider } from "./DenoisingStepsSlider"; import { useLocalSliderValue } from "../hooks/useLocalSliderValue"; @@ -69,8 +79,6 @@ interface SettingsPanelProps { loraMergeStrategy?: LoraMergeStrategy; // Input mode for conditional rendering of noise controls inputMode?: InputMode; - // Whether this pipeline supports noise controls in video mode (schema-derived) - supportsNoiseControls?: boolean; // Spout settings spoutSender?: SettingsState["spoutSender"]; onSpoutSenderChange?: (spoutSender: SettingsState["spoutSender"]) => void; @@ -106,7 +114,6 @@ export function SettingsPanel({ onLorasChange, loraMergeStrategy = "permanent_merge", inputMode, - supportsNoiseControls = false, spoutSender, onSpoutSenderChange, spoutAvailable = false, @@ -331,231 +338,221 @@ export function SettingsPanel({
)} - {(pipelineId === "longlive" || - pipelineId === "streamdiffusionv2" || - pipelineId === "krea-realtime-video" || - pipelineId === "reward-forcing") && ( -
-
-
-
-
- -
- - { - const value = parseInt(e.target.value); - if (!isNaN(value)) { - handleResolutionChange("height", value); - } - }} - disabled={isStreaming} - className="text-center border-0 focus-visible:ring-0 focus-visible:ring-offset-0 h-8 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none" - min={MIN_DIMENSION} - max={2048} - /> - -
-
- {heightError && ( -

{heightError}

- )} -
- -
-
- -
- - { - const value = parseInt(e.target.value); - if (!isNaN(value)) { - handleResolutionChange("width", value); - } - }} - disabled={isStreaming} - className="text-center border-0 focus-visible:ring-0 focus-visible:ring-offset-0 h-8 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none" - min={MIN_DIMENSION} - max={2048} - /> - -
-
- {widthError && ( -

{widthError}

- )} -
- -
-
- -
- - { - const value = parseInt(e.target.value); - if (!isNaN(value)) { - handleSeedChange(value); - } - }} - disabled={isStreaming} - className="text-center border-0 focus-visible:ring-0 focus-visible:ring-offset-0 h-8 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none" - min={0} - max={2147483647} - /> - -
-
- {seedError && ( -

{seedError}

- )} + {/* Resolution controls */} + {pipelineShowsResolutionControl(pipelineId) && ( +
+
+
+ +
+ + { + const value = parseInt(e.target.value); + if (!isNaN(value)) { + handleResolutionChange("height", value); + } + }} + disabled={isStreaming} + className="text-center border-0 focus-visible:ring-0 focus-visible:ring-offset-0 h-8 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none" + min={MIN_DIMENSION} + max={2048} + /> +
+ {heightError && ( +

{heightError}

+ )}
-
- )} - - {(pipelineId === "longlive" || - pipelineId === "streamdiffusionv2" || - pipelineId === "krea-realtime-video" || - pipelineId === "reward-forcing") && ( -
-
-
- {pipelineId === "krea-realtime-video" && ( - parseFloat(v) || 1.0} - /> - )} -
- - {})} - variant="outline" - size="sm" - className="h-7" +
+
+ +
+
- -
- + + { + const value = parseInt(e.target.value); + if (!isNaN(value)) { + handleResolutionChange("width", value); + } + }} + disabled={isStreaming} + className="text-center border-0 focus-visible:ring-0 focus-visible:ring-offset-0 h-8 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none" + min={MIN_DIMENSION} + max={2048} />
+ {widthError && ( +

{widthError}

+ )}
)} - {(pipelineId === "longlive" || - pipelineId === "streamdiffusionv2" || - pipelineId === "krea-realtime-video" || - pipelineId === "reward-forcing") && ( + {/* Seed control */} + {pipelineShowsSeedControl(pipelineId) && ( +
+
+ +
+ + { + const value = parseInt(e.target.value); + if (!isNaN(value)) { + handleSeedChange(value); + } + }} + disabled={isStreaming} + className="text-center border-0 focus-visible:ring-0 focus-visible:ring-offset-0 h-8 [appearance:textfield] [&::-webkit-outer-spin-button]:appearance-none [&::-webkit-inner-spin-button]:appearance-none" + min={0} + max={2147483647} + /> + +
+
+ {seedError && ( +

{seedError}

+ )} +
+ )} + + {/* KV Cache Attention Bias (krea-realtime-video specific) */} + {pipelineShowsKvCacheAttentionBias(pipelineId) && ( + parseFloat(v) || 1.0} + /> + )} + + {/* Cache management controls */} + {pipelineShowsCacheManagement(pipelineId) && ( +
+
+ + {})} + variant="outline" + size="sm" + className="h-7" + > + {manageCache ? "ON" : "OFF"} + +
+ +
+ + +
+
+ )} + + {/* Denoising steps slider */} + {pipelineShowsDenoisingSteps(pipelineId) && ( {})} @@ -564,84 +561,72 @@ export function SettingsPanel({ /> )} - {/* Noise controls - show for video mode on supported pipelines (schema-derived) */} - {inputMode === "video" && supportsNoiseControls && ( -
-
-
-
- - {})} - disabled={isStreaming} - variant="outline" - size="sm" - className="h-7" - > - {noiseController ? "ON" : "OFF"} - -
-
- - parseFloat(v) || 0.0} + {/* Noise controls - show for video mode on pipelines that support it */} + {inputMode === "video" && pipelineShowsNoiseControls(pipelineId) && ( +
+
+ + {})} + disabled={isStreaming} + variant="outline" + size="sm" + className="h-7" + > + {noiseController ? "ON" : "OFF"} +
+ + parseFloat(v) || 0.0} + />
)} - {(pipelineId === "longlive" || - pipelineId === "streamdiffusionv2" || - pipelineId === "krea-realtime-video" || - pipelineId === "reward-forcing") && ( -
-
-
-
- - -
-
-
+ {/* Quantization selector */} + {pipelineShowsQuantization(pipelineId) && ( +
+ +
)} diff --git a/frontend/src/data/pipelines.ts b/frontend/src/data/pipelines.ts index e380950e..52dbd320 100644 --- a/frontend/src/data/pipelines.ts +++ b/frontend/src/data/pipelines.ts @@ -8,6 +8,20 @@ export const DEFAULT_PROMPTS: Record = { "A 3D animated scene. A **panda** sitting in the grass, looking around.", }; +// UI capability flags - controls which UI elements are shown for each pipeline +export interface PipelineUICapabilities { + showTimeline?: boolean; // Show prompt timeline (default: true) + showPromptInput?: boolean; // Show prompt input controls (default: true) + showResolutionControl?: boolean; // Show resolution width/height controls + showSeedControl?: boolean; // Show seed control + showDenoisingSteps?: boolean; // Show denoising steps slider + showCacheManagement?: boolean; // Show cache management toggle and reset + showQuantization?: boolean; // Show quantization selector + showKvCacheAttentionBias?: boolean; // Show KV cache attention bias slider (krea-realtime-video) + showNoiseControls?: boolean; // Show noise scale and controller (video mode) + canChangeReferenceWhileStreaming?: boolean; // Allow changing reference image during stream +} + export interface PipelineInfo { name: string; about: string; @@ -19,10 +33,14 @@ export interface PipelineInfo { defaultTemporalInterpolationSteps?: number; // Default number of steps for temporal interpolation supportsLoRA?: boolean; // Whether this pipeline supports LoRA adapters requiresReferenceImage?: boolean; // Whether this pipeline requires a reference image (e.g. PersonaLive) + referenceImageDescription?: string; // Description of what the reference image is for // Multi-mode support supportedModes: InputMode[]; defaultMode: InputMode; + + // UI capabilities - controls which settings/controls are shown + ui?: PipelineUICapabilities; } export const PIPELINES: Record = { @@ -38,9 +56,18 @@ export const PIPELINES: Record = { defaultTemporalInterpolationMethod: "slerp", defaultTemporalInterpolationSteps: 0, supportsLoRA: true, - // Multi-mode support supportedModes: ["text", "video"], defaultMode: "video", + ui: { + showTimeline: true, + showPromptInput: true, + showResolutionControl: true, + showSeedControl: true, + showDenoisingSteps: true, + showCacheManagement: true, + showQuantization: true, + showNoiseControls: true, + }, }, longlive: { name: "LongLive", @@ -54,9 +81,18 @@ export const PIPELINES: Record = { defaultTemporalInterpolationMethod: "slerp", defaultTemporalInterpolationSteps: 0, supportsLoRA: true, - // Multi-mode support supportedModes: ["text", "video"], defaultMode: "text", + ui: { + showTimeline: true, + showPromptInput: true, + showResolutionControl: true, + showSeedControl: true, + showDenoisingSteps: true, + showCacheManagement: true, + showQuantization: true, + showNoiseControls: true, + }, }, "krea-realtime-video": { name: "Krea Realtime Video", @@ -70,9 +106,19 @@ export const PIPELINES: Record = { defaultTemporalInterpolationMethod: "linear", defaultTemporalInterpolationSteps: 4, supportsLoRA: true, - // Multi-mode support supportedModes: ["text", "video"], defaultMode: "text", + ui: { + showTimeline: true, + showPromptInput: true, + showResolutionControl: true, + showSeedControl: true, + showDenoisingSteps: true, + showCacheManagement: true, + showQuantization: true, + showKvCacheAttentionBias: true, + showNoiseControls: true, + }, }, "reward-forcing": { name: "RewardForcing", @@ -86,18 +132,36 @@ export const PIPELINES: Record = { defaultTemporalInterpolationMethod: "slerp", defaultTemporalInterpolationSteps: 0, supportsLoRA: true, - // Multi-mode support supportedModes: ["text", "video"], defaultMode: "text", + ui: { + showTimeline: true, + showPromptInput: true, + showResolutionControl: true, + showSeedControl: true, + showDenoisingSteps: true, + showCacheManagement: true, + showQuantization: true, + showNoiseControls: true, + }, }, passthrough: { name: "Passthrough", about: "A pipeline that returns the input video without any processing that is useful for testing and debugging.", requiresModels: false, - // Video-only pipeline supportedModes: ["video"], defaultMode: "video", + ui: { + showTimeline: false, + showPromptInput: false, + showResolutionControl: false, + showSeedControl: false, + showDenoisingSteps: false, + showCacheManagement: false, + showQuantization: false, + showNoiseControls: false, + }, }, personalive: { name: "PersonaLive", @@ -106,11 +170,22 @@ export const PIPELINES: Record = { "Real-time portrait animation pipeline from GVCLab. Animates a reference portrait image using driving video frames to transfer expressions and head movements.", estimatedVram: 12, requiresModels: true, - // Video-only pipeline - requires reference image + driving video supportedModes: ["video"], defaultMode: "video", - // PersonaLive requires a reference image to be uploaded first requiresReferenceImage: true, + referenceImageDescription: + "Portrait image to animate. Expressions and head movements from the driving video will be transferred to this image.", + ui: { + showTimeline: false, + showPromptInput: false, + showResolutionControl: true, + showSeedControl: true, + showDenoisingSteps: false, + showCacheManagement: false, + showQuantization: false, + showNoiseControls: false, + canChangeReferenceWhileStreaming: false, + }, }, }; @@ -141,3 +216,53 @@ export function getDefaultPromptForMode(mode: InputMode): string { export function pipelineRequiresReferenceImage(pipelineId: string): boolean { return PIPELINES[pipelineId]?.requiresReferenceImage === true; } + +export function getPipelineReferenceImageDescription( + pipelineId: string +): string | undefined { + return PIPELINES[pipelineId]?.referenceImageDescription; +} + +// UI capability helper functions + +export function pipelineShowsTimeline(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showTimeline !== false; +} + +export function pipelineShowsPromptInput(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showPromptInput !== false; +} + +export function pipelineShowsResolutionControl(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showResolutionControl === true; +} + +export function pipelineShowsSeedControl(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showSeedControl === true; +} + +export function pipelineShowsDenoisingSteps(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showDenoisingSteps === true; +} + +export function pipelineShowsCacheManagement(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showCacheManagement === true; +} + +export function pipelineShowsQuantization(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showQuantization === true; +} + +export function pipelineShowsKvCacheAttentionBias(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showKvCacheAttentionBias === true; +} + +export function pipelineShowsNoiseControls(pipelineId: string): boolean { + return PIPELINES[pipelineId]?.ui?.showNoiseControls === true; +} + +export function pipelineCanChangeReferenceWhileStreaming( + pipelineId: string +): boolean { + return PIPELINES[pipelineId]?.ui?.canChangeReferenceWhileStreaming !== false; +} diff --git a/frontend/src/pages/StreamPage.tsx b/frontend/src/pages/StreamPage.tsx index 0996119b..56698b7b 100644 --- a/frontend/src/pages/StreamPage.tsx +++ b/frontend/src/pages/StreamPage.tsx @@ -17,6 +17,7 @@ import { getPipelineDefaultMode, getDefaultPromptForMode, pipelineRequiresReferenceImage, + pipelineShowsTimeline, } from "../data/pipelines"; import type { InputMode, @@ -31,6 +32,7 @@ import { uploadPersonaLiveReference, } from "../lib/api"; import { sendLoRAScaleUpdates } from "../utils/loraHelpers"; +import { toast } from "sonner"; // Delay before resetting video reinitialization flag (ms) // This allows useVideoSource to detect the flag change and trigger reinitialization @@ -55,13 +57,8 @@ function buildLoRAParams( export function StreamPage() { // Use the stream state hook for settings management - const { - settings, - updateSettings, - getDefaults, - supportsNoiseControls, - spoutAvailable, - } = useStreamState(); + const { settings, updateSettings, getDefaults, spoutAvailable } = + useStreamState(); // Prompt state - use unified default prompts based on mode const initialMode = @@ -238,6 +235,10 @@ export function StreamPage() { }; const handlePipelineIdChange = (pipelineId: PipelineId) => { + console.log( + `[handlePipelineIdChange] Switching to pipeline: ${pipelineId}` + ); + // Stop the stream if it's currently running if (isStreaming) { stopStream(); @@ -634,6 +635,11 @@ export function StreamPage() { // Use override pipeline ID if provided, otherwise use current settings const pipelineIdToUse = overridePipelineId || settings.pipelineId; + // Debug: Log which pipeline is being started + console.log( + `[handleStartStream] Starting pipeline: ${pipelineIdToUse} (settings.pipelineId: ${settings.pipelineId}, override: ${overridePipelineId})` + ); + try { // Check if models are needed but not downloaded const pipelineInfo = PIPELINES[pipelineIdToUse]; @@ -642,6 +648,9 @@ export function StreamPage() { const status = await checkModelStatus(pipelineIdToUse); if (!status.downloaded) { // Show download dialog + console.log( + `[handleStartStream] Models not downloaded for: ${pipelineIdToUse}, showing download dialog` + ); setPipelineNeedsModels(pipelineIdToUse); setShowDownloadDialog(true); return false; // Stream did not start @@ -652,9 +661,6 @@ export function StreamPage() { } } - // Always load pipeline with current parameters - backend will handle the rest - console.log(`Loading ${pipelineIdToUse} pipeline...`); - // Determine current input mode const currentMode = settings.inputMode || getPipelineDefaultMode(pipelineIdToUse) || "text"; @@ -723,9 +729,11 @@ export function StreamPage() { } else if (pipelineIdToUse === "personalive" && resolution) { // PersonaLive requires a reference image if (!referenceImage) { - console.error( - "PersonaLive requires a reference image. Please upload one first." - ); + toast.error("Reference Image Required", { + description: + "Please upload a reference portrait image before starting PersonaLive.", + duration: 5000, + }); return false; } loadParams = { @@ -929,15 +937,27 @@ export function StreamPage() { isPlaying={!settings.paused} isDownloading={isDownloading} onPlayPauseToggle={() => { - // Use timeline's play/pause handler instead of direct video toggle - if (timelinePlayPauseRef.current) { + // For pipelines with timeline, use timeline's play/pause handler + // For pipelines without timeline (PersonaLive, Passthrough), toggle directly + if ( + pipelineShowsTimeline(settings.pipelineId) && + timelinePlayPauseRef.current + ) { timelinePlayPauseRef.current(); + } else { + handlePlayPauseToggle(); } }} onStartStream={() => { - // Use timeline's play/pause handler to start stream - if (timelinePlayPauseRef.current) { + // For pipelines with timeline, use timeline's play/pause handler + // For pipelines without timeline (PersonaLive, Passthrough), start directly + if ( + pipelineShowsTimeline(settings.pipelineId) && + timelinePlayPauseRef.current + ) { timelinePlayPauseRef.current(); + } else { + handleStartStream(); } }} onVideoPlaying={() => { @@ -949,111 +969,113 @@ export function StreamPage() { }} />
- {/* Timeline area - compact, always visible */} -
- { - // Update the left panel's prompt state to reflect current timeline prompt - const prompts = [{ text, weight: 100 }]; - setPromptItems(prompts); - - // Send to backend - use transition if streaming and transition steps > 0 - if (isStreaming && transitionSteps > 0) { - sendParameterUpdate({ - transition: { - target_prompts: prompts, - num_steps: transitionSteps, - temporal_interpolation_method: - temporalInterpolationMethod, - }, - }); - } else { - // Send direct prompts without transition - sendParameterUpdate({ - prompts, - prompt_interpolation_method: interpolationMethod, - denoising_step_list: settings.denoisingSteps || [700, 500], - }); - } - }} - onPromptItemsSubmit={( - prompts, - blockTransitionSteps, - blockTemporalInterpolationMethod - ) => { - // Update the left panel's prompt state to reflect current timeline prompt blend - setPromptItems(prompts); - - // Use transition params from block if provided, otherwise use global settings - const effectiveTransitionSteps = - blockTransitionSteps ?? transitionSteps; - const effectiveTemporalInterpolationMethod = - blockTemporalInterpolationMethod ?? - temporalInterpolationMethod; - - // Update the left panel's transition settings to reflect current block's values - if (blockTransitionSteps !== undefined) { - setTransitionSteps(blockTransitionSteps); + {/* Timeline area - only show for pipelines that support text prompts */} + {pipelineShowsTimeline(settings.pipelineId) && ( +
+ { + // Update the left panel's prompt state to reflect current timeline prompt + const prompts = [{ text, weight: 100 }]; + setPromptItems(prompts); + + // Send to backend - use transition if streaming and transition steps > 0 + if (isStreaming && transitionSteps > 0) { + sendParameterUpdate({ + transition: { + target_prompts: prompts, + num_steps: transitionSteps, + temporal_interpolation_method: + temporalInterpolationMethod, + }, + }); + } else { + // Send direct prompts without transition + sendParameterUpdate({ + prompts, + prompt_interpolation_method: interpolationMethod, + denoising_step_list: settings.denoisingSteps || [ + 700, 500, + ], + }); + } + }} + onPromptItemsSubmit={( + prompts, + blockTransitionSteps, + blockTemporalInterpolationMethod + ) => { + // Update the left panel's prompt state to reflect current timeline prompt blend + setPromptItems(prompts); + + // Use transition params from block if provided, otherwise use global settings + const effectiveTransitionSteps = + blockTransitionSteps ?? transitionSteps; + const effectiveTemporalInterpolationMethod = + blockTemporalInterpolationMethod ?? + temporalInterpolationMethod; + + // Update the left panel's transition settings to reflect current block's values + if (blockTransitionSteps !== undefined) { + setTransitionSteps(blockTransitionSteps); + } + if (blockTemporalInterpolationMethod !== undefined) { + setTemporalInterpolationMethod( + blockTemporalInterpolationMethod + ); + } + + // Send to backend - use transition if streaming and transition steps > 0 + if (isStreaming && effectiveTransitionSteps > 0) { + sendParameterUpdate({ + transition: { + target_prompts: prompts, + num_steps: effectiveTransitionSteps, + temporal_interpolation_method: + effectiveTemporalInterpolationMethod, + }, + }); + } else { + // Send direct prompts without transition + sendParameterUpdate({ + prompts, + prompt_interpolation_method: interpolationMethod, + denoising_step_list: settings.denoisingSteps || [ + 700, 500, + ], + }); + } + }} + disabled={ + isPipelineLoading || isConnecting || showDownloadDialog } - if (blockTemporalInterpolationMethod !== undefined) { - setTemporalInterpolationMethod( - blockTemporalInterpolationMethod - ); - } - - // Send to backend - use transition if streaming and transition steps > 0 - if (isStreaming && effectiveTransitionSteps > 0) { - sendParameterUpdate({ - transition: { - target_prompts: prompts, - num_steps: effectiveTransitionSteps, - temporal_interpolation_method: - effectiveTemporalInterpolationMethod, - }, - }); - } else { - // Send direct prompts without transition - sendParameterUpdate({ - prompts, - prompt_interpolation_method: interpolationMethod, - denoising_step_list: settings.denoisingSteps || [700, 500], - }); - } - }} - disabled={ - settings.pipelineId === "passthrough" || - settings.pipelineId === "personalive" || - isPipelineLoading || - isConnecting || - showDownloadDialog - } - isStreaming={isStreaming} - isVideoPaused={settings.paused} - timelineRef={timelineRef} - onLiveStateChange={setIsLive} - onLivePromptSubmit={handleLivePromptSubmit} - onDisconnect={stopStream} - onStartStream={handleStartStream} - onVideoPlayPauseToggle={handlePlayPauseToggle} - onPromptEdit={handleTimelinePromptEdit} - isCollapsed={isTimelineCollapsed} - onCollapseToggle={setIsTimelineCollapsed} - externalSelectedPromptId={externalSelectedPromptId} - settings={settings} - onSettingsImport={updateSettings} - onPlayPauseRef={timelinePlayPauseRef} - onVideoPlayingCallbackRef={onVideoPlayingCallbackRef} - onResetCache={handleResetCache} - onTimelinePromptsChange={handleTimelinePromptsChange} - onTimelineCurrentTimeChange={handleTimelineCurrentTimeChange} - onTimelinePlayingChange={handleTimelinePlayingChange} - isDownloading={isDownloading} - /> -
+ isStreaming={isStreaming} + isVideoPaused={settings.paused} + timelineRef={timelineRef} + onLiveStateChange={setIsLive} + onLivePromptSubmit={handleLivePromptSubmit} + onDisconnect={stopStream} + onStartStream={handleStartStream} + onVideoPlayPauseToggle={handlePlayPauseToggle} + onPromptEdit={handleTimelinePromptEdit} + isCollapsed={isTimelineCollapsed} + onCollapseToggle={setIsTimelineCollapsed} + externalSelectedPromptId={externalSelectedPromptId} + settings={settings} + onSettingsImport={updateSettings} + onPlayPauseRef={timelinePlayPauseRef} + onVideoPlayingCallbackRef={onVideoPlayingCallbackRef} + onResetCache={handleResetCache} + onTimelinePromptsChange={handleTimelinePromptsChange} + onTimelineCurrentTimeChange={handleTimelineCurrentTimeChange} + onTimelinePlayingChange={handleTimelinePlayingChange} + isDownloading={isDownloading} + /> +
+ )}
{/* Right Panel - Settings */} @@ -1104,7 +1126,6 @@ export function StreamPage() { onLorasChange={handleLorasChange} loraMergeStrategy={settings.loraMergeStrategy ?? "permanent_merge"} inputMode={settings.inputMode} - supportsNoiseControls={supportsNoiseControls(settings.pipelineId)} spoutSender={settings.spoutSender} onSpoutSenderChange={handleSpoutSenderChange} spoutAvailable={spoutAvailable} diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index ac4c0c74..88d661c3 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -76,6 +76,20 @@ export interface SettingsState { }; } +// UI capability flags - controls which UI elements are shown for each pipeline +export interface PipelineUICapabilities { + showTimeline?: boolean; // Show prompt timeline (default: true) + showPromptInput?: boolean; // Show prompt input controls (default: true) + showResolutionControl?: boolean; // Show resolution width/height controls + showSeedControl?: boolean; // Show seed control + showDenoisingSteps?: boolean; // Show denoising steps slider + showCacheManagement?: boolean; // Show cache management toggle and reset + showQuantization?: boolean; // Show quantization selector + showKvCacheAttentionBias?: boolean; // Show KV cache attention bias slider (krea-realtime-video) + showNoiseControls?: boolean; // Show noise scale and controller (video mode) + canChangeReferenceWhileStreaming?: boolean; // Allow changing reference image during stream +} + export interface PipelineInfo { name: string; about: string; @@ -89,8 +103,12 @@ export interface PipelineInfo { defaultTemporalInterpolationSteps?: number; supportsLoRA?: boolean; requiresReferenceImage?: boolean; // Whether this pipeline requires a reference image (e.g. PersonaLive) + referenceImageDescription?: string; // Description of what the reference image is for // Multi-mode support supportedModes: InputMode[]; defaultMode: InputMode; + + // UI capabilities - controls which settings/controls are shown + ui?: PipelineUICapabilities; } diff --git a/src/scope/core/pipelines/personalive/__init__.py b/src/scope/core/pipelines/personalive/__init__.py index f3995ef9..266517a5 100644 --- a/src/scope/core/pipelines/personalive/__init__.py +++ b/src/scope/core/pipelines/personalive/__init__.py @@ -3,3 +3,4 @@ from .pipeline import PersonaLivePipeline __all__ = ["PersonaLivePipeline"] + diff --git a/src/scope/core/pipelines/personalive/blocks/__init__.py b/src/scope/core/pipelines/personalive/blocks/__init__.py index fdb7b942..ebc248ee 100644 --- a/src/scope/core/pipelines/personalive/blocks/__init__.py +++ b/src/scope/core/pipelines/personalive/blocks/__init__.py @@ -1 +1,2 @@ """PersonaLive modular blocks for pipeline composition.""" + diff --git a/src/scope/core/pipelines/personalive/components/face_detector.py b/src/scope/core/pipelines/personalive/components/face_detector.py index 9724c95c..e011aac3 100644 --- a/src/scope/core/pipelines/personalive/components/face_detector.py +++ b/src/scope/core/pipelines/personalive/components/face_detector.py @@ -197,3 +197,4 @@ def close(self): def __del__(self): self.close() + diff --git a/src/scope/core/pipelines/personalive/docs/README.md b/src/scope/core/pipelines/personalive/docs/README.md index da470be1..affdfad0 100644 --- a/src/scope/core/pipelines/personalive/docs/README.md +++ b/src/scope/core/pipelines/personalive/docs/README.md @@ -122,3 +122,4 @@ PersonaLive uses a multi-component architecture: ## Credits Based on the PersonaLive implementation by GVCLab. + diff --git a/src/scope/core/pipelines/personalive/tensorrt/builder.py b/src/scope/core/pipelines/personalive/tensorrt/builder.py index 12e63f16..a97204be 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/builder.py +++ b/src/scope/core/pipelines/personalive/tensorrt/builder.py @@ -241,3 +241,4 @@ def is_engine_available(model_dir: Path, height: int = 512, width: int = 512) -> """ engine_path = get_engine_path(model_dir, height, width) return engine_path.exists() + From edf2b786909ff10e2c0ed472695e1e76799ea426 Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Wed, 17 Dec 2025 17:38:56 -0500 Subject: [PATCH 6/7] Linting and formatting. Signed-off-by: BuffMcBigHuge --- frontend/src/components/SettingsPanel.tsx | 4 +- .../core/pipelines/personalive/__init__.py | 1 - .../pipelines/personalive/blocks/__init__.py | 1 - .../personalive/components/face_detector.py | 8 +- .../core/pipelines/personalive/pipeline.py | 210 +++++++++++++----- .../personalive/tensorrt/__init__.py | 14 +- .../pipelines/personalive/tensorrt/builder.py | 24 +- .../pipelines/personalive/tensorrt/convert.py | 19 +- .../personalive/tensorrt/engine_model.py | 82 +++++-- .../personalive/tensorrt/framed_model.py | 91 ++++++-- .../pipelines/personalive/tensorrt/runner.py | 26 ++- .../personalive/tensorrt/unet_explicit.py | 29 ++- src/scope/server/app.py | 11 +- src/scope/server/download_models.py | 4 +- src/scope/server/models_config.py | 48 +++- src/scope/server/schema.py | 4 +- 16 files changed, 428 insertions(+), 148 deletions(-) diff --git a/frontend/src/components/SettingsPanel.tsx b/frontend/src/components/SettingsPanel.tsx index 2d7213c5..613118c7 100644 --- a/frontend/src/components/SettingsPanel.tsx +++ b/frontend/src/components/SettingsPanel.tsx @@ -622,9 +622,7 @@ export function SettingsPanel({ None - - fp8_e4m3fn (Dynamic) - + fp8_e4m3fn (Dynamic)
diff --git a/src/scope/core/pipelines/personalive/__init__.py b/src/scope/core/pipelines/personalive/__init__.py index 266517a5..f3995ef9 100644 --- a/src/scope/core/pipelines/personalive/__init__.py +++ b/src/scope/core/pipelines/personalive/__init__.py @@ -3,4 +3,3 @@ from .pipeline import PersonaLivePipeline __all__ = ["PersonaLivePipeline"] - diff --git a/src/scope/core/pipelines/personalive/blocks/__init__.py b/src/scope/core/pipelines/personalive/blocks/__init__.py index ebc248ee..fdb7b942 100644 --- a/src/scope/core/pipelines/personalive/blocks/__init__.py +++ b/src/scope/core/pipelines/personalive/blocks/__init__.py @@ -1,2 +1 @@ """PersonaLive modular blocks for pipeline composition.""" - diff --git a/src/scope/core/pipelines/personalive/components/face_detector.py b/src/scope/core/pipelines/personalive/components/face_detector.py index e011aac3..bea0e23c 100644 --- a/src/scope/core/pipelines/personalive/components/face_detector.py +++ b/src/scope/core/pipelines/personalive/components/face_detector.py @@ -10,6 +10,7 @@ try: import mediapipe as mp + MEDIAPIPE_AVAILABLE = True except ImportError: MEDIAPIPE_AVAILABLE = False @@ -58,9 +59,7 @@ def detect_face_landmarks(self, image: np.ndarray) -> list | None: return None def get_face_bounding_box( - self, - image: np.ndarray, - padding_ratio: float = 0.2 + self, image: np.ndarray, padding_ratio: float = 0.2 ) -> Tuple[int, int, int, int] | None: """Get bounding box for the detected face. @@ -192,9 +191,8 @@ def crop_face_tensor( def close(self): """Release resources.""" - if hasattr(self, 'face_mesh') and self.face_mesh: + if hasattr(self, "face_mesh") and self.face_mesh: self.face_mesh.close() def __del__(self): self.close() - diff --git a/src/scope/core/pipelines/personalive/pipeline.py b/src/scope/core/pipelines/personalive/pipeline.py index 85bfae4e..c8904e8f 100644 --- a/src/scope/core/pipelines/personalive/pipeline.py +++ b/src/scope/core/pipelines/personalive/pipeline.py @@ -84,6 +84,7 @@ def _profile_stage(name: str): elapsed = time.perf_counter() - start _profiling_timings[name].append(elapsed) + # TensorRT support (optional) try: from .tensorrt import TRT_AVAILABLE, TRTRunner, get_engine_path, PYCUDA_AVAILABLE @@ -199,7 +200,9 @@ def __init__( # Load pose encoder (motion extractor) start = time.time() - self.pose_encoder = MotionExtractor(num_kp=21).to(device=device, dtype=dtype).eval() + self.pose_encoder = ( + MotionExtractor(num_kp=21).to(device=device, dtype=dtype).eval() + ) pose_encoder_path = personalive_weights / "motion_extractor.pth" if pose_encoder_path.exists(): state_dict = torch.load(pose_encoder_path, map_location="cpu") @@ -291,7 +294,9 @@ def __init__( ) self.clip_image_processor = CLIPImageProcessor() self.cond_image_processor = VaeImageProcessor( - vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=True + vae_scale_factor=self.vae_scale_factor, + do_convert_rgb=True, + do_normalize=True, ) # Face detector @@ -365,14 +370,18 @@ def _init_tensorrt(self, model_dir: Path): self._use_engine_model = True # Setup output-to-input bindings for recurrent state (zero-copy) - self.trt_runner.bind({ - "motion_hidden_states_out": "motion_hidden_states", - "pose_cond_fea_out": "pose_cond_fea", - "latents": "sample", - }) + self.trt_runner.bind( + { + "motion_hidden_states_out": "motion_hidden_states", + "pose_cond_fea_out": "pose_cond_fea", + "latents": "sample", + } + ) self.use_tensorrt = True - logger.info(f"TensorRT engine loaded with pycuda (EngineModel) from {engine_path}") + logger.info( + f"TensorRT engine loaded with pycuda (EngineModel) from {engine_path}" + ) return except Exception as e: @@ -385,14 +394,18 @@ def _init_tensorrt(self, model_dir: Path): self._use_engine_model = False # Setup output-to-input bindings for recurrent state - self.trt_runner.bind({ - "motion_hidden_states_out": "motion_hidden_states", - "pose_cond_fea_out": "pose_cond_fea", - "latents": "sample", - }) + self.trt_runner.bind( + { + "motion_hidden_states_out": "motion_hidden_states", + "pose_cond_fea_out": "pose_cond_fea", + "latents": "sample", + } + ) self.use_tensorrt = True - logger.info(f"TensorRT engine loaded with polygraphy (TRTRunner) from {engine_path}") + logger.info( + f"TensorRT engine loaded with polygraphy (TRTRunner) from {engine_path}" + ) except Exception as e: logger.warning(f"Failed to load TensorRT engine: {e}") @@ -417,7 +430,9 @@ def prepare(self, **kwargs) -> Requirements | None: return Requirements(input_size=self.temporal_window_size) return None - def _fast_resize(self, images: torch.Tensor, height: int, width: int) -> torch.Tensor: + def _fast_resize( + self, images: torch.Tensor, height: int, width: int + ) -> torch.Tensor: """Fast bilinear resize of image tensor.""" return F.interpolate( images, @@ -426,7 +441,9 @@ def _fast_resize(self, images: torch.Tensor, height: int, width: int) -> torch.T align_corners=False, ) - def _interpolate_tensors(self, a: torch.Tensor, b: torch.Tensor, num: int) -> torch.Tensor: + def _interpolate_tensors( + self, a: torch.Tensor, b: torch.Tensor, num: int + ) -> torch.Tensor: """Linear interpolation between tensors a and b.""" if a.shape != b.shape: raise ValueError(f"Shape mismatch: a.shape={a.shape}, b.shape={b.shape}") @@ -447,7 +464,9 @@ def _calculate_dis(self, A: torch.Tensor, B: torch.Tensor, threshold: float = 10 dist = torch.cdist(B_flat.to(torch.float32), A_flat.to(torch.float32), p=2) min_dist, min_idx = dist.min(dim=1) - idx_to_add = torch.nonzero(min_dist[:1] > threshold, as_tuple=False).squeeze(1).tolist() + idx_to_add = ( + torch.nonzero(min_dist[:1] > threshold, as_tuple=False).squeeze(1).tolist() + ) if len(idx_to_add) > 0: B_to_add = B[:, idx_to_add] @@ -499,7 +518,9 @@ def fuse_reference(self, ref_image: Image.Image): self.encoder_hidden_states = clip_image_embeds.unsqueeze(1) # Encode reference image - ref_image_tensor = ref_image_tensor.to(dtype=self.vae.dtype, device=self.vae.device) + ref_image_tensor = ref_image_tensor.to( + dtype=self.vae.dtype, device=self.vae.device + ) self.ref_image_tensor = ref_image_tensor.squeeze(0) ref_image_latents = self.vae.encode(ref_image_tensor).latent_dist.mean ref_image_latents = ref_image_latents * 0.18215 @@ -509,7 +530,9 @@ def fuse_reference(self, ref_image: Image.Image): self.reference_control_writer.clear() self.reference_unet( ref_image_latents.to(self.reference_unet.device), - torch.zeros((self.batch_size,), dtype=self.dtype, device=self.reference_unet.device), + torch.zeros( + (self.batch_size,), dtype=self.dtype, device=self.reference_unet.device + ), encoder_hidden_states=self.encoder_hidden_states, return_dict=False, ) @@ -521,14 +544,28 @@ def fuse_reference(self, ref_image: Image.Image): self._reference_hidden_states = reference_hidden_states # Clear previous prefill values (TRTRunner only - EngineModel persists buffers) - if hasattr(self.trt_runner, 'clear_prefill'): + if hasattr(self.trt_runner, "clear_prefill"): self.trt_runner.clear_prefill() self.trt_runner.prefill(encoder_hidden_states=self.encoder_hidden_states) # Prefill reference hidden states ref_hidden_names = [ - "d00", "d01", "d10", "d11", "d20", "d21", "m", - "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" + "d00", + "d01", + "d10", + "d11", + "d20", + "d21", + "m", + "u10", + "u11", + "u12", + "u20", + "u21", + "u22", + "u30", + "u31", + "u32", ] for name in ref_hidden_names: if name in reference_hidden_states: @@ -559,8 +596,12 @@ def fuse_reference(self, ref_image: Image.Image): padding_num = (self.temporal_adaptive_step - 1) * self.temporal_window_size init_latents = ref_image_latents.unsqueeze(2).repeat(1, 1, padding_num, 1, 1) noise = torch.randn_like(init_latents) - init_timesteps = reversed(self.timesteps).repeat_interleave(self.temporal_window_size, dim=0) - noisy_latents_first = self.scheduler.add_noise(init_latents, noise, init_timesteps[:padding_num]) + init_timesteps = reversed(self.timesteps).repeat_interleave( + self.temporal_window_size, dim=0 + ) + noisy_latents_first = self.scheduler.add_noise( + init_latents, noise, init_timesteps[:padding_num] + ) for i in range(self.temporal_adaptive_step - 1): l = i * self.temporal_window_size @@ -575,9 +616,13 @@ def fuse_reference(self, ref_image: Image.Image): 1, 1, self.temporal_window_size, 1, 1 ) noise = torch.randn_like(new_latents) - new_latents = self.scheduler.add_noise(new_latents, noise, self.timesteps[-1:]) + new_latents = self.scheduler.add_noise( + new_latents, noise, self.timesteps[-1:] + ) sample = torch.cat([sample, new_latents], dim=2) - self.trt_runner.prefill(sample=sample) # Use 'sample' (input name), not 'latents' (output name) + self.trt_runner.prefill( + sample=sample + ) # Use 'sample' (input name), not 'latents' (output name) self.reference_fused = True logger.info("Reference image fused successfully") @@ -683,8 +728,10 @@ def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: tgt_cond_tensor = self._fast_resize(images, 256, 256).clamp(0, 1) if self.first_frame: - mot_bbox_param, kps_ref, kps_frame1, kps_dri = self.pose_encoder.interpolate_kps_online( - self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + mot_bbox_param, kps_ref, kps_frame1, kps_dri = ( + self.pose_encoder.interpolate_kps_online( + self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + ) ) self.kps_ref = kps_ref self.kps_frame1 = kps_frame1 @@ -697,7 +744,7 @@ def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: with _profile_stage("keypoint_drawing"): keypoints = draw_keypoints(mot_bbox_param, device=device) boxes = get_boxes(kps_dri) - keypoints = rearrange(keypoints.unsqueeze(2), 'f c b h w -> b c f h w') + keypoints = rearrange(keypoints.unsqueeze(2), "f c b h w -> b c f h w") keypoints = keypoints.to(device=device, dtype=self.dtype) # Face cropping for motion input @@ -733,7 +780,9 @@ def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: # Prefill TensorRT with initial motion and pose states self.trt_runner.prefill( motion_hidden_states=init_motion_hidden_states, - pose_cond_fea=pose_fea[:, :, :temporal_window_size * (temporal_adaptive_step - 1)], + pose_cond_fea=pose_fea[ + :, :, : temporal_window_size * (temporal_adaptive_step - 1) + ], ) self.motion_bank = ref_motion @@ -750,9 +799,13 @@ def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: # Prepare noise with _profile_stage("noise_generation"): new_noise = torch.randn( - batch_size, 4, temporal_window_size, - self.height // 8, self.width // 8, - device=device, dtype=self.dtype + batch_size, + 4, + temporal_window_size, + self.height // 8, + self.width // 8, + device=device, + dtype=self.dtype, ) # TensorRT inference (includes motion_encoder, pose_guider, denoising_unet, vae_decode) @@ -801,15 +854,38 @@ def _process_frames_trt(self, images: torch.Tensor) -> torch.Tensor: # Concatenate new keyframe features ref_hidden_names = [ - "d00", "d01", "d10", "d11", "d20", "d21", "m", - "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" + "d00", + "d01", + "d10", + "d11", + "d20", + "d21", + "m", + "u10", + "u11", + "u12", + "u20", + "u21", + "u22", + "u30", + "u31", + "u32", ] for name in ref_hidden_names: - if name in self._reference_hidden_states and name in new_ref_hidden: + if ( + name in self._reference_hidden_states + and name in new_ref_hidden + ): self._reference_hidden_states[name] = torch.cat( - [self._reference_hidden_states[name], new_ref_hidden[name]], dim=1 + [ + self._reference_hidden_states[name], + new_ref_hidden[name], + ], + dim=1, + ) + self.trt_runner.prefill( + **{name: self._reference_hidden_states[name]} ) - self.trt_runner.prefill(**{name: self._reference_hidden_states[name]}) logger.debug("Added history keyframe (TensorRT)") self.num_khf += 1 @@ -866,8 +942,10 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: # Get keypoints if self.first_frame: - mot_bbox_param, kps_ref, kps_frame1, kps_dri = self.pose_encoder.interpolate_kps_online( - self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + mot_bbox_param, kps_ref, kps_frame1, kps_dri = ( + self.pose_encoder.interpolate_kps_online( + self.ref_cond_tensor, tgt_cond_tensor, num_interp=12 + 1 + ) ) self.kps_ref = kps_ref self.kps_frame1 = kps_frame1 @@ -879,7 +957,7 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: # Draw keypoints and get bounding boxes keypoints = draw_keypoints(mot_bbox_param, device=device) boxes = get_boxes(kps_dri) - keypoints = rearrange(keypoints.unsqueeze(2), 'f c b h w -> b c f h w') + keypoints = rearrange(keypoints.unsqueeze(2), "f c b h w -> b c f h w") keypoints = keypoints.to(device=device, dtype=self.pose_guider.dtype) # Process motion features @@ -926,7 +1004,9 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: self.pose_pile.append(pose_fea) # Prepare noisy latents for new frames - latents = self.ref_image_latents.unsqueeze(2).repeat(1, 1, temporal_window_size, 1, 1) + latents = self.ref_image_latents.unsqueeze(2).repeat( + 1, 1, temporal_window_size, 1, 1 + ) noise = torch.randn_like(latents) latents = self.scheduler.add_noise(latents, noise, self.timesteps[:1]) self.latents_pile.append(latents) @@ -946,9 +1026,11 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: # Denoise latents_model_input = torch.cat(list(self.latents_pile), dim=2) for j in range(jump): - timesteps = reversed(self.timesteps[j::jump]).repeat_interleave(temporal_window_size, dim=0) + timesteps = reversed(self.timesteps[j::jump]).repeat_interleave( + temporal_window_size, dim=0 + ) timesteps = torch.stack([timesteps] * batch_size) - timesteps = rearrange(timesteps, 'b f -> (b f)') + timesteps = rearrange(timesteps, "b f -> (b f)") noise_pred = self.denoising_unet( latents_model_input, @@ -962,18 +1044,28 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: )[0] clip_length = noise_pred.shape[2] - mid_noise_pred = rearrange(noise_pred, 'b c f h w -> (b f) c h w') - mid_latents = rearrange(latents_model_input, 'b c f h w -> (b f) c h w') + mid_noise_pred = rearrange(noise_pred, "b c f h w -> (b f) c h w") + mid_latents = rearrange(latents_model_input, "b c f h w -> (b f) c h w") latents_model_input, pred_original_sample = self.scheduler.step( - mid_noise_pred, timesteps, mid_latents, - generator=self.generator, return_dict=False + mid_noise_pred, + timesteps, + mid_latents, + generator=self.generator, + return_dict=False, + ) + latents_model_input = rearrange( + latents_model_input, "(b f) c h w -> b c f h w", f=clip_length + ) + pred_original_sample = rearrange( + pred_original_sample, "(b f) c h w -> b c f h w", f=clip_length + ) + latents_model_input = torch.cat( + [ + pred_original_sample[:, :, :temporal_window_size], + latents_model_input[:, :, temporal_window_size:], + ], + dim=2, ) - latents_model_input = rearrange(latents_model_input, '(b f) c h w -> b c f h w', f=clip_length) - pred_original_sample = rearrange(pred_original_sample, '(b f) c h w -> b c f h w', f=clip_length) - latents_model_input = torch.cat([ - pred_original_sample[:, :, :temporal_window_size], - latents_model_input[:, :, temporal_window_size:] - ], dim=2) latents_model_input = latents_model_input.to(dtype=self.dtype) # History keyframe mechanism @@ -981,7 +1073,9 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: self.reference_control_writer.clear() self.reference_unet( pred_original_sample[:, :, 0].to(self.reference_unet.dtype), - torch.zeros((batch_size,), dtype=self.dtype, device=self.reference_unet.device), + torch.zeros( + (batch_size,), dtype=self.dtype, device=self.reference_unet.device + ), encoder_hidden_states=self.encoder_hidden_states, return_dict=False, ) @@ -992,7 +1086,11 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: # Update latent piles for i in range(len(self.latents_pile)): self.latents_pile[i] = latents_model_input[ - :, :, i * temporal_adaptive_step:(i + 1) * temporal_adaptive_step, :, : + :, + :, + i * temporal_adaptive_step : (i + 1) * temporal_adaptive_step, + :, + :, ] # Pop oldest latents and decode diff --git a/src/scope/core/pipelines/personalive/tensorrt/__init__.py b/src/scope/core/pipelines/personalive/tensorrt/__init__.py index d1bae587..acb1b235 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/__init__.py +++ b/src/scope/core/pipelines/personalive/tensorrt/__init__.py @@ -30,6 +30,7 @@ # Check for zero-copy CUDA buffer support (polygraphy) try: from polygraphy.cuda import DeviceView # noqa: F401 + CUDA_BUFFERS_AVAILABLE = True except ImportError: pass @@ -38,6 +39,7 @@ try: import pycuda.driver # noqa: F401 import pycuda.autoinit # noqa: F401 + PYCUDA_AVAILABLE = True except ImportError: pass @@ -68,17 +70,23 @@ else: # Provide stub functions when TRT is not available def get_engine_path(*args, **kwargs): - raise RuntimeError("TensorRT not available. Install with: pip install daydream-scope[tensorrt]") + raise RuntimeError( + "TensorRT not available. Install with: pip install daydream-scope[tensorrt]" + ) def build_engine(*args, **kwargs): - raise RuntimeError("TensorRT not available. Install with: pip install daydream-scope[tensorrt]") + raise RuntimeError( + "TensorRT not available. Install with: pip install daydream-scope[tensorrt]" + ) def is_engine_available(*args, **kwargs): return False class TRTRunner: def __init__(self, *args, **kwargs): - raise RuntimeError("TensorRT not available. Install with: pip install daydream-scope[tensorrt]") + raise RuntimeError( + "TensorRT not available. Install with: pip install daydream-scope[tensorrt]" + ) __all__ = [ "build_engine", diff --git a/src/scope/core/pipelines/personalive/tensorrt/builder.py b/src/scope/core/pipelines/personalive/tensorrt/builder.py index a97204be..67dfbcbd 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/builder.py +++ b/src/scope/core/pipelines/personalive/tensorrt/builder.py @@ -85,7 +85,9 @@ def create_optimization_profile( Polygraphy Profile object. """ if not TRT_AVAILABLE: - raise RuntimeError("TensorRT is not available. Install with: pip install daydream-scope[tensorrt]") + raise RuntimeError( + "TensorRT is not available. Install with: pip install daydream-scope[tensorrt]" + ) # Derived dimensions tb = temporal_window_size * temporal_adaptive_step # temporal batch size (16) @@ -105,9 +107,20 @@ def create_optimization_profile( fixed_inputs = { "sample": (batch_size, lc, tb, lh, lw), "encoder_hidden_states": (batch_size, 1, emb), - "motion_hidden_states": (batch_size, temporal_window_size * (temporal_adaptive_step - 1), ml, mc), + "motion_hidden_states": ( + batch_size, + temporal_window_size * (temporal_adaptive_step - 1), + ml, + mc, + ), "motion": (batch_size, ic, temporal_window_size, mh, mw), - "pose_cond_fea": (batch_size, cd0, temporal_window_size * (temporal_adaptive_step - 1), lh, lw), + "pose_cond_fea": ( + batch_size, + cd0, + temporal_window_size * (temporal_adaptive_step - 1), + lh, + lw, + ), "pose": (batch_size, ic, temporal_window_size, height, width), "new_noise": (batch_size, lc, temporal_window_size, lh, lw), } @@ -177,7 +190,9 @@ def build_engine( FileNotFoundError: If ONNX model doesn't exist. """ if not TRT_AVAILABLE: - raise RuntimeError("TensorRT is not available. Install with: pip install daydream-scope[tensorrt]") + raise RuntimeError( + "TensorRT is not available. Install with: pip install daydream-scope[tensorrt]" + ) onnx_path = Path(onnx_path) engine_path = Path(engine_path) @@ -241,4 +256,3 @@ def is_engine_available(model_dir: Path, height: int = 512, width: int = 512) -> """ engine_path = get_engine_path(model_dir, height, width) return engine_path.exists() - diff --git a/src/scope/core/pipelines/personalive/tensorrt/convert.py b/src/scope/core/pipelines/personalive/tensorrt/convert.py index 390e75f5..a0e260a5 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/convert.py +++ b/src/scope/core/pipelines/personalive/tensorrt/convert.py @@ -35,7 +35,9 @@ def log_gpu_memory(msg: str = ""): if torch.cuda.is_available(): allocated = torch.cuda.memory_allocated() / 1024**3 reserved = torch.cuda.memory_reserved() / 1024**3 - logger.info(f"GPU Memory {msg}: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved") + logger.info( + f"GPU Memory {msg}: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved" + ) def parse_args(): @@ -103,6 +105,7 @@ def load_models(model_dir: Path, device: torch.device, dtype: torch.dtype): MotEncoder, PoseGuider, ) + # Use explicit reference UNet for TensorRT export from .unet_explicit import UNet3DConditionModelExplicit @@ -182,7 +185,9 @@ def cleanup(): # Load VAE logger.info("Loading VAE...") - vae = AutoencoderKL.from_pretrained(str(vae_path)).to(device=device, dtype=export_dtype) + vae = AutoencoderKL.from_pretrained(str(vae_path)).to( + device=device, dtype=export_dtype + ) vae.set_default_attn_processor() cleanup() @@ -295,7 +300,9 @@ def main(): logger.error(f"Model directory not found: {args.model_dir}") sys.exit(1) - personalive_weights = args.model_dir / "PersonaLive" / "pretrained_weights" / "personalive" + personalive_weights = ( + args.model_dir / "PersonaLive" / "pretrained_weights" / "personalive" + ) if not personalive_weights.exists(): logger.error(f"PersonaLive weights not found at: {personalive_weights}") logger.error("Please download PersonaLive models first using: download_models") @@ -323,7 +330,9 @@ def main(): logger.info(f"Converting PersonaLive models to TensorRT") logger.info(f" Model directory: {args.model_dir}") logger.info(f" Target resolution: {args.height}x{args.width}") - logger.info(f" ONNX export resolution: {onnx_height}x{onnx_width} (to save memory)") + logger.info( + f" ONNX export resolution: {onnx_height}x{onnx_width} (to save memory)" + ) logger.info(f" Precision: {'FP32' if args.fp32 else 'FP16'}") logger.info(f" Device: {device}") @@ -342,7 +351,7 @@ def main(): models=models, onnx_path=onnx_path, height=onnx_height, # 256 to save memory - width=onnx_width, # 256 to save memory + width=onnx_width, # 256 to save memory batch_size=args.batch_size, device=device, dtype=dtype, diff --git a/src/scope/core/pipelines/personalive/tensorrt/engine_model.py b/src/scope/core/pipelines/personalive/tensorrt/engine_model.py index 6068e0d1..7c87a467 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/engine_model.py +++ b/src/scope/core/pipelines/personalive/tensorrt/engine_model.py @@ -117,28 +117,46 @@ def get_safe_shape(engine, name): # Use max profile shape for dynamic dimensions profile = engine.get_tensor_profile_shape(name, 0) if profile: - logger.debug(f"Dynamic shape for {name}: {shape} -> Max: {profile[2]}") + logger.debug( + f"Dynamic shape for {name}: {shape} -> Max: {profile[2]}" + ) return profile[2] return shape # Get shapes and dtypes for all tensors - self.input_shapes = {name: get_safe_shape(self.engine, name) for name in self.input_names} - self.input_dtypes = {name: self.engine.get_tensor_dtype(name) for name in self.input_names} + self.input_shapes = { + name: get_safe_shape(self.engine, name) for name in self.input_names + } + self.input_dtypes = { + name: self.engine.get_tensor_dtype(name) for name in self.input_names + } self.input_nbytes = { - name: trt.volume(self.input_shapes[name]) * trt.nptype(self.input_dtypes[name])().itemsize + name: trt.volume(self.input_shapes[name]) + * trt.nptype(self.input_dtypes[name])().itemsize for name in self.input_names } - self.output_shapes = {name: get_safe_shape(self.engine, name) for name in self.output_names} - self.output_dtypes = {name: self.engine.get_tensor_dtype(name) for name in self.output_names} + self.output_shapes = { + name: get_safe_shape(self.engine, name) for name in self.output_names + } + self.output_dtypes = { + name: self.engine.get_tensor_dtype(name) for name in self.output_names + } self.output_nbytes = { - name: trt.volume(self.output_shapes[name]) * trt.nptype(self.output_dtypes[name])().itemsize + name: trt.volume(self.output_shapes[name]) + * trt.nptype(self.output_dtypes[name])().itemsize for name in self.output_names } # Allocate CUDA device memory for inputs and outputs - self.dinputs = {name: cuda.mem_alloc(self.input_nbytes[name]) for name in self.input_names} - self.doutputs = {name: cuda.mem_alloc(self.output_nbytes[name]) for name in self.output_names} + self.dinputs = { + name: cuda.mem_alloc(self.input_nbytes[name]) + for name in self.input_names + } + self.doutputs = { + name: cuda.mem_alloc(self.output_nbytes[name]) + for name in self.output_names + } # Create execution context self.context = self.engine.create_execution_context() @@ -156,12 +174,14 @@ def get_safe_shape(engine, name): self.houtputs = { name: cuda.pagelocked_empty( trt.volume(self.output_shapes[name]), - dtype=trt.nptype(self.output_dtypes[name]) + dtype=trt.nptype(self.output_dtypes[name]), ) for name in self.output_names } - logger.info(f"TensorRT engine loaded successfully ({len(self.input_names)} inputs, {len(self.output_names)} outputs)") + logger.info( + f"TensorRT engine loaded successfully ({len(self.input_names)} inputs, {len(self.output_names)} outputs)" + ) except Exception as e: self.ctx.pop() @@ -197,14 +217,18 @@ def __call__( if name not in self.input_names: continue - if isinstance(hinput, torch.Tensor) and hinput.is_cuda and hinput.device.index == self.device_int: + if ( + isinstance(hinput, torch.Tensor) + and hinput.is_cuda + and hinput.device.index == self.device_int + ): # GPU tensor on same device - device-to-device copy (fast!) hinput_con = hinput.contiguous() cuda.memcpy_dtod_async( self.dinputs[name], hinput_con.data_ptr(), self.input_nbytes[name], - self.stream + self.stream, ) else: # CPU tensor or numpy array - host-to-device copy @@ -228,20 +252,24 @@ def __call__( t = torch.zeros( trt.volume(self.output_shapes[name]), device=self._torch_device, - dtype=_numpy_to_torch_dtype(trt.nptype(self.output_dtypes[name])) + dtype=_numpy_to_torch_dtype( + trt.nptype(self.output_dtypes[name]) + ), ) cuda.memcpy_dtod_async( t.data_ptr(), self.doutputs[name], self.output_nbytes[name], - self.stream + self.stream, ) t = t.reshape(tuple(self.output_shapes[name])) result[name] = t else: # Device-to-host copy to numpy for name in output_names: - cuda.memcpy_dtoh_async(self.houtputs[name], self.doutputs[name], self.stream) + cuda.memcpy_dtoh_async( + self.houtputs[name], self.doutputs[name], self.stream + ) result[name] = self.houtputs[name].reshape(self.output_shapes[name]) # Synchronize stream @@ -287,10 +315,16 @@ def prefill(self, **inputs: torch.Tensor | np.ndarray) -> bool: real_nbytes = hinput.nbytes # Copy to device - if isinstance(hinput, torch.Tensor) and hinput.is_cuda and hinput.device.index == self.device_int: + if ( + isinstance(hinput, torch.Tensor) + and hinput.is_cuda + and hinput.device.index == self.device_int + ): # GPU tensor - device-to-device copy hinput_con = hinput.contiguous() - cuda.memcpy_dtod_async(dst_ptr, hinput_con.data_ptr(), real_nbytes, self.stream) + cuda.memcpy_dtod_async( + dst_ptr, hinput_con.data_ptr(), real_nbytes, self.stream + ) else: # CPU tensor/numpy - host-to-device copy if isinstance(hinput, torch.Tensor): @@ -333,7 +367,9 @@ def bind(self, output_to_input: dict[str, str]) -> bool: continue # Point input address to output buffer - self.context.set_tensor_address(input_name, int(self.doutputs[output_name])) + self.context.set_tensor_address( + input_name, int(self.doutputs[output_name]) + ) except Exception as e: logger.error(f"Bind failed: {e}") @@ -367,12 +403,12 @@ def get_output_shapes(self) -> dict[str, tuple]: def __del__(self): """Clean up CUDA resources.""" try: - if hasattr(self, 'ctx'): + if hasattr(self, "ctx"): self.ctx.push() # Free CUDA memory - for ptr in getattr(self, 'dinputs', {}).values(): + for ptr in getattr(self, "dinputs", {}).values(): ptr.free() - for ptr in getattr(self, 'doutputs', {}).values(): + for ptr in getattr(self, "doutputs", {}).values(): ptr.free() self.ctx.pop() except Exception: @@ -390,5 +426,3 @@ def __repr__(self) -> str: r += f" {name}: {dtype}{tuple(self.output_shapes[name])},\n" r += " ]\n)" return r - - diff --git a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py index 85533c1d..2d5bb0cf 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py +++ b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py @@ -13,6 +13,7 @@ try: from polygraphy.backend.trt import Profile + POLYGRAPHY_AVAILABLE = True except ImportError: POLYGRAPHY_AVAILABLE = False @@ -121,7 +122,9 @@ def forward( # Encode new motion features and concatenate new_motion_hidden_states = self.motion_encoder(motion) - motion_hidden_states = torch.cat([motion_hidden_states, new_motion_hidden_states], dim=1) + motion_hidden_states = torch.cat( + [motion_hidden_states, new_motion_hidden_states], dim=1 + ) # Prepare encoder hidden states for UNet [clip_embeds, motion_features] encoder_hidden_states_combined = [encoder_hidden_states, motion_hidden_states] @@ -132,13 +135,27 @@ def forward( self.timesteps, encoder_hidden_states_combined, pose_cond_fea, - d00, d01, d10, d11, d20, d21, m, - u10, u11, u12, u20, u21, u22, u30, u31, u32, + d00, + d01, + d10, + d11, + d20, + d21, + m, + u10, + u11, + u12, + u20, + u21, + u22, + u30, + u31, + u32, ) # Scheduler step - score = rearrange(score, 'b c f h w -> (b f) c h w') - sample_flat = rearrange(sample, 'b c f h w -> (b f) c h w') + score = rearrange(score, "b c f h w -> (b f) c h w") + sample_flat = rearrange(sample, "b c f h w -> (b f) c h w") latents_model_input, pred_original_sample = self.scheduler.step( score, self.timesteps, sample_flat, return_dict=False @@ -148,7 +165,9 @@ def forward( pred_original_sample = pred_original_sample.to(sample.dtype) # Reshape back to 5D - latents_model_input = rearrange(latents_model_input, '(b f) c h w -> b c f h w', f=16) + latents_model_input = rearrange( + latents_model_input, "(b f) c h w -> b c f h w", f=16 + ) # Decode first 4 frames (temporal_window_size) pred_video = self.decode_slice(self.vae, pred_original_sample[:4]) @@ -167,7 +186,14 @@ def forward( # First frame latent for potential keyframe update latent_first = pred_original_sample[:1] - return pred_video, latents, pose_cond_fea_out, motion_hidden_states_out, motion_out, latent_first + return ( + pred_video, + latents, + pose_cond_fea_out, + motion_hidden_states_out, + motion_out, + latent_first, + ) def get_sample_input(self, batchsize: int, height: int, width: int, dtype, device): """Generate sample inputs for ONNX export. @@ -182,11 +208,23 @@ def get_sample_input(self, batchsize: int, height: int, width: int, dtype, devic Returns: Dictionary of sample tensors for ONNX export. """ - tw, ts, tb = 4, 4, 16 # temporal_window_size, temporal_adaptive_steps, temporal_batch + tw, ts, tb = ( + 4, + 4, + 16, + ) # temporal_window_size, temporal_adaptive_steps, temporal_batch ml, mc, mh, mw = 32, 16, 224, 224 # motion latent dims b, h, w = batchsize, height, width lh, lw = height // 8, width // 8 # latent height/width - cd0, cd1, cd2, cm, cu1, cu2, cu3 = 320, 640, 1280, 1280, 1280, 640, 320 # unet channels + cd0, cd1, cd2, cm, cu1, cu2, cu3 = ( + 320, + 640, + 1280, + 1280, + 1280, + 640, + 320, + ) # unet channels emb = 768 # CLIP embedding dims lc, ic = 4, 3 # latent/image channels @@ -221,18 +259,41 @@ def get_sample_input(self, batchsize: int, height: int, width: int, dtype, devic def get_input_names(): """Get ordered input names for ONNX export.""" return [ - "sample", "encoder_hidden_states", "motion_hidden_states", - "motion", "pose_cond_fea", "pose", "new_noise", - "d00", "d01", "d10", "d11", "d20", "d21", "m", - "u10", "u11", "u12", "u20", "u21", "u22", "u30", "u31", "u32" + "sample", + "encoder_hidden_states", + "motion_hidden_states", + "motion", + "pose_cond_fea", + "pose", + "new_noise", + "d00", + "d01", + "d10", + "d11", + "d20", + "d21", + "m", + "u10", + "u11", + "u12", + "u20", + "u21", + "u22", + "u30", + "u31", + "u32", ] @staticmethod def get_output_names(): """Get ordered output names for ONNX export.""" return [ - "pred_video", "latents", "pose_cond_fea_out", - "motion_hidden_states_out", "motion_out", "latent_first" + "pred_video", + "latents", + "pose_cond_fea_out", + "motion_hidden_states_out", + "motion_out", + "latent_first", ] @staticmethod diff --git a/src/scope/core/pipelines/personalive/tensorrt/runner.py b/src/scope/core/pipelines/personalive/tensorrt/runner.py index 2362478c..1416de2b 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/runner.py +++ b/src/scope/core/pipelines/personalive/tensorrt/runner.py @@ -26,6 +26,7 @@ # Suppress polygraphy deprecation warnings about DeviceView.dtype # (warns about future API change: DataType.from_dtype(device_view.dtype).numpy()) from polygraphy.logger import G_LOGGER + G_LOGGER.severity = G_LOGGER.ERROR # Only show errors, not warnings TRT_AVAILABLE = True @@ -33,9 +34,12 @@ # Check for CUDA buffer support (polygraphy >= 0.47) try: from polygraphy.cuda import DeviceView + CUDA_BUFFERS_AVAILABLE = True except ImportError: - logger.warning("polygraphy.cuda.DeviceView not available - using numpy path (slower)") + logger.warning( + "polygraphy.cuda.DeviceView not available - using numpy path (slower)" + ) except ImportError: pass @@ -68,7 +72,9 @@ def __init__( FileNotFoundError: If engine file doesn't exist. """ if not TRT_AVAILABLE: - raise RuntimeError("TensorRT is not available. Install with: pip install daydream-scope[tensorrt]") + raise RuntimeError( + "TensorRT is not available. Install with: pip install daydream-scope[tensorrt]" + ) engine_path = Path(engine_path) if not engine_path.exists(): @@ -92,7 +98,9 @@ def __init__( if CUDA_BUFFERS_AVAILABLE: logger.info("TensorRT runner using CUDA device buffers (zero-copy)") else: - logger.warning("TensorRT runner using numpy path (slower - install polygraphy>=0.47 for zero-copy)") + logger.warning( + "TensorRT runner using numpy path (slower - install polygraphy>=0.47 for zero-copy)" + ) def _load_engine(self): """Load the TensorRT engine.""" @@ -261,7 +269,7 @@ def _call_cuda_buffers( # Convert outputs to PyTorch tensors (they come back as DeviceView or numpy) result = {} for name, output in outputs.items(): - if hasattr(output, 'numpy'): + if hasattr(output, "numpy"): # DeviceView - create tensor from same memory arr = output.numpy() result[name] = torch.from_numpy(arr).to(self.device) @@ -352,13 +360,19 @@ def get_input_shapes(self) -> dict[str, tuple]: """Get input tensor shapes.""" if self._runner is None: return {} - return {name: tuple(meta.shape) for name, meta in self._runner.get_input_metadata().items()} + return { + name: tuple(meta.shape) + for name, meta in self._runner.get_input_metadata().items() + } def get_output_shapes(self) -> dict[str, tuple]: """Get output tensor shapes.""" if self._runner is None: return {} - return {name: tuple(meta.shape) for name, meta in self._runner.get_output_metadata().items()} + return { + name: tuple(meta.shape) + for name, meta in self._runner.get_output_metadata().items() + } def __del__(self): """Clean up resources.""" diff --git a/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py b/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py index d3d0dd15..7cc80db2 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py +++ b/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py @@ -25,7 +25,11 @@ from safetensors.torch import load_file from ..modules.resnet import InflatedConv3d, InflatedGroupNorm -from ..modules.unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block +from ..modules.unet_3d_blocks import ( + UNetMidBlock3DCrossAttn, + get_down_block, + get_up_block, +) logger = logging.get_logger(__name__) @@ -41,6 +45,7 @@ class UNet3DConditionModelExplicit(ModelMixin, ConfigMixin): This model accepts reference hidden states (d00, d01, ..., u32) as explicit forward parameters instead of using attention processor hooks. """ + _supports_gradient_checkpointing = True @register_to_config @@ -169,7 +174,7 @@ def __init__( motion_module_type=motion_module_type, temporal_module_type=temporal_module_type, motion_module_kwargs=motion_module_kwargs, - temporal_module_kwargs=temporal_module_kwargs + temporal_module_kwargs=temporal_module_kwargs, ) self.down_blocks.append(down_block) @@ -404,7 +409,9 @@ def forward( if self.class_embedding is not None: if class_labels is None: - raise ValueError("class_labels should be provided when num_class_embeds > 0") + raise ValueError( + "class_labels should be provided when num_class_embeds > 0" + ) if self.config.class_embed_type == "timestep": class_labels = self.time_proj(class_labels) class_emb = self.class_embedding(class_labels).to(dtype=self.dtype) @@ -453,7 +460,9 @@ def forward( for down_block_res_sample, down_block_additional_residual in zip( down_block_res_samples, down_block_additional_residuals ): - down_block_res_sample = down_block_res_sample + down_block_additional_residual + down_block_res_sample = ( + down_block_res_sample + down_block_additional_residual + ) new_down_block_res_samples += (down_block_res_sample,) down_block_res_samples = new_down_block_res_samples @@ -482,8 +491,10 @@ def forward( is_final_block = i == len(self.up_blocks) - 1 - res_samples = down_block_res_samples[-len(upsample_block.resnets):] - down_block_res_samples = down_block_res_samples[:-len(upsample_block.resnets)] + res_samples = down_block_res_samples[-len(upsample_block.resnets) :] + down_block_res_samples = down_block_res_samples[ + : -len(upsample_block.resnets) + ] if not is_final_block and forward_upsample_size: upsample_size = down_block_res_samples[-1].shape[2:] @@ -583,8 +594,9 @@ def from_pretrained_2d( raise RuntimeError(f"Unknown format: {motion_module_path.suffix}") motion_state_dict = { - k.replace('motion_modules.', 'temporal_modules.'): v - for k, v in motion_state_dict.items() if "pos_encoder" not in k + k.replace("motion_modules.", "temporal_modules."): v + for k, v in motion_state_dict.items() + if "pos_encoder" not in k } if mm_zero_proj_out: @@ -598,4 +610,3 @@ def from_pretrained_2d( logger.debug(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}") return model - diff --git a/src/scope/server/app.py b/src/scope/server/app.py index 35a24e76..603f796a 100644 --- a/src/scope/server/app.py +++ b/src/scope/server/app.py @@ -336,13 +336,12 @@ async def set_personalive_reference( status_info = await pipeline_manager.get_status_info_async() if status_info.get("pipeline_id") != "personalive": raise HTTPException( - status_code=400, - detail="PersonaLive pipeline must be loaded first" + status_code=400, detail="PersonaLive pipeline must be loaded first" ) if status_info.get("status") != "loaded": raise HTTPException( status_code=400, - detail=f"Pipeline not ready. Current status: {status_info.get('status')}" + detail=f"Pipeline not ready. Current status: {status_info.get('status')}", ) # Get image from request body @@ -354,7 +353,9 @@ async def set_personalive_reference( try: image = Image.open(BytesIO(body)).convert("RGB") except Exception as e: - raise HTTPException(status_code=400, detail=f"Invalid image data: {e}") from e + raise HTTPException( + status_code=400, detail=f"Invalid image data: {e}" + ) from e # Get pipeline and fuse reference pipeline = pipeline_manager.get_pipeline() @@ -366,7 +367,7 @@ async def set_personalive_reference( return PersonaLiveReferenceResponse( success=True, - message="Reference image set successfully. Ready to process driving video." + message="Reference image set successfully. Ready to process driving video.", ) except HTTPException: diff --git a/src/scope/server/download_models.py b/src/scope/server/download_models.py index 15df3dc8..d1ecc9b6 100644 --- a/src/scope/server/download_models.py +++ b/src/scope/server/download_models.py @@ -250,7 +250,9 @@ def download_personalive_pipeline() -> None: ) print(f"[OK] Downloaded sd-image-variations-diffusers to: {sd_variations_dst}") else: - print(f"[SKIP] sd-image-variations-diffusers already exists at: {sd_variations_dst}") + print( + f"[SKIP] sd-image-variations-diffusers already exists at: {sd_variations_dst}" + ) # 2. Download sd-vae-ft-mse (improved VAE) # This provides: config.json, diffusion_pytorch_model.safetensors diff --git a/src/scope/server/models_config.py b/src/scope/server/models_config.py index cfc2b133..7dc4acc3 100644 --- a/src/scope/server/models_config.py +++ b/src/scope/server/models_config.py @@ -126,14 +126,46 @@ def get_required_model_files(pipeline_id: str | None = None) -> list[Path]: # personalive pipeline if pipeline_id == "personalive": return [ - models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "denoising_unet.pth", - models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "reference_unet.pth", - models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "motion_encoder.pth", - models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "motion_extractor.pth", - models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "pose_guider.pth", - models_dir / "PersonaLive" / "pretrained_weights" / "personalive" / "temporal_module.pth", - models_dir / "PersonaLive" / "pretrained_weights" / "sd-image-variations-diffusers" / "model_index.json", - models_dir / "PersonaLive" / "pretrained_weights" / "sd-vae-ft-mse" / "config.json", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "personalive" + / "denoising_unet.pth", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "personalive" + / "reference_unet.pth", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "personalive" + / "motion_encoder.pth", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "personalive" + / "motion_extractor.pth", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "personalive" + / "pose_guider.pth", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "personalive" + / "temporal_module.pth", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "sd-image-variations-diffusers" + / "model_index.json", + models_dir + / "PersonaLive" + / "pretrained_weights" + / "sd-vae-ft-mse" + / "config.json", ] # Default: nothing is required diff --git a/src/scope/server/schema.py b/src/scope/server/schema.py index 7052358d..f17a41b4 100644 --- a/src/scope/server/schema.py +++ b/src/scope/server/schema.py @@ -391,7 +391,9 @@ class PersonaLiveLoadParams(PipelineLoadParams): class PersonaLiveReferenceResponse(BaseModel): """Response after setting PersonaLive reference image.""" - success: bool = Field(..., description="Whether the reference image was set successfully") + success: bool = Field( + ..., description="Whether the reference image was set successfully" + ) message: str = Field(..., description="Status message") From 85ce8a393d78184e1a6ef4784304d4dac8506bce Mon Sep 17 00:00:00 2001 From: BuffMcBigHuge Date: Mon, 22 Dec 2025 11:45:44 -0500 Subject: [PATCH 7/7] Partial lint. Signed-off-by: BuffMcBigHuge --- .../core/pipelines/personalive/__init__.py | 2 + .../pipelines/personalive/blocks/__init__.py | 2 + .../personalive/components/face_detector.py | 13 +++--- .../core/pipelines/personalive/docs/README.md | 2 + .../core/pipelines/personalive/pipeline.py | 17 ++++--- .../personalive/tensorrt/__init__.py | 2 +- .../pipelines/personalive/tensorrt/builder.py | 3 +- .../pipelines/personalive/tensorrt/convert.py | 2 +- .../personalive/tensorrt/engine_model.py | 6 +-- .../pipelines/personalive/tensorrt/export.py | 5 +-- .../personalive/tensorrt/framed_model.py | 2 +- .../pipelines/personalive/tensorrt/runner.py | 2 - .../personalive/tensorrt/unet_explicit.py | 44 ++++++++++--------- src/scope/server/app.py | 1 + src/scope/server/patch_xformers.py | 2 - 15 files changed, 55 insertions(+), 50 deletions(-) diff --git a/src/scope/core/pipelines/personalive/__init__.py b/src/scope/core/pipelines/personalive/__init__.py index f3995ef9..48ad939b 100644 --- a/src/scope/core/pipelines/personalive/__init__.py +++ b/src/scope/core/pipelines/personalive/__init__.py @@ -3,3 +3,5 @@ from .pipeline import PersonaLivePipeline __all__ = ["PersonaLivePipeline"] + + diff --git a/src/scope/core/pipelines/personalive/blocks/__init__.py b/src/scope/core/pipelines/personalive/blocks/__init__.py index fdb7b942..5409311e 100644 --- a/src/scope/core/pipelines/personalive/blocks/__init__.py +++ b/src/scope/core/pipelines/personalive/blocks/__init__.py @@ -1 +1,3 @@ """PersonaLive modular blocks for pipeline composition.""" + + diff --git a/src/scope/core/pipelines/personalive/components/face_detector.py b/src/scope/core/pipelines/personalive/components/face_detector.py index bea0e23c..c88cb11f 100644 --- a/src/scope/core/pipelines/personalive/components/face_detector.py +++ b/src/scope/core/pipelines/personalive/components/face_detector.py @@ -1,7 +1,6 @@ """Face detection component using MediaPipe for PersonaLive pipeline.""" import logging -from typing import Tuple import numpy as np import torch @@ -60,7 +59,7 @@ def detect_face_landmarks(self, image: np.ndarray) -> list | None: def get_face_bounding_box( self, image: np.ndarray, padding_ratio: float = 0.2 - ) -> Tuple[int, int, int, int] | None: + ) -> tuple[int, int, int, int] | None: """Get bounding box for the detected face. Args: @@ -103,7 +102,7 @@ def get_face_bounding_box( def crop_face( self, image: np.ndarray, - target_size: Tuple[int, int] = (512, 512), + target_size: tuple[int, int] = (512, 512), padding_ratio: float = 0.2, ) -> np.ndarray | None: """Crop and resize face from image. @@ -131,7 +130,7 @@ def crop_face( return np.array(face_pil) - def crop_face_from_pil(self, pil_image, target_size: Tuple[int, int] = (512, 512)): + def crop_face_from_pil(self, pil_image, target_size: tuple[int, int] = (512, 512)): """Crop face from a PIL image. Args: @@ -154,8 +153,8 @@ def crop_face_from_pil(self, pil_image, target_size: Tuple[int, int] = (512, 512 def crop_face_tensor( self, image_tensor: torch.Tensor, - boxes: Tuple[int, int, int, int], - target_size: Tuple[int, int] = (224, 224), + boxes: tuple[int, int, int, int], + target_size: tuple[int, int] = (224, 224), ) -> torch.Tensor: """Crop face region from a tensor using precomputed box. @@ -196,3 +195,5 @@ def close(self): def __del__(self): self.close() + + diff --git a/src/scope/core/pipelines/personalive/docs/README.md b/src/scope/core/pipelines/personalive/docs/README.md index affdfad0..9ce63dbf 100644 --- a/src/scope/core/pipelines/personalive/docs/README.md +++ b/src/scope/core/pipelines/personalive/docs/README.md @@ -123,3 +123,5 @@ PersonaLive uses a multi-component architecture: Based on the PersonaLive implementation by GVCLab. + + diff --git a/src/scope/core/pipelines/personalive/pipeline.py b/src/scope/core/pipelines/personalive/pipeline.py index c8904e8f..590e6c63 100644 --- a/src/scope/core/pipelines/personalive/pipeline.py +++ b/src/scope/core/pipelines/personalive/pipeline.py @@ -26,7 +26,7 @@ from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection from ..interface import Pipeline, Requirements -from ..process import preprocess_chunk, postprocess_chunk +from ..process import preprocess_chunk from ..schema import PersonaLiveConfig from ..utils import load_model_config from .components import FaceDetector @@ -42,7 +42,6 @@ get_boxes, ) - # Profiling utilities _profiling_enabled = False _profiling_timings: dict[str, list[float]] = defaultdict(list) @@ -87,7 +86,7 @@ def _profile_stage(name: str): # TensorRT support (optional) try: - from .tensorrt import TRT_AVAILABLE, TRTRunner, get_engine_path, PYCUDA_AVAILABLE + from .tensorrt import PYCUDA_AVAILABLE, TRT_AVAILABLE, TRTRunner, get_engine_path # Import EngineModel if pycuda is available (best performance) if PYCUDA_AVAILABLE: @@ -604,9 +603,9 @@ def fuse_reference(self, ref_image: Image.Image): ) for i in range(self.temporal_adaptive_step - 1): - l = i * self.temporal_window_size + left_idx = i * self.temporal_window_size r = (i + 1) * self.temporal_window_size - self.latents_pile.append(noisy_latents_first[:, :, l:r]) + self.latents_pile.append(noisy_latents_first[:, :, left_idx:r]) # For TensorRT, also prefill initial latents (using correct input name 'sample') if self.use_tensorrt: @@ -977,9 +976,9 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: ref_motion, dri_motion[:, :1], num=12 + 1 )[:, :-1] for i in range(temporal_adaptive_step - 1): - l = i * temporal_window_size + left_idx = i * temporal_window_size r = (i + 1) * temporal_window_size - self.motion_pile.append(init_motion_hidden_states[:, l:r]) + self.motion_pile.append(init_motion_hidden_states[:, left_idx:r]) self.motion_pile.append(dri_motion) self.motion_bank = ref_motion @@ -996,9 +995,9 @@ def _process_frames_pytorch(self, images: torch.Tensor) -> torch.Tensor: pose_fea = self.pose_guider(keypoints) if self.first_frame: for i in range(temporal_adaptive_step): - l = i * temporal_window_size + left_idx = i * temporal_window_size r = (i + 1) * temporal_window_size - self.pose_pile.append(pose_fea[:, :, l:r]) + self.pose_pile.append(pose_fea[:, :, left_idx:r]) self.first_frame = False else: self.pose_pile.append(pose_fea) diff --git a/src/scope/core/pipelines/personalive/tensorrt/__init__.py b/src/scope/core/pipelines/personalive/tensorrt/__init__.py index acb1b235..2a76ebf2 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/__init__.py +++ b/src/scope/core/pipelines/personalive/tensorrt/__init__.py @@ -37,8 +37,8 @@ # Check for pycuda (best performance) try: - import pycuda.driver # noqa: F401 import pycuda.autoinit # noqa: F401 + import pycuda.driver # noqa: F401 PYCUDA_AVAILABLE = True except ImportError: diff --git a/src/scope/core/pipelines/personalive/tensorrt/builder.py b/src/scope/core/pipelines/personalive/tensorrt/builder.py index 67dfbcbd..1274805f 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/builder.py +++ b/src/scope/core/pipelines/personalive/tensorrt/builder.py @@ -5,7 +5,6 @@ """ import logging -import os from pathlib import Path from typing import TYPE_CHECKING @@ -256,3 +255,5 @@ def is_engine_available(model_dir: Path, height: int = 512, width: int = 512) -> """ engine_path = get_engine_path(model_dir, height, width) return engine_path.exists() + + diff --git a/src/scope/core/pipelines/personalive/tensorrt/convert.py b/src/scope/core/pipelines/personalive/tensorrt/convert.py index a0e260a5..7de1a140 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/convert.py +++ b/src/scope/core/pipelines/personalive/tensorrt/convert.py @@ -327,7 +327,7 @@ def main(): # TRT engine is built at target resolution with dynamic shapes onnx_height, onnx_width = 256, 256 - logger.info(f"Converting PersonaLive models to TensorRT") + logger.info("Converting PersonaLive models to TensorRT") logger.info(f" Model directory: {args.model_dir}") logger.info(f" Target resolution: {args.height}x{args.width}") logger.info( diff --git a/src/scope/core/pipelines/personalive/tensorrt/engine_model.py b/src/scope/core/pipelines/personalive/tensorrt/engine_model.py index 7c87a467..d2693eea 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/engine_model.py +++ b/src/scope/core/pipelines/personalive/tensorrt/engine_model.py @@ -20,9 +20,9 @@ # Check for TensorRT and pycuda availability PYCUDA_TRT_AVAILABLE = False try: - import tensorrt as trt - import pycuda.driver as cuda import pycuda.autoinit # noqa: F401 - Required for pycuda initialization + import pycuda.driver as cuda + import tensorrt as trt PYCUDA_TRT_AVAILABLE = True TRT_LOGGER = trt.Logger(trt.Logger.WARNING) @@ -185,7 +185,7 @@ def get_safe_shape(engine, name): except Exception as e: self.ctx.pop() - raise RuntimeError(f"Failed to initialize TensorRT engine: {e}") + raise RuntimeError(f"Failed to initialize TensorRT engine: {e}") from e self.ctx.pop() diff --git a/src/scope/core/pipelines/personalive/tensorrt/export.py b/src/scope/core/pipelines/personalive/tensorrt/export.py index 8e5917d6..8e7dfa03 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/export.py +++ b/src/scope/core/pipelines/personalive/tensorrt/export.py @@ -8,7 +8,6 @@ import gc import logging -import os from contextlib import contextmanager from pathlib import Path @@ -108,8 +107,8 @@ def optimize_onnx( """ try: import onnx - except ImportError: - raise RuntimeError("onnx package required. Install with: pip install onnx") + except ImportError as e: + raise RuntimeError("onnx package required. Install with: pip install onnx") from e onnx_path = Path(onnx_path) onnx_opt_path = Path(onnx_opt_path) diff --git a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py index 2d5bb0cf..a9ef436f 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/framed_model.py +++ b/src/scope/core/pipelines/personalive/tensorrt/framed_model.py @@ -8,8 +8,8 @@ """ import torch -from torch import nn from einops import rearrange +from torch import nn try: from polygraphy.backend.trt import Profile diff --git a/src/scope/core/pipelines/personalive/tensorrt/runner.py b/src/scope/core/pipelines/personalive/tensorrt/runner.py index 1416de2b..1efee6d9 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/runner.py +++ b/src/scope/core/pipelines/personalive/tensorrt/runner.py @@ -8,7 +8,6 @@ import logging from pathlib import Path -from typing import Any import numpy as np import torch @@ -19,7 +18,6 @@ TRT_AVAILABLE = False CUDA_BUFFERS_AVAILABLE = False try: - import tensorrt as trt from polygraphy.backend.common import BytesFromPath from polygraphy.backend.trt import EngineFromBytes, TrtRunner diff --git a/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py b/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py index 7cc80db2..d9aa521b 100644 --- a/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py +++ b/src/scope/core/pipelines/personalive/tensorrt/unet_explicit.py @@ -8,11 +8,9 @@ PersonaLive/src/models/unet_3d_explicit_reference.py """ -from collections import OrderedDict from dataclasses import dataclass from os import PathLike from pathlib import Path -from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn as nn @@ -51,27 +49,27 @@ class UNet3DConditionModelExplicit(ModelMixin, ConfigMixin): @register_to_config def __init__( self, - sample_size: Optional[int] = None, + sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, - down_block_types: Tuple[str] = ( + down_block_types: tuple[str] = ( "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D", ), mid_block_type: str = "UNetMidBlock3DCrossAttn", - up_block_types: Tuple[str] = ( + up_block_types: tuple[str] = ( "UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", ), - only_cross_attention: Union[bool, Tuple[bool]] = False, - block_out_channels: Tuple[int] = (320, 640, 1280, 1280), + only_cross_attention: bool | tuple[bool] = False, + block_out_channels: tuple[int] = (320, 640, 1280, 1280), layers_per_block: int = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, @@ -79,11 +77,11 @@ def __init__( norm_num_groups: int = 32, norm_eps: float = 1e-5, cross_attention_dim: int = 1280, - attention_head_dim: Union[int, Tuple[int]] = 8, + attention_head_dim: int | tuple[int] = 8, dual_cross_attention: bool = False, use_linear_projection: bool = False, - class_embed_type: Optional[str] = None, - num_class_embeds: Optional[int] = None, + class_embed_type: str | None = None, + num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = "default", use_inflated_groupnorm=False, @@ -95,12 +93,16 @@ def __init__( motion_module_decoder_only=False, motion_module_type=None, temporal_module_type=None, - motion_module_kwargs={}, - temporal_module_kwargs={}, + motion_module_kwargs=None, + temporal_module_kwargs=None, unet_use_cross_frame_attention=None, unet_use_temporal_attention=None, ): super().__init__() + if motion_module_kwargs is None: + motion_module_kwargs = {} + if temporal_module_kwargs is None: + temporal_module_kwargs = {} self.sample_size = sample_size time_embed_dim = block_out_channels[0] * 4 @@ -283,14 +285,14 @@ def __init__( ) @property - def attn_processors(self) -> Dict[str, AttentionProcessor]: + def attn_processors(self) -> dict[str, AttentionProcessor]: """Returns attention processors.""" processors = {} def fn_recursive_add_processors( name: str, module: torch.nn.Module, - processors: Dict[str, AttentionProcessor], + processors: dict[str, AttentionProcessor], ): if hasattr(module, "set_processor"): processors[f"{name}.processor"] = module.processor @@ -308,7 +310,7 @@ def fn_recursive_add_processors( return processors def set_attn_processor( - self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]] + self, processor: AttentionProcessor | dict[str, AttentionProcessor] ): """Sets the attention processor.""" count = len(self.attn_processors.keys()) @@ -335,7 +337,7 @@ def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor): def forward( self, sample: torch.FloatTensor, - timestep: Union[torch.Tensor, float, int], + timestep: torch.Tensor | float | int, encoder_hidden_states: torch.Tensor, pose_cond_fea: torch.Tensor, # Explicit reference hidden states for TensorRT @@ -355,13 +357,13 @@ def forward( u30: torch.Tensor, u31: torch.Tensor, u32: torch.Tensor, - class_labels: Optional[torch.Tensor] = None, - attention_mask: Optional[torch.Tensor] = None, - down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None, - mid_block_additional_residual: Optional[torch.Tensor] = None, + class_labels: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + down_block_additional_residuals: tuple[torch.Tensor] | None = None, + mid_block_additional_residual: torch.Tensor | None = None, return_dict: bool = True, skip_mm: bool = False, - ) -> Union[UNet3DConditionOutput, Tuple]: + ) -> UNet3DConditionOutput | tuple: """Forward pass with explicit reference hidden states. Args: diff --git a/src/scope/server/app.py b/src/scope/server/app.py index 603f796a..21448372 100644 --- a/src/scope/server/app.py +++ b/src/scope/server/app.py @@ -329,6 +329,7 @@ async def set_personalive_reference( starting video streaming. """ from io import BytesIO + from PIL import Image try: diff --git a/src/scope/server/patch_xformers.py b/src/scope/server/patch_xformers.py index 89281ca6..983cde33 100644 --- a/src/scope/server/patch_xformers.py +++ b/src/scope/server/patch_xformers.py @@ -8,8 +8,6 @@ def find_xformers_flash_file() -> Path | None: """Find the xformers flash.py file in the installed package.""" try: - import xformers - # Get the package path spec = importlib.util.find_spec("xformers") if spec is None or spec.origin is None: