Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
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 |
-
|
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(
|
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
|
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)
|