Deadmon commited on
Commit
013bfbb
·
verified ·
1 Parent(s): 397288e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -12
app.py CHANGED
@@ -14,12 +14,7 @@ from src.flux.sampling import denoise_controlnet, get_noise, get_schedule, prepa
14
  from src.flux.util import load_ae, load_clip, load_t5, load_flow_model, load_controlnet, load_safetensors
15
 
16
  # Download and load the ControlNet model
17
- model_url = "https://huggingface.co/XLabs-AI/flux-controlnet-canny-v3/resolve/main/flux-canny-controlnet-v3.safetensors?download=true"
18
- model_path = "./flux-canny-controlnet-v3.safetensors"
19
- if not os.path.exists(model_path):
20
- response = requests.get(model_url)
21
- with open(model_path, 'wb') as f:
22
- f.write(response.content)
23
 
24
  # Source: https://github.com/XLabs-AI/x-flux.git
25
  name = "flux-dev"
@@ -35,10 +30,7 @@ def load_models():
35
  clip = load_clip(device)
36
  model = load_flow_model(name, device=device)
37
  ae = load_ae(name, device=device)
38
- controlnet = load_controlnet(name, device).to(device).to(torch.bfloat16)
39
-
40
- checkpoint = load_safetensors(model_path)
41
- controlnet.load_state_dict(checkpoint, strict=False)
42
 
43
  load_models()
44
 
@@ -98,8 +90,7 @@ def generate_image(prompt, control_image, control_mode, num_steps=50, guidance=4
98
  x = get_noise(1, height, width, device=torch_device, dtype=torch.bfloat16, seed=seed)
99
  inp_cond = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
100
 
101
- controlnet.control_mode = control_modes.index(control_mode) # Set the control mode directly in the controlnet model
102
- x = denoise_controlnet(model, **inp_cond, controlnet=controlnet, timesteps=timesteps, guidance=guidance, controlnet_cond=controlnet_cond)
103
 
104
  x = unpack(x.float(), height, width)
105
  x = ae.decode(x)
 
14
  from src.flux.util import load_ae, load_clip, load_t5, load_flow_model, load_controlnet, load_safetensors
15
 
16
  # Download and load the ControlNet model
17
+ controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 
 
 
 
 
18
 
19
  # Source: https://github.com/XLabs-AI/x-flux.git
20
  name = "flux-dev"
 
30
  clip = load_clip(device)
31
  model = load_flow_model(name, device=device)
32
  ae = load_ae(name, device=device)
33
+ controlnet = load_controlnet(controlnet_model, device).to(device).to(torch.bfloat16)
 
 
 
34
 
35
  load_models()
36
 
 
90
  x = get_noise(1, height, width, device=torch_device, dtype=torch.bfloat16, seed=seed)
91
  inp_cond = prepare(t5=t5, clip=clip, img=x, prompt=prompt)
92
 
93
+ x = denoise_controlnet(model, **inp_cond, controlnet=controlnet, timesteps=timesteps, guidance=guidance, controlnet_cond=controlnet_cond, control_mode=control_modes.index(control_mode))
 
94
 
95
  x = unpack(x.float(), height, width)
96
  x = ae.decode(x)