Remove redundant importing

Signed-off-by: mingyuanm <[email protected]>

nemo/collections/diffusion/encoders/conditioner.py

@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
-
 import torch
 import torch.nn as nn
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer

nemo/collections/diffusion/models/flux/model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Callable
 
 import torch
 from megatron.core.transformer.transformer_config import TransformerConfig

nemo/collections/diffusion/models/flux/pipeline.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import os
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import List, Optional, Union
 
 import numpy as np
 import torch
@@ -24,11 +24,11 @@
 from tqdm import tqdm
 
 from nemo.collections.diffusion.encoders.conditioner import FrozenCLIPEmbedder, FrozenT5Embedder
 from nemo.collections.diffusion.models.flux.model import Flux, FluxParams
 from nemo.collections.diffusion.sampler.flow_matching.flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
 from nemo.collections.diffusion.utils.flux_ckpt_converter import flux_transformer_converter
 from nemo.collections.diffusion.utils.flux_pipeline_utils import FluxModelParams
-from nemo.collections.diffusion.vae.autoencoder import AutoEncoder, AutoEncoderParams
+from nemo.collections.diffusion.vae.autoencoder import AutoEncoder
 
 
 class FluxInferencePipeline(nn.Module):
@@ -247,8 +247,6 @@
         offload: bool = True,
     ):
         assert device == 'cuda', 'Transformer blocks in Mcore must run on cuda devices'
-        height = height
-        width = width
 
         if prompt is not None and isinstance(prompt, str):
             batch_size = 1

nemo/collections/diffusion/utils/mcore_parallel_utils.py

@@ -20,7 +20,6 @@
 
 import megatron.core.parallel_state as ps
 import torch
-from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
 
 
 class Utils: