chengzeyi commited on
Commit
07cfdaf
·
verified ·
1 Parent(s): ea37dbc

Upload folder using huggingface_hub

Browse files
.pytest_cache/.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ # Created by pytest automatically.
2
+ *
.pytest_cache/CACHEDIR.TAG ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ Signature: 8a477f597d28d172789f06886806bc55
2
+ # This file is a cache directory tag created by pytest.
3
+ # For information about cache directory tags, see:
4
+ # https://bford.info/cachedir/spec.html
.pytest_cache/README.md ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ # pytest cache directory #
2
+
3
+ This directory contains data from the pytest's cache plugin,
4
+ which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
5
+
6
+ **Do not** commit this to version control.
7
+
8
+ See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.
.pytest_cache/v/cache/stepwise ADDED
@@ -0,0 +1 @@
 
 
1
+ []
README.md CHANGED
@@ -1,12 +1,14 @@
1
  ---
2
  title: FLUX Kontext Dev Ultra Fast
3
- emoji: 🚀
4
- colorFrom: blue
5
- colorTo: indigo
6
  sdk: gradio
7
- sdk_version: 5.34.2
8
  app_file: app.py
9
  pinned: false
 
 
10
  ---
11
 
12
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
1
  ---
2
  title: FLUX Kontext Dev Ultra Fast
3
+ emoji: 🖼
4
+ colorFrom: purple
5
+ colorTo: red
6
  sdk: gradio
7
+ sdk_version: 5.25.2
8
  app_file: app.py
9
  pinned: false
10
+ license: apache-2.0
11
+ short_description: 'Ultra Fast FLUX Kontext Dev for Image Editing'
12
  ---
13
 
