From d64e5e439d68b4236e9823dad35ce778b136bfc9 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 24 Mar 2026 16:00:24 +0530 Subject: [PATCH] Fix unguarded `torchvision` import in Cosmos (#13321) update --- .../cosmos/pipeline_cosmos2_5_predict.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py index cdea71a5ab..5817112058 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py @@ -16,22 +16,29 @@ from typing import Callable import numpy as np import torch -import torchvision -import torchvision.transforms -import torchvision.transforms.functional from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration from ...callbacks import MultiPipelineCallbacks, PipelineCallback from ...image_processor import PipelineImageInput from ...models import AutoencoderKLWan, CosmosTransformer3DModel from ...schedulers import UniPCMultistepScheduler -from ...utils import is_cosmos_guardrail_available, is_torch_xla_available, logging, replace_example_docstring +from ...utils import ( + is_cosmos_guardrail_available, + is_torch_xla_available, + is_torchvision_available, + logging, + replace_example_docstring, +) from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import CosmosPipelineOutput +if is_torchvision_available(): + import torchvision.transforms.functional + + if is_cosmos_guardrail_available(): from cosmos_guardrail import CosmosSafetyChecker else: