File size: 4,957 Bytes
c70f6b6
 
 
 
8c6f026
0cea845
c70f6b6
 
 
 
 
 
 
0d53d7d
c70f6b6
 
9102b04
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1d15b14
9102b04
 
c70f6b6
 
 
 
 
 
 
 
10af042
 
c70f6b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9102b04
c70f6b6
 
 
 
9102b04
 
 
 
 
 
 
 
 
 
 
ce1c2ea
c70f6b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8e62d24
 
 
c70f6b6
 
 
 
 
 
 
 
 
 
9102b04
 
 
 
 
 
 
 
 
 
 
 
 
c70f6b6
 
9102b04
 
 
c70f6b6
 
 
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
import gradio as gr
import numpy as np
import spaces
import torch
import random
from peft import PeftModel

from diffusers import FluxControlPipeline, FluxTransformer2DModel
from image_gen_aux import DepthPreprocessor

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

pipe = FluxControlPipeline.from_pretrained("black-forest-labs/FLUX.1-Depth-dev", torch_dtype=torch.bfloat16).to("cuda")
processor = DepthPreprocessor.from_pretrained("LiheYoung/depth-anything-large-hf")

def load_lora(lora_path):
    if not lora_path.strip():
        return "Please provide a valid LoRA path"
    try:
        pipe.load_lora_weights(lora_path)
        return f"Successfully loaded LoRA weights from {lora_path}"
    except Exception as e:
        return f"Error loading LoRA weights: {str(e)}"

def unload_lora():
    try:
        pipe.unload_lora_weights()
        return "Successfully unloaded LoRA weights"
    except Exception as e:
        return f"Error unloading LoRA weights: {str(e)}"

@spaces.GPU
def infer(control_image, prompt, seed=42, randomize_seed=False, width=1024, height=1024, 
          guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    control_image = processor(control_image)[0].convert("RGB")
    image = pipe(
        prompt=prompt,
        control_image=control_image,
        height=height,
        width=width,
        num_inference_steps=num_inference_steps,
        guidance_scale=guidance_scale,
        generator=torch.Generator().manual_seed(seed),
    ).images[0]
    return image, seed

css="""
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

with gr.Blocks(css=css) as demo:
    
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 Depth [dev] with LoRA Support
12B param rectified flow transformer structural conditioning tuned, guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
        """)

        # LoRA controls
        with gr.Row():
            lora_path = gr.Textbox(
                label="HuggingFace LoRA Path",
                placeholder="e.g., Borcherding/FLUX.1-dev-LoRA-AutumnSpringTrees"
            )
            load_lora_btn = gr.Button("Load LoRA")
            unload_lora_btn = gr.Button("Unload LoRA")
        
        lora_status = gr.Textbox(label="LoRA Status", interactive=False)

        control_image = gr.Image(label="Upload the image for control", type="pil")
        with gr.Row():
            prompt = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)
        
        with gr.Accordion("Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            with gr.Row():
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=30,
                    step=0.5,
                    value=10,
                )
  
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

    # Event handlers
    load_lora_btn.click(
        fn=load_lora,
        inputs=[lora_path],
        outputs=[lora_status]
    )
    
    unload_lora_btn.click(
        fn=unload_lora,
        inputs=[],
        outputs=[lora_status]
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[control_image, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
        outputs=[result, seed]
    )

demo.launch()