patrickvonplaten commited on
Commit
066cf1b
1 Parent(s): 8cfdc75
Files changed (2) hide show
  1. convert_flax_to_pt.py +30 -0
  2. run_control_inpaint.py +55 -0
convert_flax_to_pt.py ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ import argparse
3
+ from huggingface_hub import HfApi
4
+
5
+
6
+ def main(api, model_id):
7
+ info = api.list_repo_refs(model_id)
8
+ branches = set([b.name for b in info.branches]) - set(["main"])
9
+
10
+ return list(branches)
11
+
12
+
13
+ if __name__ == "__main__":
14
+ DESCRIPTION = """
15
+ Simple utility to get all branches from a repo
16
+ """
17
+ parser = argparse.ArgumentParser(description=DESCRIPTION)
18
+ parser.add_argument(
19
+ "--model_id",
20
+ type=str,
21
+ help="The name of the model on the hub to retrieve the branches from. E.g. `gpt2` or `facebook/wav2vec2-base-960h`",
22
+ )
23
+
24
+ args = parser.parse_args()
25
+ model_id = args.model_id
26
+ api = HfApi()
27
+ branches = main(api, model_id)
28
+
29
+ if len(branches) > 0:
30
+ print(f"{model_id}: {branches}")
run_control_inpaint.py ADDED
@@ -0,0 +1,55 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ # !pip install transformers accelerate
3
+ from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler
4
+ from diffusers.utils import load_image
5
+ import numpy as np
6
+ import torch
7
+
8
+ init_image = load_image(
9
+ "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"
10
+ )
11
+ init_image = init_image.resize((512, 512))
12
+
13
+ generator = torch.Generator(device="cpu").manual_seed(33)
14
+
15
+ mask_image = load_image(
16
+ "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"
17
+ )
18
+ mask_image = mask_image.resize((512, 512))
19
+
20
+
21
+ def make_inpaint_condition(image, image_mask):
22
+ image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
23
+ image_mask = np.array(image_mask.convert("L")).astype(np.float32) / 255.0
24
+
25
+ assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
26
+ image[image_mask > 0.5] = -1.0 # set as masked pixel
27
+ image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
28
+ image = torch.from_numpy(image)
29
+ return image
30
+
31
+
32
+ control_image = make_inpaint_condition(init_image, mask_image)
33
+
34
+ controlnet = ControlNetModel.from_pretrained(
35
+ "lllyasviel/control_v11p_sd15_inpaint", torch_dtype=torch.float16
36
+ )
37
+ pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
38
+ "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
39
+ )
40
+
41
+ # speed up diffusion process with faster scheduler and memory optimization
42
+ pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
43
+
44
+ pipe.enable_model_cpu_offload()
45
+
46
+ # generate image
47
+ image = pipe(
48
+ "a beautiful man",
49
+ num_inference_steps=20,
50
+ generator=generator,
51
+ eta=1.0,
52
+ image=init_image,
53
+ mask_image=mask_image,
54
+ control_image=control_image,
55
+ ).images[0]