Skip to content

Commit

Permalink
Remove redundant importing
Browse files Browse the repository at this point in the history
Signed-off-by: mingyuanm <mingyuanm@nvidia.com>
  • Loading branch information
Victor49152 committed Oct 16, 2024
1 parent fbd6987 commit aa9df2a
Show file tree
Hide file tree
Showing 4 changed files with 3 additions and 9 deletions.
3 changes: 0 additions & 3 deletions nemo/collections/diffusion/encoders/conditioner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion nemo/collections/diffusion/models/flux/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 2 additions & 4 deletions nemo/collections/diffusion/models/flux/pipeline.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -28,7 +28,7 @@
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):
Expand Down Expand Up @@ -247,8 +247,6 @@ def __call__(
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
Expand Down
1 change: 0 additions & 1 deletion nemo/collections/diffusion/utils/mcore_parallel_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit aa9df2a

Please sign in to comment.