Alibrown commited on
Commit
103a9de
Β·
verified Β·
1 Parent(s): 080d285

Create _app.py

Browse files
Files changed (1) hide show
  1. _app.py +160 -0
_app.py ADDED
@@ -0,0 +1,160 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright Volkan Sah! Do not steal code if its for free! Respect creators!
2
+ # You can use it for free a star or follow will be greate!
3
+ import gradio as gr
4
+ import numpy as np
5
+ import random
6
+ import os
7
+ import spaces
8
+ import time
9
+ from diffusers import AutoPipelineForText2Image
10
+ import torch
11
+ from huggingface_hub import login
12
+
13
+ login(os.environ.get("HF_TOKEN"))
14
+ device = "cuda" if torch.cuda.is_available() else "cpu"
15
+ torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
16
+
17
+ pipe = AutoPipelineForText2Image.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
18
+ pipe.load_lora_weights('enhanceaiteam/Flux-uncensored', weight_name='lora.safetensors')
19
+ pipe = pipe.to(device)
20
+
21
+ MAX_SEED = np.iinfo(np.int32).max
22
+ MAX_IMAGE_SIZE = 2048
23
+
24
+ def dev_pipeline_test(prompt, width, height, steps):
25
+ """Simulates pipeline tests without actual image generation"""
26
+ time.sleep(0.1) # Simulate processing
27
+ estimated_memory = (width * height * 3 * 4) / (1024 * 1024) # MB
28
+ return {
29
+ 'success': True,
30
+ 'memory_required': f"{estimated_memory:.2f}MB",
31
+ 'compute_units': steps * (width * height) / 1024**2,
32
+ 'would_generate': True
33
+ }
34
+
35
+ @spaces.GPU
36
+ def infer(
37
+ prompt,
38
+ seed,
39
+ randomize_seed,
40
+ width,
41
+ height,
42
+ guidance_scale,
43
+ num_inference_steps,
44
+ dev_mode,
45
+ progress=gr.Progress(track_tqdm=True),
46
+ ):
47
+ if dev_mode:
48
+ result = dev_pipeline_test(prompt, width, height, num_inference_steps)
49
+ # Create a black test image with debug info
50
+ debug_image = np.zeros((height, width, 3), dtype=np.uint8)
51
+ return debug_image, seed, f"DEV MODE: {result}"
52
+
53
+ if randomize_seed:
54
+ seed = random.randint(0, MAX_SEED)
55
+ generator = torch.Generator().manual_seed(seed)
56
+ image = pipe(
57
+ prompt=prompt,
58
+ guidance_scale=guidance_scale,
59
+ num_inference_steps=num_inference_steps,
60
+ width=width,
61
+ height=height,
62
+ generator=generator,
63
+ ).images[0]
64
+ return image, seed, "Production generation completed"
65
+
66
+ examples = [
67
+ "Tiger in a jungle, cold color palette, muted colors, detailed, 8k",
68
+ "An astronaut riding a pink horse",
69
+ "A delicious ceviche cheesecake slice",
70
+ ]
71
+
72
+ css = """
73
+ #col-container {
74
+ margin: 0 auto;
75
+ max-width: 640px;
76
+ }
77
+ """
78
+
79
+ with gr.Blocks(css=css) as demo:
80
+ with gr.Column(elem_id="col-container"):
81
+ gr.Markdown("""# [FLUX.1-dev](https://blackforestlabs.ai/)
82
+ Generate any type of image with Flux-Dev (Lora: Flux-uncensored). Note: This script works well, but please use min. ZeroGPU
83
+ """)
84
+
85
+ with gr.Row():
86
+ prompt = gr.Text(
87
+ label="Prompt",
88
+ show_label=False,
89
+ max_lines=1,
90
+ placeholder="Enter your prompt",
91
+ container=False,
92
+ )
93
+ run_button = gr.Button("Run", scale=0, variant="primary")
94
+
95
+ result = gr.Image(label="Result", show_label=False)
96
+ status_text = gr.Text(label="Status", show_label=True)
97
+
98
+ with gr.Accordion("Advanced Settings", open=False):
99
+ dev_mode = gr.Checkbox(label="Developer Mode (No actual generation)", value=False)
100
+ seed = gr.Slider(
101
+ label="Seed",
102
+ minimum=0,
103
+ maximum=MAX_SEED,
104
+ step=1,
105
+ value=0,
106
+ )
107
+ randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
108
+
109
+ with gr.Row():
110
+ width = gr.Slider(
111
+ label="Width",
112
+ minimum=256,
113
+ maximum=MAX_IMAGE_SIZE,
114
+ step=32,
115
+ value=1024,
116
+ )
117
+ height = gr.Slider(
118
+ label="Height",
119
+ minimum=256,
120
+ maximum=MAX_IMAGE_SIZE,
121
+ step=32,
122
+ value=1024,
123
+ )
124
+
125
+ with gr.Row():
126
+ guidance_scale = gr.Slider(
127
+ label="Guidance scale",
128
+ minimum=0.0,
129
+ maximum=10.0,
130
+ step=0.1,
131
+ value=3.5,
132
+ )
133
+ num_inference_steps = gr.Slider(
134
+ label="Number of inference steps",
135
+ minimum=1,
136
+ maximum=50,
137
+ step=1,
138
+ value=2,
139
+ )
140
+
141
+ gr.Examples(examples=examples, inputs=[prompt])
142
+
143
+ gr.on(
144
+ triggers=[run_button.click, prompt.submit],
145
+ fn=infer,
146
+ inputs=[
147
+ prompt,
148
+ seed,
149
+ randomize_seed,
150
+ width,
151
+ height,
152
+ guidance_scale,
153
+ num_inference_steps,
154
+ dev_mode,
155
+ ],
156
+ outputs=[result, seed, status_text],
157
+ )
158
+
159
+ if __name__ == "__main__":
160
+ demo.launch()