refactor: convert prepare-test-weights.sh to python

CONTRIBUTING.md

@@ -48,7 +48,7 @@ rye sync --features test,conversion
 Then, download and convert all the necessary weights. Be aware that this will use around 50 GB of disk space:
 
 ```bash
-./scripts/prepare-test-weights.sh
+rye run python scripts/prepare_test_weights.py
 ```
 
 Finally, run the tests:

pyproject.toml

@@ -41,6 +41,8 @@ conversion = [
     "diffusers>=0.24.0",
     "transformers>=4.35.2",
     "segment-anything-py>=1.0",
+    "requests>=2.26.0",
+    "tqdm>=4.62.3",
 ]
 
 [build-system]

scripts/conversion/convert_diffusers_autoencoder_kl.py

@@ -18,8 +18,11 @@ class Args(argparse.Namespace):
 
 def setup_converter(args: Args) -> ModelConverter:
     target = LatentDiffusionAutoencoder()
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
     source: nn.Module = AutoencoderKL.from_pretrained(  # type: ignore
-        pretrained_model_name_or_path=args.source_path, subfolder=args.subfolder
+        pretrained_model_name_or_path=args.source_path,
+        subfolder=args.subfolder,
+        low_cpu_mem_usage=False,
     )  # type: ignore
     x = torch.randn(1, 3, 512, 512)
     converter = ModelConverter(source_model=source, target_model=target, skip_output_check=True, verbose=args.verbose)

scripts/conversion/convert_diffusers_controlnet.py

@@ -22,7 +22,11 @@ class Args(argparse.Namespace):
 
 @torch.no_grad()
 def convert(args: Args) -> dict[str, torch.Tensor]:
-    controlnet_src: nn.Module = ControlNetModel.from_pretrained(pretrained_model_name_or_path=args.source_path)  # type: ignore
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
+    controlnet_src: nn.Module = ControlNetModel.from_pretrained(  # type: ignore
+        pretrained_model_name_or_path=args.source_path,
+        low_cpu_mem_usage=False,
+    )
     unet = SD1UNet(in_channels=4)
     adapter = SD1ControlnetAdapter(unet, name="mycn").inject()
     controlnet = unet.Controlnet

scripts/conversion/convert_diffusers_lora.py

@@ -40,7 +40,11 @@ class Args(argparse.Namespace):
 @torch.no_grad()
 def process(args: Args) -> None:
     diffusers_state_dict = cast(dict[str, Tensor], torch.load(args.source_path, map_location="cpu"))  # type: ignore
-    diffusers_sd = DiffusionPipeline.from_pretrained(pretrained_model_name_or_path=args.base_model)  # type: ignore
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
+    diffusers_sd = DiffusionPipeline.from_pretrained(  # type: ignore
+        pretrained_model_name_or_path=args.base_model,
+        low_cpu_mem_usage=False,
+    )
     diffusers_model = cast(fl.Module, diffusers_sd.unet)  # type: ignore
 
     refiners_model = SD1UNet(in_channels=4)

scripts/conversion/convert_diffusers_t2i_adapter.py

@@ -48,7 +48,11 @@
 
     sdxl = "xl" in args.source_path
     target = ConditionEncoderXL() if sdxl else ConditionEncoder()
-    source: nn.Module = T2IAdapter.from_pretrained(pretrained_model_name_or_path=args.source_path)  # type: ignore
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
+    source: nn.Module = T2IAdapter.from_pretrained(  # type: ignore
+        pretrained_model_name_or_path=args.source_path,
+        low_cpu_mem_usage=False,
+    )
     assert isinstance(source, nn.Module), "Source model is not a nn.Module"
 
     x = torch.randn(1, 3, 1024, 1024) if sdxl else torch.randn(1, 3, 512, 512)

scripts/conversion/convert_diffusers_unet.py

@@ -17,8 +17,11 @@ class Args(argparse.Namespace):
 
 
 def setup_converter(args: Args) -> ModelConverter:
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
     source: nn.Module = UNet2DConditionModel.from_pretrained(  # type: ignore
-        pretrained_model_name_or_path=args.source_path, subfolder="unet"
+        pretrained_model_name_or_path=args.source_path,
+        subfolder="unet",
+        low_cpu_mem_usage=False,
     )
     source_in_channels: int = source.config.in_channels  # type: ignore
     source_clip_embedding_dim: int = source.config.cross_attention_dim  # type: ignore

scripts/conversion/convert_transformers_clip_image_model.py

@@ -21,8 +21,11 @@ class Args(argparse.Namespace):
 
 
 def setup_converter(args: Args) -> ModelConverter:
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
     source: nn.Module = CLIPVisionModelWithProjection.from_pretrained(  # type: ignore
-        pretrained_model_name_or_path=args.source_path, subfolder=args.subfolder
+        pretrained_model_name_or_path=args.source_path,
+        subfolder=args.subfolder,
+        low_cpu_mem_usage=False,
     )
     assert isinstance(source, nn.Module), "Source model is not a nn.Module"
     architecture: str = source.config.architectures[0]  # type: ignore

scripts/conversion/convert_transformers_clip_text_model.py

@@ -22,8 +22,11 @@ class Args(argparse.Namespace):
 
 
 def setup_converter(args: Args) -> ModelConverter:
+    # low_cpu_mem_usage=False stops some annoying console messages us to `pip install accelerate`
     source: nn.Module = CLIPTextModelWithProjection.from_pretrained(  # type: ignore
-        pretrained_model_name_or_path=args.source_path, subfolder=args.subfolder
+        pretrained_model_name_or_path=args.source_path,
+        subfolder=args.subfolder,
+        low_cpu_mem_usage=False,
     )
     assert isinstance(source, nn.Module), "Source model is not a nn.Module"
     architecture: str = source.config.architectures[0]  # type: ignore