prithivMLmods commited on
Commit
328da4c
·
verified ·
1 Parent(s): bd634f6

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +50 -0
main.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from diffusers.utils import load_image, check_min_version
3
+ from controlnet_flux import FluxControlNetModel
4
+ from transformer_flux import FluxTransformer2DModel
5
+ from pipeline_flux_controlnet_inpaint import FluxControlNetInpaintingPipeline
6
+
7
+ check_min_version("0.30.2")
8
+
9
+ # Set image path , mask path and prompt
10
+ image_path='https://huggingface.co/alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha/resolve/main/images/bucket.png',
11
+ mask_path='https://huggingface.co/alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha/resolve/main/images/bucket_mask.jpeg',
12
+ prompt='a person wearing a white shoe, carrying a white bucket with text "FLUX" on it'
13
+
14
+ # Build pipeline
15
+ controlnet = FluxControlNetModel.from_pretrained("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha", torch_dtype=torch.bfloat16)
16
+ transformer = FluxTransformer2DModel.from_pretrained(
17
+ "black-forest-labs/FLUX.1-dev", subfolder='transformer', torch_dytpe=torch.bfloat16
18
+ )
19
+ pipe = FluxControlNetInpaintingPipeline.from_pretrained(
20
+ "black-forest-labs/FLUX.1-dev",
21
+ controlnet=controlnet,
22
+ transformer=transformer,
23
+ torch_dtype=torch.bfloat16
24
+ ).to("cuda")
25
+ pipe.transformer.to(torch.bfloat16)
26
+ pipe.controlnet.to(torch.bfloat16)
27
+
28
+ # Load image and mask
29
+ size = (768, 768)
30
+ image = load_image(image_path).convert("RGB").resize(size)
31
+ mask = load_image(mask_path).convert("RGB").resize(size)
32
+ generator = torch.Generator(device="cuda").manual_seed(24)
33
+
34
+ # Inpaint
35
+ result = pipe(
36
+ prompt=prompt,
37
+ height=size[1],
38
+ width=size[0],
39
+ control_image=image,
40
+ control_mask=mask,
41
+ num_inference_steps=28,
42
+ generator=generator,
43
+ controlnet_conditioning_scale=0.9,
44
+ guidance_scale=3.5,
45
+ negative_prompt="",
46
+ true_guidance_scale=3.5
47
+ ).images[0]
48
+
49
+ result.save('flux_inpaint.png')
50
+ print("Successfully inpaint image")