creek
The AI Image Editor of 2030git clone https://github.com/maydayv7/creek
flask/index.py (30832B)
1import os
2import io
3import sys
4import base64
5import cv2
6import re
7import numpy as np
8import torch
9import traceback
10import requests
11import scipy.ndimage
12import json
13from functools import wraps
14from flask import Flask, jsonify, request, send_file
15from flask_cors import CORS
16from PIL import Image, ImageFilter
17from torchvision import transforms
18import time
19import uuid
20from Crypto.Cipher import AES
21from dotenv import load_dotenv
22
23load_dotenv()
24
25
26class CryptoManager:
27 def __init__(self, key_base64):
28 # Decode the base64 key to raw bytes (must be 32 bytes for AES-256)
29 self.key = base64.b64decode(key_base64)
30
31 def encrypt(self, plain_text):
32 # 1. Generate a random unique Nonce (12 bytes is standard for GCM)
33 nonce = os.urandom(12)
34
35 # 2. Initialize Cipher
36 cipher = AES.new(self.key, AES.MODE_GCM, nonce=nonce)
37
38 # 3. Encrypt and get Tag (MAC)
39 ciphertext, tag = cipher.encrypt_and_digest(plain_text.encode("utf-8"))
40
41 # 4. Pack: Nonce + Ciphertext + Tag
42 combined = nonce + ciphertext + tag
43
44 # 5. Return as Base64 string
45 return base64.b64encode(combined).decode("utf-8")
46
47 def decrypt(self, encrypted_b64):
48 try:
49 # 1. Decode Base64
50 data = base64.b64decode(encrypted_b64)
51
52 # 2. Unpack (Slice the bytes)
53 nonce = data[:12]
54 tag = data[-16:]
55 ciphertext = data[12:-16]
56
57 # 3. Decrypt
58 cipher = AES.new(self.key, AES.MODE_GCM, nonce=nonce)
59 decrypted_data = cipher.decrypt_and_verify(ciphertext, tag)
60 return decrypted_data.decode("utf-8")
61 except Exception as e:
62 print(f"Decryption failed: {e}")
63 return None
64
65
66# Setup Secret Key
67SHARED_SECRET_KEY = os.getenv("SHARED_SECRET_KEY")
68if not SHARED_SECRET_KEY:
69 print("❌ Error: SHARED_SECRET_KEY not found in .env")
70 sys.exit(1)
71
72crypto = CryptoManager(SHARED_SECRET_KEY)
73
74
75# ==============================================================================
76# SECURITY DECORATOR (Middleware)
77# ==============================================================================
78def secure_endpoint(f):
79 @wraps(f)
80 def decorated_function(*args, **kwargs):
81 # --- 1. INCOMING DECRYPTION ---
82 try:
83 # Expecting JSON format: { "data": "BASE64_ENCRYPTED_STRING" }
84 incoming = request.get_json(silent=True)
85 if not incoming or "data" not in incoming:
86 return (
87 jsonify(
88 {
89 "error": "Invalid format. Expected {'data': 'encrypted_string'}"
90 }
91 ),
92 400,
93 )
94
95 encrypted_b64 = incoming["data"]
96 decrypted_json_str = crypto.decrypt(encrypted_b64)
97
98 if decrypted_json_str is None:
99 return jsonify({"error": "Decryption failed (Check Key or Nonce)"}), 403
100
101 # Parse the decrypted string back to a Python dictionary
102 decrypted_payload = json.loads(decrypted_json_str)
103
104 # OVERRIDE request.get_json() so the inner function sees the decrypted data
105 request.get_json = lambda **k: decrypted_payload
106
107 except Exception as e:
108 return jsonify({"error": f"Security Middleware Error: {str(e)}"}), 500
109
110 # --- 2. EXECUTE ORIGINAL LOGIC ---
111 response = f(*args, **kwargs)
112
113 # --- 3. OUTGOING ENCRYPTION ---
114 try:
115 # Handle Flask response tuples (e.g., jsonify(...), 500)
116 resp_obj = response
117 status_code = 200
118 if isinstance(response, tuple):
119 resp_obj = response[0]
120 if len(response) > 1:
121 status_code = response[1]
122
123 # Extract the plain JSON data from the Response object
124 if hasattr(resp_obj, "get_json"):
125 plain_data = resp_obj.get_json()
126 else:
127 # Fallback if it's not a response object yet
128 plain_data = resp_obj
129
130 # Convert dict -> JSON String -> Encrypt
131 plain_json_str = json.dumps(plain_data)
132 encrypted_response = crypto.encrypt(plain_json_str)
133
134 # Return standard encrypted wrapper
135 return jsonify({"data": encrypted_response}), status_code
136
137 except Exception as e:
138 return jsonify({"error": f"Response Encryption Error: {str(e)}"}), 500
139
140 return decorated_function
141
142
143# --- FAL.AI IMPORTS ---
144try:
145 import fal_client
146
147 # SETUP API KEY
148 if os.getenv("FAL_KEY"):
149 os.environ["FAL_KEY"] = os.getenv("FAL_KEY")
150 FAL_AVAILABLE = True
151 else:
152 print("⚠️ Warning: FAL_KEY not found in environment variables.")
153 FAL_AVAILABLE = False
154except ImportError:
155 print("⚠️ Fal.ai Client not installed. /inpainting-api will fail.")
156 FAL_AVAILABLE = False
157
158# --- FLORENCE-2 IMPORTS ---
159try:
160 import bitsandbytes
161 from transformers import AutoModelForCausalLM, AutoProcessor, BitsAndBytesConfig
162 import transformers.dynamic_module_utils
163 import torch.nn as nn
164
165 FLORENCE_AVAILABLE = True
166except ImportError as e:
167 print(
168 f"⚠️ Florence-2 Disabled: {e} (Ensure 'bitsandbytes' and 'transformers' are installed)"
169 )
170 FLORENCE_AVAILABLE = False
171except Exception as e:
172 print(f"⚠️ Florence-2 Disabled: Unexpected initialization error: {e}")
173 FLORENCE_AVAILABLE = False
174
175# --- IMPORT BIREFNET ---
176current_dir = os.path.dirname(os.path.abspath(__file__))
177birefnet_path = os.path.join(current_dir, "BiRefNet")
178
179if os.path.exists(birefnet_path):
180 if birefnet_path not in sys.path:
181 sys.path.append(birefnet_path)
182 print(f"✅ Added {birefnet_path} to system path.")
183else:
184 print(f"❌ Error: '{birefnet_path}' not found. Please clone the repository.")
185
186try:
187 from models.birefnet import BiRefNet
188
189 print("✅ BiRefNet imported successfully.")
190except ImportError as e:
191 print(f"⚠️ Import Error: {e}")
192 try:
193 import BiRefNet.models.birefnet as brn
194
195 BiRefNet = brn.BiRefNet
196 print("✅ BiRefNet imported via package path.")
197 except ImportError:
198 print(
199 "❌ Failed to import BiRefNet. Ensure 'BiRefNet/models/birefnet.py' exists."
200 )
201
202# --- IMPORT STABLE DIFFUSION ---
203try:
204 from diffusers import StableDiffusionInpaintPipeline, AutoPipelineForInpainting
205except ImportError:
206 print("⚠️ Diffusers not found. SD features disabled.")
207 StableDiffusionInpaintPipeline = None
208 AutoPipelineForInpainting = None
209
210app = Flask(__name__)
211CORS(app)
212
213DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
214print(f"🚀 Running on device: {DEVICE}")
215
216# ==============================================================================
217# 1. LOAD STABLE DIFFUSION
218# ==============================================================================
219print("⏳ Loading Stable Diffusion (Inpainting)...")
220sd_pipe = None
221try:
222 if StableDiffusionInpaintPipeline:
223 SD_MODEL_ID = (
224 "./local_inpainting_model"
225 if os.path.exists("./local_inpainting_model")
226 else "runwayml/stable-diffusion-inpainting"
227 )
228
229 sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
230 SD_MODEL_ID,
231 torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
232 use_safetensors=True,
233 ).to(DEVICE)
234 sd_pipe.enable_attention_slicing()
235 sd_pipe.enable_model_cpu_offload()
236 print("✅ Stable Diffusion Loaded!")
237except Exception as e:
238 print(f"❌ Failed to load SD: {e}")
239
240# ==============================================================================
241# 2. LOAD BIREFNET
242# ==============================================================================
243print("⏳ Loading BiRefNet...")
244birefnet_model = None
245BIREFNET_WEIGHTS = "./BiRefNet/birefnet_fp16.pt"
246BIREFNET_SIZE = (1024, 1024)
247
248try:
249 if "BiRefNet" in locals() and os.path.exists(BIREFNET_WEIGHTS):
250 birefnet_model = BiRefNet(bb_pretrained=False)
251 state_dict = torch.load(BIREFNET_WEIGHTS, map_location=DEVICE)
252 birefnet_model.load_state_dict(state_dict)
253 birefnet_model.to(DEVICE)
254 if DEVICE == "cuda":
255 birefnet_model.half()
256 birefnet_model.eval()
257 print("✅ BiRefNet Weights Loaded!")
258 else:
259 print(f"⚠️ BiRefNet skipped. Weights found: {os.path.exists(BIREFNET_WEIGHTS)}")
260except Exception as e:
261 print(f"❌ Failed to load BiRefNet: {e}")
262
263transform_birefnet = transforms.Compose(
264 [
265 transforms.Resize(BIREFNET_SIZE),
266 transforms.ToTensor(),
267 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
268 ]
269)
270
271# ==============================================================================
272# 3. LOAD FLORENCE-2 (QUANTIZED)
273# ==============================================================================
274print("⏳ Loading Florence-2...")
275florence_model = None
276florence_processor = None
277FLORENCE_PATH = os.path.join(current_dir, "Florence-2-4bit-Quantized")
278
279if FLORENCE_AVAILABLE:
280 try:
281
282 def check_imports_fixed(filename):
283 return []
284
285 transformers.dynamic_module_utils.check_imports = check_imports_fixed
286
287 _old_getattr = nn.Module.__getattr__
288
289 def _fixed_getattr(self, name):
290 if name == "_supports_sdpa":
291 return False
292 return _old_getattr(self, name)
293
294 nn.Module.__getattr__ = _fixed_getattr
295
296 if os.path.exists(FLORENCE_PATH):
297 bnb_config = BitsAndBytesConfig(
298 load_in_4bit=True,
299 bnb_4bit_quant_type="nf4",
300 bnb_4bit_compute_dtype=torch.float16,
301 )
302
303 florence_model = AutoModelForCausalLM.from_pretrained(
304 FLORENCE_PATH,
305 quantization_config=bnb_config,
306 trust_remote_code=True,
307 device_map="cuda" if DEVICE == "cuda" else "cpu",
308 local_files_only=True,
309 )
310 florence_processor = AutoProcessor.from_pretrained(
311 FLORENCE_PATH, trust_remote_code=True
312 )
313 print("✅ Florence-2 Loaded Successfully!")
314 else:
315 print(f"⚠️ Florence-2 folder not found at: {FLORENCE_PATH}")
316 except Exception as e:
317 import traceback; traceback.print_exc()
318 print(f"❌ Failed to load Florence-2: {e}")
319
320
321# ==============================================================================
322# HELPER FUNCTIONS
323# ==============================================================================
324def decode_base64_image(b64_str):
325 if "," in b64_str:
326 b64_str = b64_str.split(",")[1]
327 image_data = base64.b64decode(b64_str)
328 img = Image.open(io.BytesIO(image_data))
329
330 if img.mode in ("RGBA", "LA") or (img.mode == "P" and "transparency" in img.info):
331 background = Image.new("RGB", img.size, (255, 255, 255))
332 if img.mode == "P":
333 img = img.convert("RGBA")
334 background.paste(img, mask=img.split()[3])
335 return background
336 else:
337 return img.convert("RGB")
338
339
340def encode_image_to_base64(pil_img):
341 buffered = io.BytesIO()
342 pil_img.save(buffered, format="PNG")
343 return base64.b64encode(buffered.getvalue()).decode("utf-8")
344
345
346def process_birefnet_output(preds, original_size):
347 if isinstance(preds, (list, tuple)):
348 pred_tensor = preds[-1]
349 else:
350 pred_tensor = preds
351
352 pred_tensor = pred_tensor.sigmoid().cpu()
353 mask_np = pred_tensor.squeeze().numpy().astype(np.float32)
354
355 if len(mask_np.shape) > 2:
356 mask_np = mask_np[0]
357
358 mask_resized = cv2.resize(mask_np, original_size, interpolation=cv2.INTER_LINEAR)
359 mask = (mask_resized > 0.5).astype(np.uint8) * 255
360
361 return Image.fromarray(mask)
362
363
364def resize_to_limit(img, max_dim=1024, multiple=8):
365 w, h = img.size
366 ratio = min(max_dim / w, max_dim / h)
367 new_w = int(w * ratio)
368 new_h = int(h * ratio)
369 new_w = new_w - (new_w % multiple)
370 new_h = new_h - (new_h % multiple)
371 if new_w < multiple:
372 new_w = multiple
373 if new_h < multiple:
374 new_h = multiple
375 return img.resize((new_w, new_h), Image.LANCZOS)
376
377
378# ==============================================================================
379# ROUTES
380# ==============================================================================
381
382
383@app.route("/")
384def index():
385 return "Image Processing API is running."
386
387
388@app.route("/test-encrypt", methods=["POST"])
389def test_encrypt():
390 # Helper route to debug encryption/decryption
391 try:
392 data = request.get_json()
393 plain_text = data.get("text", "Hello, World!")
394 encrypted = crypto.encrypt(plain_text)
395 decrypted = crypto.decrypt(encrypted)
396 return jsonify(
397 {"original": plain_text, "encrypted": encrypted, "decrypted": decrypted}
398 )
399 except Exception as e:
400 return jsonify({"error": str(e)}), 500
401
402
403@app.route("/generate", methods=["POST"])
404@secure_endpoint
405def generate_image():
406 if not sd_pipe:
407 return jsonify({"error": "SD Model not loaded"}), 500
408 try:
409 data = request.get_json()
410 prompt = data.get(
411 "prompt",
412 "The image shows a river running through a lush green valley surrounded by trees, plants, grass, and poles. In the background, the sky is filled with clouds, creating a peaceful atmosphere.",
413 )
414 empty_image = Image.new("RGB", (512, 512), (0, 0, 0))
415 full_mask = Image.new("L", (512, 512), 255)
416 print(f"🎨 Generating: {prompt}")
417 image = sd_pipe(
418 prompt=prompt,
419 image=empty_image,
420 mask_image=full_mask,
421 height=512,
422 width=512,
423 num_inference_steps=30,
424 ).images[0]
425 return jsonify({"status": "success", "image": encode_image_to_base64(image)})
426 except Exception as e:
427 return jsonify({"status": "error", "message": str(e)}), 500
428
429
430@app.route("/inpainting", methods=["POST"])
431@secure_endpoint
432def inpaint_image():
433 if not sd_pipe:
434 return jsonify({"error": "SD Model not loaded"}), 500
435 try:
436 data = request.get_json()
437 user_prompt = data.get("prompt", "")
438 clean_b64 = data.get("image")
439 drawn_b64 = data.get("mask_image")
440
441 if not clean_b64 or not drawn_b64:
442 return jsonify({"error": "Missing image or mask"}), 400
443
444 # 1. Decode Images
445 raw_clean = decode_base64_image(clean_b64).convert("RGB")
446 raw_drawn = decode_base64_image(drawn_b64).convert("RGB")
447
448 # 2. Resize maintaining Aspect Ratio (Max 512 for Local SD)
449 img_clean = resize_to_limit(raw_clean, max_dim=512)
450 # Resize drawn image to match the clean image exactly
451 img_drawn = raw_drawn.resize(img_clean.size)
452
453 print(f"🔍 Calculating Robust Difference Mask (Size: {img_clean.size})...")
454 clean_blur = np.array(
455 img_clean.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16
456 )
457 drawn_blur = np.array(
458 img_drawn.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16
459 )
460 diff_arr = np.abs(drawn_blur - clean_blur)
461 mask_arr = np.max(diff_arr, axis=2)
462 mask_binary = mask_arr > 30
463 mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
464 mask_image = Image.fromarray((mask_filled * 255).astype(np.uint8))
465 mask_image = mask_image.filter(ImageFilter.MaxFilter(9))
466 print("✅ Mask calculated.")
467
468 generated_prompt = ""
469 if florence_model and florence_processor:
470 print("👁️ Generating context with Florence-2...")
471 try:
472 task_prompt = "<DETAILED_CAPTION>"
473 inputs = florence_processor(
474 text=task_prompt, images=[img_drawn], return_tensors="pt"
475 )
476 inputs["pixel_values"] = inputs["pixel_values"].to(
477 DEVICE, torch.float16
478 )
479 inputs["input_ids"] = inputs["input_ids"].to(DEVICE)
480
481 generated_ids = florence_model.generate(
482 input_ids=inputs["input_ids"],
483 pixel_values=inputs["pixel_values"],
484 max_new_tokens=128,
485 num_beams=1,
486 do_sample=False,
487 use_cache=False,
488 )
489
490 generated_text = florence_processor.batch_decode(
491 generated_ids, skip_special_tokens=False
492 )[0]
493 generated_prompt = (
494 generated_text.replace(task_prompt, "")
495 .replace("</s>", "")
496 .replace("<s>", "")
497 .strip()
498 )
499 print(f"📝 Florence Generated: {generated_prompt}")
500 except Exception as e:
501 print(f"⚠️ Florence captioning failed: {e}")
502
503 final_prompt = f"{generated_prompt} {user_prompt}".strip()
504 negative_prompt = (
505 "blurry, low quality, ugly, text, watermark, bad anatomy, deformed, noisy"
506 )
507 print(f"✨ Final Inpaint Prompt: {final_prompt}")
508
509 save_dir = "input_data"
510 os.makedirs(save_dir, exist_ok=True)
511 timestamp = int(time.time())
512 img_clean.save(os.path.join(save_dir, f"clean_{timestamp}.png"))
513 img_drawn.save(os.path.join(save_dir, f"drawn_{timestamp}.png"))
514 mask_image.save(os.path.join(save_dir, f"generated_mask_{timestamp}.png"))
515
516 print(f"🎨 Running Inference with strength=0.85...")
517 image = sd_pipe(
518 prompt=final_prompt,
519 negative_prompt=negative_prompt,
520 image=img_drawn,
521 mask_image=mask_image,
522 num_inference_steps=50,
523 strength=0.85,
524 guidance_scale=8.5,
525 ).images[0]
526
527 final_image_path = os.path.join(save_dir, f"result_{timestamp}.png")
528 image.save(final_image_path)
529 print(f"💾 Saved output to {final_image_path}")
530
531 return jsonify({"status": "success", "image": encode_image_to_base64(image)})
532
533 except Exception as e:
534 print(f"❌ Inpainting Error: {e}")
535 traceback.print_exc()
536 return jsonify({"status": "error", "message": str(e)}), 500
537
538
539@app.route("/asset", methods=["POST"])
540@secure_endpoint
541def remove_background():
542 if not birefnet_model:
543 return jsonify({"error": "BiRefNet not loaded"}), 500
544 try:
545 data = request.get_json()
546 image_b64 = data.get("image")
547 if not image_b64:
548 return jsonify({"error": "No image provided"}), 400
549
550 original_image = decode_base64_image(image_b64)
551 orig_w, orig_h = original_image.size
552
553 input_tensor = transform_birefnet(original_image).unsqueeze(0).to(DEVICE)
554 if DEVICE == "cuda":
555 input_tensor = input_tensor.half()
556
557 print("✂️ Removing background...")
558 with torch.no_grad():
559 preds = birefnet_model(input_tensor)
560
561 mask_pil = process_birefnet_output(preds, (orig_w, orig_h))
562 original_image.putalpha(mask_pil)
563
564 return jsonify(
565 {"status": "success", "image": encode_image_to_base64(original_image)}
566 )
567 except Exception as e:
568 print(f"❌ Error: {e}")
569 traceback.print_exc()
570 return jsonify({"status": "error", "message": str(e)}), 500
571
572
573@app.route("/describe", methods=["POST"])
574@secure_endpoint
575def describe_image():
576 if not florence_model or not florence_processor:
577 return jsonify({"error": "Florence-2 not loaded"}), 500
578 try:
579 data = request.get_json()
580 image_b64 = data.get("image")
581 prompt_type = data.get("prompt", "<DETAILED_CAPTION>")
582
583 if not image_b64:
584 return jsonify({"error": "No image provided"}), 400
585
586 image = decode_base64_image(image_b64)
587 print(f"👁️ Analyzing image with Florence-2...")
588
589 inputs = florence_processor(
590 text=prompt_type, images=[image], return_tensors="pt"
591 )
592 inputs["pixel_values"] = inputs["pixel_values"].to(DEVICE, torch.float16)
593 inputs["input_ids"] = inputs["input_ids"].to(DEVICE)
594
595 generated_ids = florence_model.generate(
596 input_ids=inputs["input_ids"],
597 pixel_values=inputs["pixel_values"],
598 max_new_tokens=128,
599 num_beams=1,
600 do_sample=False,
601 use_cache=False,
602 )
603
604 generated_text = florence_processor.batch_decode(
605 generated_ids, skip_special_tokens=False
606 )[0]
607 cleaned_text = (
608 generated_text.replace(prompt_type, "")
609 .replace("</s>", "")
610 .replace("<s>", "")
611 .strip()
612 )
613
614 if cleaned_text and cleaned_text[-1] not in [".", "!", "?"]:
615 last_dot = cleaned_text.rfind(".")
616 last_excl = cleaned_text.rfind("!")
617 last_ques = cleaned_text.rfind("?")
618 cut_off = max(last_dot, last_excl, last_ques)
619 if cut_off != -1:
620 cleaned_text = cleaned_text[: cut_off + 1]
621
622 final_answer = cleaned_text
623 print(final_answer)
624
625 if "<loc_" in cleaned_text or "<poly_" in cleaned_text:
626 try:
627 parsed = florence_processor.post_process_generation(
628 generated_text,
629 task=prompt_type,
630 image_size=(image.width, image.height),
631 )
632 if isinstance(parsed, dict) and prompt_type in parsed:
633 final_answer = parsed[prompt_type]
634 else:
635 final_answer = parsed
636 except Exception:
637
638 def parse_loc_manually(text, w, h):
639 locs = re.findall(r"<loc_(\d+)>", text)
640 if locs and len(locs) % 4 == 0:
641 bboxes = []
642 for i in range(0, len(locs), 4):
643 x1 = int(int(locs[i]) / 1000 * w)
644 y1 = int(int(locs[i + 1]) / 1000 * h)
645 x2 = int(int(locs[i + 2]) / 1000 * w)
646 y2 = int(int(locs[i + 3]) / 1000 * h)
647 bboxes.append([x1, y1, x2, y2])
648 clean_text = re.sub(r"<loc_\d+>", "", text).strip()
649 return {"text": clean_text, "bboxes": bboxes}
650 return text
651
652 final_answer = parse_loc_manually(
653 cleaned_text, image.width, image.height
654 )
655
656 return jsonify({"status": "success", "output": final_answer})
657
658 except Exception as e:
659 print(f"❌ Florence Error: {e}")
660 traceback.print_exc()
661 return jsonify({"status": "error", "message": str(e)}), 500
662
663
664@app.route("/inpainting-api", methods=["POST"])
665@secure_endpoint
666def inpainting_api_fal():
667 if not FAL_AVAILABLE:
668 return jsonify({"error": "Fal.ai client not installed or API Key missing"}), 500
669
670 try:
671 data = request.get_json()
672
673 clean_b64 = data.get("image")
674 drawn_b64 = data.get("mask_image")
675 prompt = data.get(
676 "prompt",
677 "The image shows a river running through a lush green valley surrounded by trees, plants, grass, and poles. In the background, the sky is filled with clouds, creating a peaceful atmosphere.",
678 )
679
680 if not clean_b64 or not drawn_b64:
681 return (
682 jsonify({"error": "Missing 'image' (clean) or 'mask_image' (drawn)"}),
683 400,
684 )
685
686 print(f"📥 Received Request: Prompt='{prompt}'")
687
688 # 1. Decode Images
689 raw_clean = decode_base64_image(clean_b64).convert("RGB")
690 raw_drawn = decode_base64_image(drawn_b64).convert("RGB")
691
692 # 2. Resize maintaining Aspect Ratio (Max 1024 for Flux)
693 img_clean = resize_to_limit(raw_clean, max_dim=1024)
694 # Resize drawn to match exactly
695 img_drawn = raw_drawn.resize(img_clean.size)
696
697 # 3. Setup Debug Directory
698 debug_dir = "debug_fal"
699 os.makedirs(debug_dir, exist_ok=True)
700 unique_id = str(int(time.time()))
701
702 clean_path = os.path.join(debug_dir, f"fal_clean_{unique_id}.png")
703 mask_path = os.path.join(debug_dir, f"fal_mask_{unique_id}.png")
704 fal_result_path = os.path.join(debug_dir, f"fal_result_{unique_id}.png")
705
706 # 4. Mask Generation
707 print(f"🛠️ Generating mask (Size: {img_clean.size})...")
708 clean_blur = np.array(
709 img_clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
710 )
711 drawn_blur = np.array(
712 img_drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
713 )
714
715 diff_arr = np.abs(drawn_blur - clean_blur)
716 mask_arr = np.max(diff_arr, axis=2)
717 mask_binary = mask_arr > 30
718
719 white_pixels = np.sum(mask_binary)
720 print(f"📊 Mask Stats: {white_pixels} changed pixels detected.")
721 if white_pixels < 10:
722 print("⚠️ WARNING: Mask is almost empty!")
723
724 mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
725 mask = Image.fromarray((mask_filled * 255).astype(np.uint8))
726 mask = mask.filter(ImageFilter.MaxFilter(9))
727
728 # 5. Save Inputs for Inspection
729 mask.save(mask_path)
730 img_clean.save(clean_path)
731 print(f"✅ Saved debug images to: {debug_dir}/")
732
733 # 6. Run Fal.ai
734 print("🚀 Uploading images to Fal.ai...")
735 image_url = fal_client.upload_file(clean_path)
736 mask_url = fal_client.upload_file(mask_path)
737
738 print("⚡ Running Flux Dev Fill...")
739 handler = fal_client.submit(
740 "fal-ai/flux-lora-fill",
741 arguments={
742 "prompt": prompt,
743 "image_url": image_url,
744 "mask_url": mask_url,
745 "guidance_scale": 30,
746 "num_inference_steps": 28,
747 "enable_safety_checker": False,
748 },
749 )
750
751 result = handler.get()
752 print("📡 Fal Response:", result)
753
754 if "images" in result and len(result["images"]) > 0:
755 output_url = result["images"][0]["url"]
756 print(f"✨ Downloading Result: {output_url}")
757
758 response = requests.get(output_url)
759 if response.status_code == 200:
760 result_img = Image.open(io.BytesIO(response.content)).convert("RGB")
761
762 # Save Debug Output
763 result_img.save(fal_result_path)
764 print(f"💾 Saved final output to {fal_result_path}")
765
766 return jsonify(
767 {"status": "success", "image": encode_image_to_base64(result_img)}
768 )
769 else:
770 print(f"❌ Failed to download image. Status: {response.status_code}")
771 return (
772 jsonify(
773 {"status": "error", "message": "Failed to download Fal output"}
774 ),
775 500,
776 )
777 else:
778 print("❌ API returned no images.")
779 return (
780 jsonify(
781 {
782 "status": "error",
783 "message": "Fal.ai returned no images",
784 "details": result,
785 }
786 ),
787 500,
788 )
789
790 except Exception as e:
791 print(f"❌ Error in /inpainting-api: {e}")
792 traceback.print_exc()
793 return jsonify({"status": "error", "message": str(e)}), 500
794
795
796@app.route("/sketch-api", methods=["POST"])
797@secure_endpoint
798def sketch_api():
799 if not FAL_AVAILABLE:
800 return jsonify({"error": "Fal.ai client not installed or API Key missing"}), 500
801
802 try:
803 data = request.get_json()
804 prompt = data.get("prompt")
805 option = data.get("option", 1)
806
807 if not prompt:
808 return jsonify({"error": "Missing prompt"}), 400
809
810 # --- ENFORCE SHARPNESS IN PROMPT ---
811 enhanced_prompt = (
812 f"{prompt}, sharp focus, high definition, 4k, vector art, crisp lines"
813 )
814
815 if int(option) == 1:
816 # Nano Banana
817 print(f"🍌 Using Nano Banana for: {prompt}")
818 model_id = "fal-ai/nano-banana"
819 arguments = {
820 "prompt": enhanced_prompt,
821 "num_images": 1,
822 "aspect_ratio": "1:1",
823 "output_format": "png",
824 }
825 elif int(option) == 2:
826 # Flux Dev
827 print(f"🚀 Using Flux Dev for: {prompt}")
828 model_id = "fal-ai/flux/dev"
829 arguments = {
830 "image_size": "square_hd",
831 "num_inference_steps": 28,
832 "guidance_scale": 3.5,
833 "safety_tolerance": "2",
834 "enable_safety_checker": False,
835 "prompt": enhanced_prompt,
836 }
837 else:
838 return (
839 jsonify(
840 {"error": "Invalid option. Use 1 for Nano Banana, 2 for Flux Dev."}
841 ),
842 400,
843 )
844
845 # Execute request
846 handler = fal_client.submit(model_id, arguments=arguments)
847 result = handler.get()
848 print("📡 Fal Response:", result)
849
850 if "images" in result and len(result["images"]) > 0:
851 image_url = result["images"][0]["url"]
852 print(f"✨ Success! Image generated: {image_url}")
853
854 response = requests.get(image_url)
855 if response.status_code == 200:
856 img = Image.open(io.BytesIO(response.content)).convert("RGB")
857 return jsonify(
858 {"status": "success", "image": encode_image_to_base64(img)}
859 )
860 else:
861 return (
862 jsonify(
863 {
864 "status": "error",
865 "message": "Failed to download image from Fal",
866 }
867 ),
868 500,
869 )
870 else:
871 return (
872 jsonify({"status": "error", "message": "No images returned from Fal"}),
873 500,
874 )
875
876 except Exception as e:
877 print(f"❌ Error in /sketch-api: {e}")
878 traceback.print_exc()
879 return jsonify({"status": "error", "message": str(e)}), 500
880
881
882if __name__ == "__main__":
883 PORT = os.getenv("PORT")
884 if not PORT:
885 PORT = 5000
886 app.run(host="0.0.0.0", port=PORT)