From ee982c5d735a25b7fdb64e5cd987e052597de1e8 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Wed, 11 Dec 2024 22:44:09 +0100
Subject: [PATCH] Remove timm as a dependency

---
 PixArt/model/pixart.py | 20 +++++++++++++-------
 requirements.txt       |  1 -
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/PixArt/model/pixart.py b/PixArt/model/pixart.py
index f200177..a5ca872 100644
--- a/PixArt/model/pixart.py
+++ b/PixArt/model/pixart.py
@@ -8,17 +8,23 @@
 # GLIDE: https://github.com/openai/glide-text2im
 # MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
 # --------------------------------------------------------
-import math
 import torch
 import torch.nn as nn
-import os
 import numpy as np
-from timm.models.layers import DropPath
-from timm.models.vision_transformer import PatchEmbed, Mlp
-
 
 from .utils import auto_grad_checkpoint, to_2tuple
-from .blocks import t2i_modulate, CaptionEmbedder, AttentionKVCompress, MultiHeadCrossAttention, T2IFinalLayer, TimestepEmbedder, LabelEmbedder, FinalLayer
+from .blocks import (
+    t2i_modulate,
+    CaptionEmbedder,
+    AttentionKVCompress,
+    MultiHeadCrossAttention,
+    T2IFinalLayer,
+    TimestepEmbedder,
+    LabelEmbedder,
+    FinalLayer,
+    Mlp
+)
+from comfy.ldm.modules.diffusionmodules.mmdit import PatchEmbed
 
 
 class PixArtBlock(nn.Module):
@@ -37,7 +43,7 @@ class PixArtBlock(nn.Module):
         # to be compatible with lower version pytorch
         approx_gelu = lambda: nn.GELU(approximate="tanh")
         self.mlp = Mlp(in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu, drop=0)
-        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.drop_path = nn.Identity() #DropPath(drop_path) if drop_path > 0. else nn.Identity()
         self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5)
         self.sampling = sampling
         self.sr_ratio = sr_ratio
diff --git a/requirements.txt b/requirements.txt
index 1e54893..017b5b8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,3 @@
-timm>=0.6.13
 sentencepiece>=0.1.97
 transformers>=4.34.1
 accelerate>=0.23.0