File size: 4,936 Bytes
103a9de
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
157
158
159
160
# Copyright Volkan Sah! Do not steal code if its for free! Respect creators!
# You can use it for free a star or follow will be greate!
import gradio as gr
import numpy as np
import random
import os
import spaces
import time
from diffusers import AutoPipelineForText2Image
import torch
from huggingface_hub import login

login(os.environ.get("HF_TOKEN"))
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

pipe = AutoPipelineForText2Image.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
pipe.load_lora_weights('enhanceaiteam/Flux-uncensored', weight_name='lora.safetensors')
pipe = pipe.to(device)

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

def dev_pipeline_test(prompt, width, height, steps):
    """Simulates pipeline tests without actual image generation"""
    time.sleep(0.1)  # Simulate processing
    estimated_memory = (width * height * 3 * 4) / (1024 * 1024)  # MB
    return {
        'success': True,
        'memory_required': f"{estimated_memory:.2f}MB",
        'compute_units': steps * (width * height) / 1024**2,
        'would_generate': True
    }

@spaces.GPU
def infer(
    prompt,
    seed,
    randomize_seed,
    width,
    height,
    guidance_scale,
    num_inference_steps,
    dev_mode,
    progress=gr.Progress(track_tqdm=True),
):
    if dev_mode:
        result = dev_pipeline_test(prompt, width, height, num_inference_steps)
        # Create a black test image with debug info
        debug_image = np.zeros((height, width, 3), dtype=np.uint8)
        return debug_image, seed, f"DEV MODE: {result}"
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt=prompt,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        generator=generator,
    ).images[0]
    return image, seed, "Production generation completed"

examples = [
    "Tiger in a jungle, cold color palette, muted colors, detailed, 8k",
    "An astronaut riding a pink horse",
    "A delicious ceviche cheesecake slice",
]

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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""# [FLUX.1-dev](https://blackforestlabs.ai/)
Generate any type of image with Flux-Dev (Lora: Flux-uncensored). Note: This script works well, but please use min. ZeroGPU
""")
        
        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, variant="primary")
        
        result = gr.Image(label="Result", show_label=False)
        status_text = gr.Text(label="Status", show_label=True)
        
        with gr.Accordion("Advanced Settings", open=False):
            dev_mode = gr.Checkbox(label="Developer Mode (No actual generation)", value=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=0.0,
                    maximum=10.0,
                    step=0.1,
                    value=3.5,
                )
                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=2,
                )
        
        gr.Examples(examples=examples, inputs=[prompt])
        
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            seed,
            randomize_seed,
            width,
            height,
            guidance_scale,
            num_inference_steps,
            dev_mode,
        ],
        outputs=[result, seed, status_text],
    )

if __name__ == "__main__":
    demo.launch()