Refactor import paths and update requirements to include timm

models/infinity.py

@@ -18,15 +18,15 @@ from PIL import Image
 import numpy as np
 from torch.nn.attention.flex_attention import flex_attention
 
-import infinity.utils.dist as dist
-from infinity.utils.dist import for_visualize
-from infinity.models.basic import flash_attn_func, flash_fused_op_installed, AdaLNBeforeHead, CrossAttnBlock, SelfAttnBlock, CrossAttention, FastRMSNorm, precompute_rope2d_freqs_grid
-from infinity.utils import misc
-from infinity.models.flex_attn import FlexAttn
-from infinity.utils.dynamic_resolution import dynamic_resolution_h_w, h_div_w_templates
+import utils.dist as dist
+from utils.dist import for_visualize
+from models.basic import flash_attn_func, flash_fused_op_installed, AdaLNBeforeHead, CrossAttnBlock, SelfAttnBlock, CrossAttention, FastRMSNorm, precompute_rope2d_freqs_grid
+from utils import misc
+from models.flex_attn import FlexAttn
+from utils.dynamic_resolution import dynamic_resolution_h_w, h_div_w_templates
 
 try:
-    from infinity.models.fused_op import fused_ada_layer_norm, fused_ada_rms_norm
+    from models.fused_op import fused_ada_layer_norm, fused_ada_rms_norm
 except:
     fused_ada_layer_norm, fused_ada_rms_norm = None, None
 

requirements.txt

@@ -5,4 +5,5 @@ huggingface-hub
 transformers
 argparse
 spaces
-torchvision
+torchvision
+timm