File size: 2,359 Bytes
89b3db2
754b60e
 
1368e65
 
754b60e
 
 
 
 
1368e65
754b60e
 
 
 
1368e65
 
 
 
 
 
 
 
 
 
 
92e2f62
14d5805
 
 
 
 
 
92e2f62
14d5805
 
 
92e2f62
14d5805
9c0158e
d1be92b
1368e65
 
 
14d5805
 
1368e65
 
14d5805
1368e65
 
 
 
 
 
 
 
 
 
 
 
ba6d92c
14d5805
1368e65
ba6d92c
1368e65
 
 
 
 
 
14d5805
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import torch
from diffusers.utils import load_image
from diffusers import FluxControlNetPipeline, FluxControlNetModel, FluxMultiControlNetModel
import gradio as gr
import spaces

base_model = 'black-forest-labs/FLUX.1-dev'
controlnet_model_union = 'InstantX/FLUX.1-dev-Controlnet-Union'

controlnet_union = FluxControlNetModel.from_pretrained(controlnet_model_union, torch_dtype=torch.bfloat16)
controlnet = FluxMultiControlNetModel([controlnet_union])  # we always recommend loading via FluxMultiControlNetModel

pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
pipe.to("cuda")

control_modes = [
    "canny",
    "tile",
    "depth",
    "blur",
    "pose",
    "gray",
    "lq"
]

@spaces.GPU
def generate_image(prompt, control_image_depth, control_mode_depth_index, use_depth, control_image_canny, control_mode_canny_index):
    control_images = []
    control_modes = []
    conditioning_scales = []

    if use_depth:
        control_images.append(control_image_depth)
        control_modes.append(control_mode_depth_index)
        conditioning_scales.append(0.2)

    control_images.append(control_image_canny)
    control_modes.append(control_mode_canny_index)
    conditioning_scales.append(0.4)

    width, height = control_image_canny.shape[:2]

    image = pipe(
        prompt,
        control_image=control_images,
        control_mode=control_modes,
        width=width,
        height=height,
        controlnet_conditioning_scale=conditioning_scales,
        num_inference_steps=24,
        guidance_scale=3.5,
        generator=torch.manual_seed(42),
    ).images[0]

    return image

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Text(label="Prompt"),
        gr.Image(label="Control Image (Depth)"),
        gr.Dropdown(choices=control_modes, value=control_modes.index("depth"), label="Control Mode (Depth)"),
        gr.Checkbox(label="Use Depth Control Image", value=True),
        gr.Image(label="Control Image (Canny)"),
        gr.Dropdown(choices=control_modes, value=control_modes.index("canny"), label="Control Mode (Canny)")
    ],
    outputs=gr.Image(label="Generated Image"),
    title="FluxControlNet Image Generation",
    description="Generate an image using FluxControlNet with depth and canny control images.",
)

iface.launch(share=True)