Skip to content

Commit d64e5e4

Browse files
committed
Fix unguarded torchvision import in Cosmos (#13321)
update
1 parent e8ecef1 commit d64e5e4

File tree

1 file changed

+11
-4
lines changed

1 file changed

+11
-4
lines changed

src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,22 +16,29 @@
1616

1717
import numpy as np
1818
import torch
19-
import torchvision
20-
import torchvision.transforms
21-
import torchvision.transforms.functional
2219
from transformers import AutoTokenizer, Qwen2_5_VLForConditionalGeneration
2320

2421
from ...callbacks import MultiPipelineCallbacks, PipelineCallback
2522
from ...image_processor import PipelineImageInput
2623
from ...models import AutoencoderKLWan, CosmosTransformer3DModel
2724
from ...schedulers import UniPCMultistepScheduler
28-
from ...utils import is_cosmos_guardrail_available, is_torch_xla_available, logging, replace_example_docstring
25+
from ...utils import (
26+
is_cosmos_guardrail_available,
27+
is_torch_xla_available,
28+
is_torchvision_available,
29+
logging,
30+
replace_example_docstring,
31+
)
2932
from ...utils.torch_utils import randn_tensor
3033
from ...video_processor import VideoProcessor
3134
from ..pipeline_utils import DiffusionPipeline
3235
from .pipeline_output import CosmosPipelineOutput
3336

3437

38+
if is_torchvision_available():
39+
import torchvision.transforms.functional
40+
41+
3542
if is_cosmos_guardrail_available():
3643
from cosmos_guardrail import CosmosSafetyChecker
3744
else:

0 commit comments

Comments
 (0)