creek
The AI Image Editor of 2030git clone https://github.com/maydayv7/creek
flask/modal_app.py (17258B)
1import os
2import io
3import sys
4import base64
5import json
6import modal
7
8# ==============================================================================
9# 1. DEFINE THE CLOUD ENVIRONMENT
10# ==============================================================================
11image = (
12 modal.Image.debian_slim(python_version="3.11")
13 .apt_install(
14 "libgl1-mesa-glx",
15 "libglib2.0-0",
16 "libstdc++6",
17 "libxext6",
18 "libsm6",
19 "libxrender1",
20 )
21 .pip_install(
22 "torch",
23 "torchvision",
24 "transformers",
25 "diffusers",
26 "accelerate",
27 "safetensors",
28 "opencv-python-headless",
29 "pillow",
30 "numpy",
31 "scipy",
32 "bitsandbytes",
33 "timm",
34 "einops",
35 "kornia",
36 "flask-cors",
37 "fastapi[standard]",
38 "fal-client",
39 "requests",
40 "pycryptodome",
41 )
42 # --- MOUNT LOCAL MODELS ---
43 .add_local_dir("local_inpainting_model", remote_path="/models/sd-inpainting")
44 .add_local_dir("Florence-2-4bit-Quantized", remote_path="/models/florence-2")
45 .add_local_dir("BiRefNet", remote_path="/root/BiRefNet")
46)
47
48app = modal.App("creekui", image=image)
49
50
51# ==============================================================================
52# 2. THE BACKEND SERVER CLASS
53# ==============================================================================
54@app.cls(
55 gpu="A10G",
56 scaledown_window=300,
57 secrets=[modal.Secret.from_name("creek-secrets")],
58)
59class ModelBackend:
60
61 @modal.enter()
62 def load_models(self):
63 """Runs once when container starts."""
64 print("⏳ Loading models into GPU memory...")
65 import torch
66 import torch.nn as nn
67 import sys
68 from Crypto.Cipher import AES
69
70 self.device = "cuda"
71
72 # --- 1. SETUP CRYPTO ---
73 secret_key_b64 = os.environ.get("SHARED_SECRET_KEY")
74 if not secret_key_b64:
75 raise ValueError("SHARED_SECRET_KEY not set in Modal Secrets")
76
77 # Define CryptoManager inside container
78 class CryptoManager:
79 def __init__(self, key_base64):
80 self.key = base64.b64decode(key_base64)
81
82 def decrypt(self, encrypted_b64):
83 try:
84 data = base64.b64decode(encrypted_b64)
85 nonce = data[:12]
86 tag = data[-16:]
87 ciphertext = data[12:-16]
88 cipher = AES.new(self.key, AES.MODE_GCM, nonce=nonce)
89 return cipher.decrypt_and_verify(ciphertext, tag).decode("utf-8")
90 except Exception as e:
91 print(f"Decryption failed: {e}")
92 return None
93
94 def encrypt(self, plain_text):
95 nonce = os.urandom(12)
96 cipher = AES.new(self.key, AES.MODE_GCM, nonce=nonce)
97 ciphertext, tag = cipher.encrypt_and_digest(plain_text.encode("utf-8"))
98 combined = nonce + ciphertext + tag
99 return base64.b64encode(combined).decode("utf-8")
100
101 self.crypto = CryptoManager(secret_key_b64)
102 print("✅ Crypto Initialized")
103
104 # Check FAL KEY
105 if "FAL_KEY" not in os.environ:
106 print("❌ Error: FAL_KEY secret not found!")
107 else:
108 print("✅ FAL_KEY loaded securely.")
109
110 # --- 2. SETUP PATHS & MODELS ---
111 sys.path.append("/root/BiRefNet")
112
113 # Load Stable Diffusion
114 from diffusers import StableDiffusionInpaintPipeline
115
116 self.sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
117 "/models/sd-inpainting",
118 torch_dtype=torch.float16,
119 use_safetensors=True,
120 local_files_only=True,
121 ).to(self.device)
122 self.sd_pipe.enable_attention_slicing()
123 print("✅ Stable Diffusion Loaded")
124
125 # Load Florence-2
126 import transformers.dynamic_module_utils
127
128 transformers.dynamic_module_utils.check_imports = lambda f: []
129
130 # Patch for _supports_sdpa
131 _old_getattr = nn.Module.__getattr__
132
133 def _fixed_getattr(self, name):
134 if name == "_supports_sdpa":
135 return False
136 return _old_getattr(self, name)
137
138 nn.Module.__getattr__ = _fixed_getattr
139
140 from transformers import AutoModelForCausalLM, AutoProcessor, BitsAndBytesConfig
141
142 bnb_config = BitsAndBytesConfig(
143 load_in_4bit=True,
144 bnb_4bit_quant_type="nf4",
145 bnb_4bit_compute_dtype=torch.float16,
146 )
147 self.florence_model = AutoModelForCausalLM.from_pretrained(
148 "/models/florence-2",
149 quantization_config=bnb_config,
150 trust_remote_code=True,
151 local_files_only=True,
152 ).to(self.device)
153 self.florence_processor = AutoProcessor.from_pretrained(
154 "/models/florence-2", trust_remote_code=True, local_files_only=True
155 )
156 print("✅ Florence-2 Loaded")
157
158 # Load BiRefNet
159 try:
160 from models.birefnet import BiRefNet
161
162 self.birefnet = BiRefNet(bb_pretrained=False)
163 weight_path = "/root/BiRefNet/birefnet_fp16.pt"
164 state_dict = torch.load(weight_path, map_location=self.device)
165 self.birefnet.load_state_dict(state_dict)
166 self.birefnet.to(self.device).half().eval()
167 print(f"✅ BiRefNet Loaded")
168 except Exception as e:
169 print(f"❌ BiRefNet Error: {e}")
170 self.birefnet = None
171
172 from torchvision import transforms
173
174 self.transform_birefnet = transforms.Compose(
175 [
176 transforms.Resize((1024, 1024)),
177 transforms.ToTensor(),
178 transforms.Normalize(
179 mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
180 ),
181 ]
182 )
183
184 # --- SECURITY WRAPPER ---
185 def _handle_secure_request(self, item: dict, logic_func):
186 """Decrypts input -> Runs Logic -> Encrypts Output"""
187 try:
188 # 1. Decrypt Incoming
189 if "data" not in item:
190 return {"error": "Invalid format. Expected {'data': ...}"}
191
192 decrypted_json_str = self.crypto.decrypt(item["data"])
193 if decrypted_json_str is None:
194 return {"error": "Decryption failed (Check Key)"}
195
196 payload = json.loads(decrypted_json_str)
197
198 # 2. Run Actual Logic
199 result = logic_func(payload)
200
201 # 3. Encrypt Outgoing
202 encrypted_response = self.crypto.encrypt(json.dumps(result))
203 return {"data": encrypted_response}
204
205 except Exception as e:
206 print(f"Request Error: {e}")
207 return {"error": str(e)}
208
209 # ==========================================================================
210 # 3. ENDPOINTS
211 # ==========================================================================
212
213 @modal.fastapi_endpoint(method="POST")
214 def generate(self, item: dict):
215 def logic(data):
216 from PIL import Image
217
218 prompt = data.get("prompt", "A luxury watch")
219 print(f"🎨 Generating: {prompt}")
220 empty = Image.new("RGB", (512, 512))
221 mask = Image.new("L", (512, 512), 255)
222 img = self.sd_pipe(
223 prompt=prompt,
224 image=empty,
225 mask_image=mask,
226 height=512,
227 width=512,
228 num_inference_steps=30,
229 ).images[0]
230 return {"status": "success", "image": self._to_base64(img)}
231
232 return self._handle_secure_request(item, logic)
233
234 @modal.fastapi_endpoint(method="POST")
235 def inpainting(self, item: dict):
236 def logic(data):
237 from PIL import Image, ImageFilter
238 import numpy as np
239 import scipy.ndimage
240 import torch
241
242 prompt = data.get("prompt", "")
243 img_b64 = data.get("image")
244 mask_b64 = data.get("mask_image")
245
246 clean = self._decode_base64(img_b64).convert("RGB")
247 drawn = self._decode_base64(mask_b64).convert("RGB")
248
249 clean = self._resize_to_limit(clean, 512)
250 drawn = drawn.resize(clean.size)
251
252 # Robust Masking
253 clean_blur = np.array(
254 clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
255 )
256 drawn_blur = np.array(
257 drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
258 )
259 mask_arr = np.max(np.abs(drawn_blur - clean_blur), axis=2)
260 mask = Image.fromarray(
261 (scipy.ndimage.binary_fill_holes(mask_arr > 30) * 255).astype(np.uint8)
262 ).filter(ImageFilter.MaxFilter(9))
263
264 # Florence Context
265 inputs = self.florence_processor(
266 text="<DETAILED_CAPTION>", images=[drawn], return_tensors="pt"
267 )
268 inputs = {
269 k: v.to(self.device, torch.float16 if k == "pixel_values" else None)
270 for k, v in inputs.items()
271 }
272 gen_ids = self.florence_model.generate(
273 **inputs, max_new_tokens=128, num_beams=1, use_cache=False
274 )
275 context = (
276 self.florence_processor.batch_decode(
277 gen_ids, skip_special_tokens=False
278 )[0]
279 .replace("</s>", "")
280 .replace("<s>", "")
281 .replace("<DETAILED_CAPTION>", "")
282 .strip()
283 )
284
285 full_prompt = f"{context} {prompt}".strip()
286
287 output = self.sd_pipe(
288 prompt=full_prompt,
289 negative_prompt="blurry, low quality, ugly, text, watermark, bad anatomy, deformed, noisy",
290 image=drawn,
291 mask_image=mask,
292 num_inference_steps=50,
293 strength=0.85,
294 guidance_scale=8.5,
295 ).images[0]
296
297 return {"status": "success", "image": self._to_base64(output)}
298
299 return self._handle_secure_request(item, logic)
300
301 @modal.fastapi_endpoint(method="POST")
302 def inpainting_api(self, item: dict):
303 def logic(data):
304 import fal_client, requests, uuid, scipy.ndimage
305 from PIL import Image, ImageFilter
306 import numpy as np
307
308 prompt = data.get("prompt", "High quality image")
309 img_b64 = data.get("image")
310 mask_b64 = data.get("mask_image")
311
312 clean = self._decode_base64(img_b64).convert("RGB")
313 drawn = self._decode_base64(mask_b64).convert("RGB")
314
315 clean = self._resize_to_limit(clean, 1024)
316 drawn = drawn.resize(clean.size)
317
318 # Masking
319 clean_blur = np.array(
320 clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
321 )
322 drawn_blur = np.array(
323 drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
324 )
325 mask = Image.fromarray(
326 (
327 scipy.ndimage.binary_fill_holes(
328 np.max(np.abs(drawn_blur - clean_blur), axis=2) > 30
329 )
330 * 255
331 ).astype(np.uint8)
332 ).filter(ImageFilter.MaxFilter(9))
333
334 clean_p, mask_p = (
335 f"/tmp/c_{uuid.uuid4()}.png",
336 f"/tmp/m_{uuid.uuid4()}.png",
337 )
338 clean.save(clean_p)
339 mask.save(mask_p)
340
341 try:
342 res = fal_client.submit(
343 "fal-ai/flux-lora-fill",
344 arguments={
345 "prompt": prompt,
346 "image_url": fal_client.upload_file(clean_p),
347 "mask_url": fal_client.upload_file(mask_p),
348 "guidance_scale": 30,
349 "num_inference_steps": 28,
350 "enable_safety_checker": False,
351 },
352 ).get()
353
354 if "images" in res:
355 img_resp = requests.get(res["images"][0]["url"])
356 img = Image.open(io.BytesIO(img_resp.content))
357 return {"status": "success", "image": self._to_base64(img)}
358 return {"status": "error", "message": "No images from Fal"}
359 finally:
360 if os.path.exists(clean_p):
361 os.remove(clean_p)
362 if os.path.exists(mask_p):
363 os.remove(mask_p)
364
365 return self._handle_secure_request(item, logic)
366
367 @modal.fastapi_endpoint(method="POST")
368 def sketch_api(self, item: dict):
369 def logic(data):
370 import fal_client, requests
371 from PIL import Image
372
373 prompt = data.get("prompt", "")
374 option = int(data.get("option", 1))
375
376 enhanced_prompt = (
377 f"{prompt}, sharp focus, high definition, 4k, vector art, crisp lines"
378 )
379
380 if option == 1:
381 # Nano Banana
382 print(f"🍌 Using Nano Banana for: {prompt}")
383 model_id = "fal-ai/nano-banana"
384 arguments = {
385 "prompt": enhanced_prompt,
386 "num_images": 1,
387 "aspect_ratio": "1:1",
388 "output_format": "png",
389 }
390 else:
391 # Flux Dev
392 print(f"🚀 Using Flux Dev for: {prompt}")
393 model_id = "fal-ai/flux/dev"
394 arguments = {
395 "image_size": "square_hd",
396 "num_inference_steps": 28,
397 "guidance_scale": 3.5,
398 "safety_tolerance": "2",
399 "enable_safety_checker": False,
400 "prompt": enhanced_prompt,
401 }
402
403 res = fal_client.submit(model_id, arguments=arguments).get()
404 if "images" in res:
405 img_resp = requests.get(res["images"][0]["url"])
406 img = Image.open(io.BytesIO(img_resp.content)).convert("RGB")
407 return {"status": "success", "image": self._to_base64(img)}
408 return {"status": "error", "message": "No images returned"}
409
410 return self._handle_secure_request(item, logic)
411
412 @modal.fastapi_endpoint(method="POST")
413 def asset(self, item: dict):
414 def logic(data):
415 import torch, numpy as np
416 from PIL import Image
417
418 img_b64 = data.get("image")
419 img = self._decode_base64(img_b64)
420 w, h = img.size
421
422 inp = self.transform_birefnet(img).unsqueeze(0).to(self.device).half()
423 with torch.no_grad():
424 preds = self.birefnet(inp)[-1].sigmoid()
425
426 import torch.nn.functional as F
427
428 res = F.interpolate(preds, size=(h, w), mode="bilinear", align_corners=True)
429 mask = Image.fromarray((res.squeeze().cpu().numpy() * 255).astype(np.uint8))
430 img.putalpha(mask)
431
432 return {"status": "success", "image": self._to_base64(img)}
433
434 return self._handle_secure_request(item, logic)
435
436 @modal.fastapi_endpoint(method="POST")
437 def describe(self, item: dict):
438 def logic(data):
439 import torch
440
441 img_b64 = data.get("image")
442 img = self._decode_base64(img_b64)
443 prompt = data.get("prompt", "<DETAILED_CAPTION>")
444
445 inputs = self.florence_processor(
446 text=prompt, images=img, return_tensors="pt"
447 )
448 inputs = {
449 k: v.to(self.device, torch.float16 if k == "pixel_values" else None)
450 for k, v in inputs.items()
451 }
452
453 gen_ids = self.florence_model.generate(
454 **inputs, max_new_tokens=1024, num_beams=3, use_cache=False
455 )
456 txt = self.florence_processor.batch_decode(
457 gen_ids, skip_special_tokens=False
458 )[0]
459 clean_txt = (
460 txt.replace("</s>", "").replace("<s>", "").replace(prompt, "").strip()
461 )
462
463 return {"status": "success", "output": clean_txt}
464
465 return self._handle_secure_request(item, logic)
466
467 # --- HELPERS ---
468 def _decode_base64(self, b64_str):
469 from PIL import Image
470
471 if "," in b64_str:
472 b64_str = b64_str.split(",")[1]
473 return Image.open(io.BytesIO(base64.b64decode(b64_str))).convert("RGB")
474
475 def _to_base64(self, img):
476 buffered = io.BytesIO()
477 img.save(buffered, format="PNG")
478 return base64.b64encode(buffered.getvalue()).decode("utf-8")
479
480 def _resize_to_limit(self, img, max_dim=1024, multiple=8):
481 from PIL import Image
482
483 w, h = img.size
484 ratio = min(max_dim / w, max_dim / h)
485 new_w = int(w * ratio)
486 new_h = int(h * ratio)
487 new_w = new_w - (new_w % multiple)
488 new_h = new_h - (new_h % multiple)
489 if new_w < multiple:
490 new_w = multiple
491 if new_h < multiple:
492 new_h = multiple
493 return img.resize((new_w, new_h), Image.LANCZOS)