14
+ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
app.py ADDED
@@ -0,0 +1,292 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import requests
3
+ import time
4
+ import threading
5
+ import uuid
6
+ import base64
7
+ from pathlib import Path
8
+ from dotenv import load_dotenv
9
+ import gradio as gr
10
+ import random
11
+ import torch
12
+ from PIL import Image, ImageDraw, ImageFont
13
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
14
+
15
+ load_dotenv()
16
+ API_KEY = os.getenv("WAVESPEED_API_KEY")
17
+ if not API_KEY:
18
+ raise ValueError("WAVESPEED_API_KEY is not set in environment variables")
19
+
20
+ MODEL_URL = "TostAI/nsfw-text-detection-large"
21
+ CLASS_NAMES = {0: "✅ SAFE", 1: "⚠️ QUESTIONABLE", 2: "🚫 UNSAFE"}
22
+
23
+ try:
24
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
25
+ model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
26
+ except Exception as e:
27
+ raise RuntimeError(f"Failed to load safety model: {str(e)}")
28
+
29
+
30
+ class SessionManager:
31
+ _instances = {}
32
+ _lock = threading.Lock()
33
+
34
+ @classmethod
35
+ def get_session(cls, session_id):
36
+ with cls._lock:
37
+ if session_id not in cls._instances:
38
+ cls._instances[session_id] = {
39
+ 'count': 0,
40
+ 'history': [],
41
+ 'last_active': time.time()
42
+ }
43
+ return cls._instances[session_id]
44
+
45
+ @classmethod
46
+ def cleanup_sessions(cls):
47
+ with cls._lock:
48
+ now = time.time()
49
+ expired = [
50
+ k for k, v in cls._instances.items()
51
+ if now - v['last_active'] > 3600
52
+ ]
53
+ for k in expired:
54
+ del cls._instances[k]
55
+
56
+
57
+ class RateLimiter:
58
+
59
+ def __init__(self):
60
+ self.clients = {}
61
+ self.lock = threading.Lock()
62
+
63
+ def check(self, client_id):
64
+ with self.lock:
65
+ now = time.time()
66
+ if client_id not in self.clients:
67
+ self.clients[client_id] = {'count': 1, 'reset': now + 3600}
68
+ return True
69
+ if now > self.clients[client_id]['reset']:
70
+ self.clients[client_id] = {'count': 1, 'reset': now + 3600}
71
+ return True
72
+ if self.clients[client_id]['count'] >= 20:
73
+ return False
74
+ self.clients[client_id]['count'] += 1
75
+ return True
76
+
77
+
78
+ session_manager = SessionManager()
79
+ rate_limiter = RateLimiter()
80
+
81
+
82
+ def create_error_image(message):
83
+ img = Image.new("RGB", (512, 512), "#ffdddd")
84
+ try:
85
+ font = ImageFont.truetype("arial.ttf", 24)
86
+ except:
87
+ font = ImageFont.load_default()
88
+ draw = ImageDraw.Draw(img)
89
+ text = f"Error: {message[:60]}..." if len(message) > 60 else message
90
+ draw.text((50, 200), text, fill="#ff0000", font=font)
91
+ return img
92
+
93
+
94
+ @torch.no_grad()
95
+ def classify_prompt(prompt):
96
+ inputs = tokenizer(prompt,
97
+ return_tensors="pt",
98
+ truncation=True,
99
+ max_length=512)
100
+ outputs = model(**inputs)
101
+ return torch.argmax(outputs.logits).item()
102
+
103
+
104
+ def image_to_base64(file_path):
105
+ with open(file_path, "rb") as f:
106
+ return base64.b64encode(f.read()).decode()
107
+
108
+
109
+ def generate_image(image_file,
110
+ prompt,
111
+ seed,
112
+ session_id,
113
+ enable_safety_checker=True):
114
+ try:
115
+ if enable_safety_checker:
116
+ safety_level = classify_prompt(prompt)
117
+ if safety_level != 0:
118
+ error_img = create_error_image(CLASS_NAMES[safety_level])
119
+ yield f"❌ Blocked: {CLASS_NAMES[safety_level]}", error_img, ""
120
+ return
121
+
122
+ if not rate_limiter.check(session_id):
123
+ error_img = create_error_image(
124
+ "Hourly limit exceeded (20 requests)")
125
+ yield "❌ Too many requests, please try again later", error_img, ""
126
+ return
127
+
128
+ session = session_manager.get_session(session_id)
129
+ session['last_active'] = time.time()
130
+ session['count'] += 1
131
+
132
+ error_messages = []
133
+ if not image_file:
134
+ error_messages.append("Please upload an image file")
135
+ elif not Path(image_file).exists():
136
+ error_messages.append("File does not exist")
137
+ if not prompt.strip():
138
+ error_messages.append("Prompt cannot be empty")
139
+ if error_messages:
140
+ error_img = create_error_image(" | ".join(error_messages))
141
+ yield "❌ Input validation failed", error_img, ""
142
+ return
143
+
144
+ try:
145
+ base64_image = image_to_base64(image_file)
146
+ except Exception as e:
147
+ error_img = create_error_image(f"File processing failed: {str(e)}")
148
+ yield "❌ File processing failed", error_img, ""
149
+ return
150
+
151
+ headers = {
152
+ "Content-Type": "application/json",
153
+ "Authorization": f"Bearer {API_KEY}",
154
+ }
155
+ payload = {
156
+ "enable_safety_checker": enable_safety_checker,
157
+ "image": base64_image,
158
+ "prompt": prompt,
159
+ "seed": int(seed) if seed != -1 else random.randint(0, 999999)
160
+ }
161
+
162
+ response = requests.post(
163
+ "https://api.wavespeed.ai/api/v3/wavespeed-ai/flux-kontext-dev-ultra-fast",
164
+ headers=headers,
165
+ json=payload,
166
+ timeout=30)
167
+ response.raise_for_status()
168
+
169
+ request_id = response.json()["data"]["id"]
170
+ result_url = f"https://api.wavespeed.ai/api/v3/predictions/{request_id}/result"
171
+ start_time = time.time()
172
+
173
+ for _ in range(60):
174
+ time.sleep(1)
175
+ resp = requests.get(result_url, headers=headers)
176
+ resp.raise_for_status()
177
+
178
+ data = resp.json()["data"]
179
+ status = data["status"]
180
+
181
+ if status == "completed":
182
+ elapsed = time.time() - start_time
183
+ image_url = data["outputs"][0]
184
+ session["history"].append(image_url)
185
+ yield f"🎉 Generation successful! Time taken {elapsed:.1f}s", image_url, image_url
186
+ return
187
+ elif status == "failed":
188
+ raise Exception(data.get("error", "Unknown error"))
189
+ else:
190
+ yield f"⏳ Current status: {status.capitalize()}...", None, None
191
+
192
+ raise Exception("Generation timed out")
193
+
194
+ except Exception as e:
195
+ error_img = create_error_image(str(e))
196
+ yield f"❌ Generation failed: {str(e)}", error_img, ""
197
+
198
+
199
+ def cleanup_task():
200
+ while True:
201
+ session_manager.cleanup_sessions()
202
+ time.sleep(3600)
203
+
204
+
205
+ with gr.Blocks(theme=gr.themes.Soft(),
206
+ css="""
207
+ .status-box { padding: 10px; border-radius: 5px; margin: 5px; }
208
+ .safe { background: #e8f5e9; border: 1px solid #a5d6a7; }
209
+ .warning { background: #fff3e0; border: 1px solid #ffcc80; }
210
+ .error { background: #ffebee; border: 1px solid #ef9a9a; }
211
+ """) as app:
212
+
213
+ session_id = gr.State(str(uuid.uuid4()))
214
+
215
+ gr.Markdown("# 🖼️FLUX Kontext Dev Ultra Fast Live On Wavespeed AI")
216
+ gr.Markdown(
217
+ "FLUX Kontext Dev is a new SOTA image editing model published by Black Forest Labs. We have deployed it on [WaveSpeedAI](https://wavespeed.ai/) for ultra-fast image editing. You can use it to edit images in various styles, add objects, or even change the mood of the image. It supports both text prompts and image inputs."
218
+ )
219
+ gr.Markdown(
220
+ "[FLUX Kontext Dev on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev)"
221
+ )
222
+ gr.Markdown(
223
+ "[FLUX Kontext Dev Ultra Fast on WaveSpeedAI](https://wavespeed.ai/models/wavespeed-ai/flux-kontext-dev-ultra-fast)"
224
+ )
225
+
226
+ with gr.Row():
227
+ with gr.Column(scale=1):
228
+ image_file = gr.Image(label="Upload Image",
229
+ type="filepath",
230
+ sources=["upload"],
231
+ interactive=True,
232
+ image_mode="RGB")
233
+ prompt = gr.Textbox(label="Prompt",
234
+ placeholder="Please enter your prompt...",
235
+ lines=3)
236
+ seed = gr.Number(label="seed",
237
+ value=-1,
238
+ minimum=-1,
239
+ maximum=999999,
240
+ step=1)
241
+ random_btn = gr.Button("random🎲seed", variant="secondary")
242
+ enable_safety = gr.Checkbox(label="🔒 Enable Safety Checker",
243
+ value=True,
244
+ interactive=False)
245
+ with gr.Column(scale=1):
246
+ output_image = gr.Image(label="Generated Result")
247
+ output_url = gr.Textbox(label="Image URL",
248
+ interactive=True,
249
+ visible=False)
250
+ status = gr.Textbox(label="Status", elem_classes=["status-box"])
251
+ submit_btn = gr.Button("Start Generation", variant="primary")
252
+ gr.Examples(
253
+ examples=[
254
+ [
255
+ "Convert the image into Claymation style.",
256
+ "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/penguin.png"
257
+ ],
258
+ [
259
+ "Convert the image into a Ghibli style.",
260
+ "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flux_ip_adapter_input.jpg"
261
+ ],
262
+ [
263
+ "Add sunglasses to the face of the girl.",
264
+ "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/ip_mask_girl2.png"
265
+ ],
266
+ # [
267
+ # 'Convert the image into an ink sketch style.',
268
+ # "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"
269
+ # ],
270
+ # [
271
+ # 'Add a butterfly to the scene.',
272
+ # "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_depth_result.png"
273
+ # ]
274
+ ],
275
+ inputs=[prompt, image_file],
276
+ label="Examples")
277
+
278
+ random_btn.click(fn=lambda: random.randint(0, 999999), outputs=seed)
279
+
280
+ submit_btn.click(
281
+ generate_image,
282
+ inputs=[image_file, prompt, seed, session_id, enable_safety],
283
+ outputs=[status, output_image, output_url],
284
+ api_name=False,
285
+ )
286
+
287
+ if __name__ == "__main__":
288
+ threading.Thread(target=cleanup_task, daemon=True).start()
289
+ app.queue(max_size=8).launch(
290
+ server_name="0.0.0.0",
291
+ share=False,
292
+ )
requirements.txt ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ gradio
2
+ aiohttp
3
+ plotly
4
+ python-dotenv
5
+ pydantic==2.8.2
6
+ torch
7
+ transformers==4.37.2