commit 7bec11a0160201da9290d6b3b97fa1101c2b8870
parent 471a07f56fb160ea08c39a315b4d1605dd8a070f
Author: maydayv7 <maydayv7@gmail.com>
Date: Wed, 3 Dec 2025 23:22:35 +0530
Add flask server in-tree
Diffstat:
5 files changed, 1191 insertions(+), 0 deletions(-)
diff --git a/flask/.gitignore b/flask/.gitignore
@@ -0,0 +1,5 @@
+__pycache__
+venv
+local_inpainting_model
+Florence-2-4bit-Quantized
+BiRefNet
diff --git a/flask/README.md b/flask/README.md
@@ -0,0 +1,57 @@
+# AI Image Generator Backend
+
+This project sets up a local Flask server that uses the Stable Diffusion v1.5 model to generate images from text prompts.
+Prerequisites
+Before you begin, ensure you have the following installed:
+Python 3.10 or 3.11 (Recommended for compatibility)
+
+### Prerequisites
+
+- NVIDIA GPU (Recommended)
+ You need a GPU with at least 4GB VRAM
+ Ensure you have the latest NVIDIA drivers installed
+ Check CUDA availability: Open terminal and run `nvidia-smi`
+
+## Installation & Setup
+
+### A. Install Standard Packages
+
+```
+pip install flask flask-cors diffusers transformers accelerate safetensors
+```
+
+### B. Install PyTorch & xFormers (GPU Acceleration)
+
+This step differs slightly depending on your OS and GPU
+
+#### Windows (NVIDIA GPU)
+
+To ensure you get the version compatible with your GPU, run:
+
+```
+pip install torch torchvision xformers --index-url https://download.pytorch.org/whl/cu118
+```
+
+If you have a very new GPU, you can try `cu121` instead of `cu118`
+
+#### Mac (M1/M2/M3 Silicon)
+
+Macs use "MPS" (Metal Performance Shaders) instead of CUDA
+
+```
+pip install torch torchvision torchaudio
+```
+
+### C. Models
+
+Download the requisite models from [here](https://github.com/adobeinter/Adobe-Models/tree/main)
+
+## Running the Server
+
+```shell
+python index.py
+```
+
+# Server Deployment
+
+To deploy to [Modal](https://modal.com/) server, simply modify `modal_app.py` according to your requirements and run `modal deploy modal_app.py`
diff --git a/flask/index.py b/flask/index.py
@@ -0,0 +1,613 @@
+import os
+import io
+import sys
+import base64
+import cv2
+import re
+import numpy as np
+import torch
+import traceback
+import requests
+import scipy.ndimage
+from flask import Flask, jsonify, request, send_file
+from flask_cors import CORS
+from PIL import Image, ImageFilter
+from torchvision import transforms
+import time
+import uuid
+
+# --- FAL.AI IMPORTS ---
+try:
+ import fal_client
+ # 1. SETUP API KEY
+ os.environ["FAL_KEY"] = "f040803a-2cc4-4210-86f2-53b4a0e33354:335fb1972606d25f80004bd3bd11d935"
+ FAL_AVAILABLE = True
+except ImportError:
+ print("⚠️ Fal.ai Client not installed. /inpainting-api will fail.")
+ FAL_AVAILABLE = False
+
+# --- FLORENCE-2 IMPORTS ---
+try:
+ import bitsandbytes
+ from transformers import AutoModelForCausalLM, AutoProcessor, BitsAndBytesConfig
+ import transformers.dynamic_module_utils
+ import torch.nn as nn
+ FLORENCE_AVAILABLE = True
+except ImportError as e:
+ print(f"⚠️ Florence-2 Disabled: {e} (Ensure 'bitsandbytes' and 'transformers' are installed)")
+ FLORENCE_AVAILABLE = False
+except Exception as e:
+ print(f"⚠️ Florence-2 Disabled: Unexpected initialization error: {e}")
+ FLORENCE_AVAILABLE = False
+
+# --- IMPORT BIREFNET ---
+current_dir = os.path.dirname(os.path.abspath(__file__))
+birefnet_path = os.path.join(current_dir, "BiRefNet")
+
+if os.path.exists(birefnet_path):
+ if birefnet_path not in sys.path:
+ sys.path.append(birefnet_path)
+ print(f"✅ Added {birefnet_path} to system path.")
+else:
+ print(f"❌ Error: '{birefnet_path}' not found. Please clone the repository.")
+
+try:
+ from models.birefnet import BiRefNet
+ print("✅ BiRefNet imported successfully.")
+except ImportError as e:
+ print(f"⚠️ Import Error: {e}")
+ try:
+ import BiRefNet.models.birefnet as brn
+ BiRefNet = brn.BiRefNet
+ print("✅ BiRefNet imported via package path.")
+ except ImportError:
+ print("❌ Failed to import BiRefNet. Ensure 'BiRefNet/models/birefnet.py' exists.")
+
+# --- IMPORT STABLE DIFFUSION ---
+try:
+ from diffusers import StableDiffusionInpaintPipeline, AutoPipelineForInpainting
+except ImportError:
+ print("⚠️ Diffusers not found. SD features disabled.")
+ StableDiffusionInpaintPipeline = None
+ AutoPipelineForInpainting = None
+
+app = Flask(__name__)
+CORS(app)
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"🚀 Running on device: {DEVICE}")
+
+# ==============================================================================
+# 1. LOAD STABLE DIFFUSION
+# ==============================================================================
+print("⏳ Loading Stable Diffusion (Inpainting)...")
+sd_pipe = None
+try:
+ if StableDiffusionInpaintPipeline:
+ SD_MODEL_ID = "./local_inpainting_model" if os.path.exists("./local_inpainting_model") else "runwayml/stable-diffusion-inpainting"
+
+ sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+ SD_MODEL_ID,
+ torch_dtype=torch.float16 if DEVICE == 'cuda' else torch.float32,
+ use_safetensors=True,
+ ).to(DEVICE)
+ sd_pipe.enable_attention_slicing()
+ sd_pipe.enable_model_cpu_offload()
+ print("✅ Stable Diffusion Loaded!")
+except Exception as e:
+ print(f"❌ Failed to load SD: {e}")
+
+# ==============================================================================
+# 2. LOAD BIREFNET
+# ==============================================================================
+print("⏳ Loading BiRefNet...")
+birefnet_model = None
+BIREFNET_WEIGHTS = "./BiRefNet/birefnet_fp16.pt"
+BIREFNET_SIZE = (1024, 1024)
+
+try:
+ if 'BiRefNet' in locals() and os.path.exists(BIREFNET_WEIGHTS):
+ birefnet_model = BiRefNet(bb_pretrained=False)
+ state_dict = torch.load(BIREFNET_WEIGHTS, map_location=DEVICE)
+ birefnet_model.load_state_dict(state_dict)
+ birefnet_model.to(DEVICE)
+ if DEVICE == 'cuda':
+ birefnet_model.half()
+ birefnet_model.eval()
+ print("✅ BiRefNet Weights Loaded!")
+ else:
+ print(f"⚠️ BiRefNet skipped. Weights found: {os.path.exists(BIREFNET_WEIGHTS)}")
+except Exception as e:
+ print(f"❌ Failed to load BiRefNet: {e}")
+
+transform_birefnet = transforms.Compose([
+ transforms.Resize(BIREFNET_SIZE),
+ transforms.ToTensor(),
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+# ==============================================================================
+# 3. LOAD FLORENCE-2 (QUANTIZED)
+# ==============================================================================
+print("⏳ Loading Florence-2...")
+florence_model = None
+florence_processor = None
+FLORENCE_PATH = os.path.join(current_dir, "Florence-2-4bit-Quantized")
+
+if FLORENCE_AVAILABLE:
+ try:
+ def check_imports_fixed(filename): return []
+ transformers.dynamic_module_utils.check_imports = check_imports_fixed
+
+ _old_getattr = nn.Module.__getattr__
+ def _fixed_getattr(self, name):
+ if name == "_supports_sdpa":
+ return False
+ return _old_getattr(self, name)
+ nn.Module.__getattr__ = _fixed_getattr
+
+ if os.path.exists(FLORENCE_PATH):
+ bnb_config = BitsAndBytesConfig(
+ load_in_4bit=True,
+ bnb_4bit_quant_type="nf4",
+ bnb_4bit_compute_dtype=torch.float16
+ )
+
+ florence_model = AutoModelForCausalLM.from_pretrained(
+ FLORENCE_PATH,
+ quantization_config=bnb_config,
+ trust_remote_code=True,
+ device_map="cuda" if DEVICE == 'cuda' else 'cpu',
+ local_files_only=True
+ )
+ florence_processor = AutoProcessor.from_pretrained(FLORENCE_PATH, trust_remote_code=True)
+ print("✅ Florence-2 Loaded Successfully!")
+ else:
+ print(f"⚠️ Florence-2 folder not found at: {FLORENCE_PATH}")
+ except Exception as e:
+ print(f"❌ Failed to load Florence-2: {e}")
+
+# ==============================================================================
+# HELPER FUNCTIONS
+# ==============================================================================
+def decode_base64_image(b64_str):
+ if "," in b64_str:
+ b64_str = b64_str.split(",")[1]
+ image_data = base64.b64decode(b64_str)
+ img = Image.open(io.BytesIO(image_data))
+
+ if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
+ background = Image.new('RGB', img.size, (255, 255, 255))
+ if img.mode == 'P': img = img.convert('RGBA')
+ background.paste(img, mask=img.split()[3])
+ return background
+ else:
+ return img.convert("RGB")
+
+def encode_image_to_base64(pil_img):
+ buffered = io.BytesIO()
+ pil_img.save(buffered, format="PNG")
+ return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+def process_birefnet_output(preds, original_size):
+ if isinstance(preds, (list, tuple)):
+ pred_tensor = preds[-1]
+ else:
+ pred_tensor = preds
+
+ pred_tensor = pred_tensor.sigmoid().cpu()
+ mask_np = pred_tensor.squeeze().numpy().astype(np.float32)
+
+ if len(mask_np.shape) > 2:
+ mask_np = mask_np[0]
+
+ mask_resized = cv2.resize(mask_np, original_size, interpolation=cv2.INTER_LINEAR)
+ mask = (mask_resized > 0.5).astype(np.uint8) * 255
+
+ return Image.fromarray(mask)
+
+def resize_to_limit(img, max_dim=1024, multiple=8):
+ w, h = img.size
+ ratio = min(max_dim / w, max_dim / h)
+ new_w = int(w * ratio)
+ new_h = int(h * ratio)
+ new_w = new_w - (new_w % multiple)
+ new_h = new_h - (new_h % multiple)
+ if new_w < multiple: new_w = multiple
+ if new_h < multiple: new_h = multiple
+ return img.resize((new_w, new_h), Image.LANCZOS)
+
+# ==============================================================================
+# ROUTES
+# ==============================================================================
+@app.route('/')
+def index():
+ return "Image Processing API is running."
+
+@app.route('/generate', methods=['POST'])
+def generate_image():
+ if not sd_pipe: return jsonify({"error": "SD Model not loaded"}), 500
+ try:
+ data = request.get_json()
+ prompt = data.get('prompt', '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.')
+ empty_image = Image.new("RGB", (512, 512), (0, 0, 0))
+ full_mask = Image.new("L", (512, 512), 255)
+ print(f"🎨 Generating: {prompt}")
+ image = sd_pipe(prompt=prompt, image=empty_image, mask_image=full_mask, height=512, width=512, num_inference_steps=30).images[0]
+ return jsonify({"status": "success", "image": encode_image_to_base64(image)})
+ except Exception as e:
+ return jsonify({"status": "error", "message": str(e)}), 500
+
+@app.route('/inpainting', methods=['POST'])
+def inpaint_image():
+ if not sd_pipe: return jsonify({"error": "SD Model not loaded"}), 500
+ try:
+ data = request.get_json()
+ user_prompt = data.get('prompt', '')
+ clean_b64 = data.get('image')
+ drawn_b64 = data.get('mask_image')
+
+ if not clean_b64 or not drawn_b64: return jsonify({"error": "Missing image or mask"}), 400
+
+ # 1. Decode Images
+ raw_clean = decode_base64_image(clean_b64).convert("RGB")
+ raw_drawn = decode_base64_image(drawn_b64).convert("RGB")
+
+ # 2. Resize maintaining Aspect Ratio (Max 512 for Local SD)
+ img_clean = resize_to_limit(raw_clean, max_dim=512)
+ # Resize drawn image to match the clean image exactly
+ img_drawn = raw_drawn.resize(img_clean.size)
+
+ print(f"🔍 Calculating Robust Difference Mask (Size: {img_clean.size})...")
+ clean_blur = np.array(img_clean.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16)
+ drawn_blur = np.array(img_drawn.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16)
+ diff_arr = np.abs(drawn_blur - clean_blur)
+ mask_arr = np.max(diff_arr, axis=2)
+ mask_binary = mask_arr > 30
+ mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
+ mask_image = Image.fromarray((mask_filled * 255).astype(np.uint8))
+ mask_image = mask_image.filter(ImageFilter.MaxFilter(9))
+ print("✅ Mask calculated.")
+
+ generated_prompt = ""
+ if florence_model and florence_processor:
+ print("👁️ Generating context with Florence-2...")
+ try:
+ task_prompt = '<DETAILED_CAPTION>'
+ inputs = florence_processor(text=task_prompt, images=[img_drawn], return_tensors="pt")
+ inputs["pixel_values"] = inputs["pixel_values"].to(DEVICE, torch.float16)
+ inputs["input_ids"] = inputs["input_ids"].to(DEVICE)
+
+ generated_ids = florence_model.generate(
+ input_ids=inputs["input_ids"],
+ pixel_values=inputs["pixel_values"],
+ max_new_tokens=128,
+ num_beams=1,
+ do_sample=False,
+ use_cache=False
+ )
+
+ generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+ generated_prompt = generated_text.replace(task_prompt, "").replace("</s>", "").replace("<s>", "").strip()
+ print(f"📝 Florence Generated: {generated_prompt}")
+ except Exception as e:
+ print(f"⚠️ Florence captioning failed: {e}")
+
+ final_prompt = f"{generated_prompt} {user_prompt}".strip()
+ negative_prompt = "blurry, low quality, ugly, text, watermark, bad anatomy, deformed, noisy"
+ print(f"✨ Final Inpaint Prompt: {final_prompt}")
+
+ save_dir = "input_data"
+ os.makedirs(save_dir, exist_ok=True)
+ timestamp = int(time.time())
+ img_clean.save(os.path.join(save_dir, f"clean_{timestamp}.png"))
+ img_drawn.save(os.path.join(save_dir, f"drawn_{timestamp}.png"))
+ mask_image.save(os.path.join(save_dir, f"generated_mask_{timestamp}.png"))
+
+ print(f"🎨 Running Inference with strength=0.85...")
+ image = sd_pipe(
+ prompt=final_prompt,
+ negative_prompt=negative_prompt,
+ image=img_drawn,
+ mask_image=mask_image,
+ num_inference_steps=50,
+ strength=0.85,
+ guidance_scale=8.5
+ ).images[0]
+
+ final_image_path = os.path.join(save_dir, f"result_{timestamp}.png")
+ image.save(final_image_path)
+ print(f"💾 Saved output to {final_image_path}")
+
+ return jsonify({"status": "success", "image": encode_image_to_base64(image)})
+
+ except Exception as e:
+ print(f"❌ Inpainting Error: {e}")
+ traceback.print_exc()
+ return jsonify({"status": "error", "message": str(e)}), 500
+
+@app.route('/asset', methods=['POST'])
+def remove_background():
+ if not birefnet_model:
+ return jsonify({"error": "BiRefNet not loaded"}), 500
+ try:
+ data = request.get_json()
+ image_b64 = data.get('image')
+ if not image_b64: return jsonify({"error": "No image provided"}), 400
+
+ original_image = decode_base64_image(image_b64)
+ orig_w, orig_h = original_image.size
+
+ input_tensor = transform_birefnet(original_image).unsqueeze(0).to(DEVICE)
+ if DEVICE == 'cuda':
+ input_tensor = input_tensor.half()
+
+ print("✂️ Removing background...")
+ with torch.no_grad():
+ preds = birefnet_model(input_tensor)
+
+ mask_pil = process_birefnet_output(preds, (orig_w, orig_h))
+ original_image.putalpha(mask_pil)
+
+ return jsonify({"status": "success", "image": encode_image_to_base64(original_image)})
+ except Exception as e:
+ print(f"❌ Error: {e}")
+ traceback.print_exc()
+ return jsonify({"status": "error", "message": str(e)}), 500
+
+@app.route('/describe', methods=['POST'])
+def describe_image():
+ if not florence_model or not florence_processor:
+ return jsonify({"error": "Florence-2 not loaded"}), 500
+ try:
+ data = request.get_json()
+ image_b64 = data.get('image')
+ prompt_type = data.get('prompt', '<DETAILED_CAPTION>')
+
+ if not image_b64: return jsonify({"error": "No image provided"}), 400
+
+ image = decode_base64_image(image_b64)
+ print(f"👁️ Analyzing image with Florence-2...")
+
+ inputs = florence_processor(text=prompt_type, images=[image], return_tensors="pt")
+ inputs["pixel_values"] = inputs["pixel_values"].to(DEVICE, torch.float16)
+ inputs["input_ids"] = inputs["input_ids"].to(DEVICE)
+
+ generated_ids = florence_model.generate(
+ input_ids=inputs["input_ids"],
+ pixel_values=inputs["pixel_values"],
+ max_new_tokens=128,
+ num_beams=1,
+ do_sample=False,
+ use_cache=False
+ )
+
+ generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+ cleaned_text = generated_text.replace(prompt_type, "").replace("</s>", "").replace("<s>", "").strip()
+
+ if cleaned_text and cleaned_text[-1] not in ['.', '!', '?']:
+ last_dot = cleaned_text.rfind('.')
+ last_excl = cleaned_text.rfind('!')
+ last_ques = cleaned_text.rfind('?')
+ cut_off = max(last_dot, last_excl, last_ques)
+ if cut_off != -1:
+ cleaned_text = cleaned_text[:cut_off+1]
+
+ final_answer = cleaned_text
+ print(final_answer)
+
+ if "<loc_" in cleaned_text or "<poly_" in cleaned_text:
+ try:
+ parsed = florence_processor.post_process_generation(
+ generated_text,
+ task=prompt_type,
+ image_size=(image.width, image.height)
+ )
+ if isinstance(parsed, dict) and prompt_type in parsed:
+ final_answer = parsed[prompt_type]
+ else:
+ final_answer = parsed
+ except Exception:
+ def parse_loc_manually(text, w, h):
+ locs = re.findall(r'<loc_(\d+)>', text)
+ if locs and len(locs) % 4 == 0:
+ bboxes = []
+ for i in range(0, len(locs), 4):
+ x1 = int(int(locs[i]) / 1000 * w)
+ y1 = int(int(locs[i+1]) / 1000 * h)
+ x2 = int(int(locs[i+2]) / 1000 * w)
+ y2 = int(int(locs[i+3]) / 1000 * h)
+ bboxes.append([x1, y1, x2, y2])
+ clean_text = re.sub(r'<loc_\d+>', '', text).strip()
+ return {"text": clean_text, "bboxes": bboxes}
+ return text
+ final_answer = parse_loc_manually(cleaned_text, image.width, image.height)
+
+ return jsonify({"status": "success", "output": final_answer})
+
+ except Exception as e:
+ print(f"❌ Florence Error: {e}")
+ traceback.print_exc()
+ return jsonify({"status": "error", "message": str(e)}), 500
+
+@app.route('/inpainting-api', methods=['POST'])
+def inpainting_api_fal():
+ if not FAL_AVAILABLE:
+ return jsonify({"error": "Fal.ai client not installed or API Key missing"}), 500
+
+ try:
+ data = request.get_json()
+
+ clean_b64 = data.get('image')
+ drawn_b64 = data.get('mask_image')
+ prompt = data.get('prompt', "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.")
+
+ if not clean_b64 or not drawn_b64:
+ return jsonify({"error": "Missing 'image' (clean) or 'mask_image' (drawn)"}), 400
+
+ print(f"📥 Received Request: Prompt='{prompt}'")
+
+ # 1. Decode Images
+ raw_clean = decode_base64_image(clean_b64).convert("RGB")
+ raw_drawn = decode_base64_image(drawn_b64).convert("RGB")
+
+ # 2. Resize maintaining Aspect Ratio (Max 1024 for Flux)
+ img_clean = resize_to_limit(raw_clean, max_dim=1024)
+ # Resize drawn to match exactly
+ img_drawn = raw_drawn.resize(img_clean.size)
+
+ # 3. Setup Debug Directory
+ debug_dir = "debug_fal"
+ os.makedirs(debug_dir, exist_ok=True)
+ unique_id = str(int(time.time()))
+
+ clean_path = os.path.join(debug_dir, f"fal_clean_{unique_id}.png")
+ mask_path = os.path.join(debug_dir, f"fal_mask_{unique_id}.png")
+ fal_result_path = os.path.join(debug_dir, f"fal_result_{unique_id}.png")
+
+ # 4. Mask Generation
+ print(f"🛠️ Generating mask (Size: {img_clean.size})...")
+ clean_blur = np.array(img_clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16)
+ drawn_blur = np.array(img_drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16)
+
+ diff_arr = np.abs(drawn_blur - clean_blur)
+ mask_arr = np.max(diff_arr, axis=2)
+ mask_binary = mask_arr > 30
+
+ white_pixels = np.sum(mask_binary)
+ print(f"📊 Mask Stats: {white_pixels} changed pixels detected.")
+ if white_pixels < 10:
+ print("⚠️ WARNING: Mask is almost empty!")
+
+ mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
+ mask = Image.fromarray((mask_filled * 255).astype(np.uint8))
+ mask = mask.filter(ImageFilter.MaxFilter(9))
+
+ # 5. Save Inputs for Inspection
+ mask.save(mask_path)
+ img_clean.save(clean_path)
+ print(f"✅ Saved debug images to: {debug_dir}/")
+
+ # 6. Run Fal.ai
+ print("🚀 Uploading images to Fal.ai...")
+ image_url = fal_client.upload_file(clean_path)
+ mask_url = fal_client.upload_file(mask_path)
+
+ print("⚡ Running Flux Dev Fill...")
+ handler = fal_client.submit(
+ "fal-ai/flux-lora-fill",
+ arguments={
+ "prompt": prompt,
+ "image_url": image_url,
+ "mask_url": mask_url,
+ "guidance_scale": 30,
+ "num_inference_steps": 28,
+ "enable_safety_checker": False
+ }
+ )
+
+ result = handler.get()
+ print("📡 Fal Response:", result)
+
+ if 'images' in result and len(result['images']) > 0:
+ output_url = result['images'][0]['url']
+ print(f"✨ Downloading Result: {output_url}")
+
+ response = requests.get(output_url)
+ if response.status_code == 200:
+ result_img = Image.open(io.BytesIO(response.content)).convert("RGB")
+
+ # Save Debug Output
+ result_img.save(fal_result_path)
+ print(f"💾 Saved final output to {fal_result_path}")
+
+ return jsonify({"status": "success", "image": encode_image_to_base64(result_img)})
+ else:
+ print(f"❌ Failed to download image. Status: {response.status_code}")
+ return jsonify({"status": "error", "message": "Failed to download Fal output"}), 500
+ else:
+ print("❌ API returned no images.")
+ return jsonify({"status": "error", "message": "Fal.ai returned no images", "details": result}), 500
+
+ except Exception as e:
+ print(f"❌ Error in /inpainting-api: {e}")
+ traceback.print_exc()
+ return jsonify({"status": "error", "message": str(e)}), 500
+
+# ==============================================================================
+# 5. NEW ROUTE: SKETCH API (Text-to-Image)
+# ==============================================================================
+@app.route('/sketch-api', methods=['POST'])
+def sketch_api():
+ if not FAL_AVAILABLE:
+ return jsonify({"error": "Fal.ai client not installed or API Key missing"}), 500
+
+ try:
+ data = request.get_json()
+ prompt = data.get('prompt')
+ option = data.get('option', 1)
+
+ if not prompt:
+ return jsonify({"error": "Missing prompt"}), 400
+
+ # --- ENFORCE SHARPNESS IN PROMPT ---
+ # We append these words to ensure the "cartoon" style isn't blurry
+ enhanced_prompt = f"{prompt}, sharp focus, high definition, 4k, vector art, crisp lines"
+
+ if int(option) == 1:
+ # Flux Schnell (No change needed, it ignores guidance)
+ print(f"🚀 Using Flux Schnell for: {prompt}")
+ model_id = "fal-ai/flux/schnell"
+ arguments = {
+ "image_size": "square_hd",
+ "num_inference_steps": 4,
+ "enable_safety_checker": False,
+ "prompt": enhanced_prompt
+ }
+ elif int(option) == 2:
+ # Flux Dev (THIS IS THE FIX)
+ print(f"🚀 Using Flux Dev for: {prompt}")
+ model_id = "fal-ai/flux/dev"
+ arguments = {
+ "image_size": "square_hd",
+ "num_inference_steps": 28,
+
+ # 👇 KEY FIX: Set this to 3.5 or 4.0 for sharp images
+ "guidance_scale": 3.5,
+
+ # 👇 OPTIONAL: Makes it less likely to return black images
+ "safety_tolerance": "2",
+
+ "enable_safety_checker": False,
+ "prompt": enhanced_prompt
+ }
+ else:
+ return jsonify({"error": "Invalid option. Use 1 for Schnell, 2 for Dev."}), 400
+
+ handler = fal_client.submit(
+ model_id,
+ arguments=arguments
+ )
+ result = handler.get()
+ print("📡 Fal Response:", result)
+
+ if 'images' in result and len(result['images']) > 0:
+ image_url = result['images'][0]['url']
+ print(f"✨ Success! Image generated: {image_url}")
+
+ response = requests.get(image_url)
+ if response.status_code == 200:
+ img = Image.open(io.BytesIO(response.content)).convert("RGB")
+ return jsonify({"status": "success", "image": encode_image_to_base64(img)})
+ else:
+ return jsonify({"status": "error", "message": "Failed to download image from Fal"}), 500
+ else:
+ return jsonify({"status": "error", "message": "No images returned from Fal"}), 500
+
+ except Exception as e:
+ print(f"❌ Error in /sketch-api: {e}")
+ traceback.print_exc()
+ return jsonify({"status": "error", "message": str(e)}), 500
+
+if __name__ == "__main__":
+ app.run(host='0.0.0.0', port=5000)
+\ No newline at end of file
diff --git a/flask/modal_app.py b/flask/modal_app.py
@@ -0,0 +1,495 @@
+import os
+import io
+import sys
+import base64
+import modal
+
+# ==============================================================================
+# 1. DEFINE THE CLOUD ENVIRONMENT
+# ==============================================================================
+image = (
+ modal.Image.debian_slim(python_version="3.11")
+ .apt_install(
+ "libgl1-mesa-glx",
+ "libglib2.0-0",
+ "libstdc++6",
+ "libxext6",
+ "libsm6",
+ "libxrender1",
+ )
+ .pip_install(
+ "torch",
+ "torchvision",
+ "transformers",
+ "diffusers",
+ "accelerate",
+ "safetensors",
+ "opencv-python-headless",
+ "pillow",
+ "numpy",
+ "scipy",
+ "bitsandbytes",
+ "timm",
+ "einops",
+ "kornia",
+ "flask-cors",
+ "fastapi[standard]",
+ "fal-client",
+ "requests",
+ )
+ .env(
+ {
+ "FAL_KEY": "f040803a-2cc4-4210-86f2-53b4a0e33354:335fb1972606d25f80004bd3bd11d935"
+ }
+ )
+ # --- MOUNT LOCAL MODELS ---
+ .add_local_dir("local_inpainting_model", remote_path="/models/sd-inpainting")
+ .add_local_dir("Florence-2-4bit-Quantized", remote_path="/models/florence-2")
+ .add_local_dir("BiRefNet", remote_path="/root/BiRefNet")
+)
+
+app = modal.App("adobe-flask", image=image)
+
+
+# ==============================================================================
+# 2. THE BACKEND SERVER CLASS
+# ==============================================================================
+@app.cls(gpu="any", scaledown_window=300)
+class ModelBackend:
+
+ @modal.enter()
+ def load_models(self):
+ """Runs once when container starts."""
+ print("⏳ Loading models into GPU memory...")
+ import torch
+ import torch.nn as nn
+ import sys
+
+ self.device = "cuda"
+
+ # --- 1. SETUP BiRefNet PATHS ---
+ sys.path.append("/root/BiRefNet")
+
+ # --- 2. LOAD STABLE DIFFUSION ---
+ from diffusers import StableDiffusionInpaintPipeline
+
+ self.sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+ "/models/sd-inpainting",
+ torch_dtype=torch.float16,
+ use_safetensors=True,
+ local_files_only=True,
+ ).to(self.device)
+ self.sd_pipe.enable_attention_slicing()
+ print("✅ Stable Diffusion Loaded")
+
+ # --- 3. LOAD FLORENCE-2 ---
+ import transformers.dynamic_module_utils
+
+ # Patch 1: Fix import check
+ def check_imports_fixed(filename):
+ return []
+
+ transformers.dynamic_module_utils.check_imports = check_imports_fixed
+
+ # Patch 2: Fix '_supports_sdpa' error
+ _old_getattr = nn.Module.__getattr__
+
+ def _fixed_getattr(self, name):
+ if name == "_supports_sdpa":
+ return False
+ return _old_getattr(self, name)
+
+ nn.Module.__getattr__ = _fixed_getattr
+
+ from transformers import AutoModelForCausalLM, AutoProcessor, BitsAndBytesConfig
+
+ bnb_config = BitsAndBytesConfig(
+ load_in_4bit=True,
+ bnb_4bit_quant_type="nf4",
+ bnb_4bit_compute_dtype=torch.float16,
+ )
+ self.florence_model = AutoModelForCausalLM.from_pretrained(
+ "/models/florence-2",
+ quantization_config=bnb_config,
+ trust_remote_code=True,
+ local_files_only=True,
+ ).to(self.device)
+ self.florence_processor = AutoProcessor.from_pretrained(
+ "/models/florence-2", trust_remote_code=True, local_files_only=True
+ )
+ print("✅ Florence-2 Loaded")
+
+ # --- 4. LOAD BIREFNET ---
+ try:
+ from models.birefnet import BiRefNet
+
+ self.birefnet = BiRefNet(bb_pretrained=False)
+
+ weight_path = "/root/BiRefNet/birefnet_fp16.pt"
+ state_dict = torch.load(weight_path, map_location=self.device)
+ self.birefnet.load_state_dict(state_dict)
+ self.birefnet.to(self.device).half().eval()
+ print(f"✅ BiRefNet Loaded from {weight_path}")
+ except Exception as e:
+ print(f"❌ BiRefNet Error: {e}")
+ self.birefnet = None
+
+ from torchvision import transforms
+
+ self.transform_birefnet = transforms.Compose(
+ [
+ transforms.Resize((1024, 1024)),
+ transforms.ToTensor(),
+ transforms.Normalize(
+ mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+ ),
+ ]
+ )
+
+ # ==========================================================================
+ # 3. ENDPOINTS
+ # ==========================================================================
+
+ @modal.fastapi_endpoint(method="POST")
+ def generate(self, item: dict):
+ from PIL import Image
+
+ prompt = item.get("prompt", "A luxury watch")
+ print(f"🎨 Generating: {prompt}")
+
+ empty_image = Image.new("RGB", (512, 512), (0, 0, 0))
+ full_mask = Image.new("L", (512, 512), 255)
+
+ image = self.sd_pipe(
+ prompt=prompt,
+ image=empty_image,
+ mask_image=full_mask,
+ height=512,
+ width=512,
+ num_inference_steps=30,
+ ).images[0]
+
+ return {"status": "success", "image": self._to_base64(image)}
+
+ @modal.fastapi_endpoint(method="POST")
+ def inpainting(self, item: dict):
+ """Local Stable Diffusion Inpainting"""
+ from PIL import Image, ImageFilter
+ import numpy as np
+ import scipy.ndimage
+ import torch
+
+ user_prompt = item.get("prompt", "")
+ img_b64 = item.get("image")
+ mask_b64 = item.get("mask_image")
+
+ if not img_b64 or not mask_b64:
+ return {"status": "error", "message": "Missing image or mask"}
+
+ # 1. Decode Images
+ raw_clean = self._decode_base64(img_b64).convert("RGB")
+ raw_drawn = self._decode_base64(mask_b64).convert("RGB")
+
+ # 2. Resize maintaining Aspect Ratio (Max 512 for Local SD)
+ img_clean = self._resize_to_limit(raw_clean, max_dim=512)
+ # Resize drawn image to match the clean image exactly
+ img_drawn = raw_drawn.resize(img_clean.size)
+
+ print(f"🔍 Calculating Robust Difference Mask (Size: {img_clean.size})...")
+
+ # --- ROBUST MASKING ---
+ clean_blur = np.array(
+ img_clean.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16
+ )
+ drawn_blur = np.array(
+ img_drawn.filter(ImageFilter.GaussianBlur(radius=2)), dtype=np.int16
+ )
+
+ diff_arr = np.abs(drawn_blur - clean_blur)
+ mask_arr = np.max(diff_arr, axis=2)
+ mask_binary = mask_arr > 30
+
+ mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
+ mask_image = Image.fromarray((mask_filled * 255).astype(np.uint8))
+ mask_image = mask_image.filter(ImageFilter.MaxFilter(9))
+ print("✅ Mask calculated.")
+
+ # --- FLORENCE-2 CONTEXT GENERATION ---
+ generated_prompt = ""
+ if self.florence_model and self.florence_processor:
+ print("👁️ Generating context with Florence-2...")
+ try:
+ task_prompt = "<DETAILED_CAPTION>"
+ # Use img_drawn (sketch) for context analysis
+ inputs = self.florence_processor(
+ text=task_prompt, images=[img_drawn], return_tensors="pt"
+ )
+ inputs["pixel_values"] = inputs["pixel_values"].to(
+ self.device, torch.float16
+ )
+ inputs["input_ids"] = inputs["input_ids"].to(self.device)
+
+ generated_ids = self.florence_model.generate(
+ input_ids=inputs["input_ids"],
+ pixel_values=inputs["pixel_values"],
+ max_new_tokens=128,
+ num_beams=1,
+ do_sample=False,
+ use_cache=False, # Fix for transformers crash
+ )
+
+ generated_text = self.florence_processor.batch_decode(
+ generated_ids, skip_special_tokens=False
+ )[0]
+ generated_prompt = (
+ generated_text.replace(task_prompt, "")
+ .replace("</s>", "")
+ .replace("<s>", "")
+ .strip()
+ )
+ print(f"📝 Florence Generated: {generated_prompt}")
+ except Exception as e:
+ print(f"⚠️ Florence captioning failed: {e}")
+
+ final_prompt = f"{generated_prompt} {user_prompt}".strip()
+ negative_prompt = (
+ "blurry, low quality, ugly, text, watermark, bad anatomy, deformed, noisy"
+ )
+
+ # --- INFERENCE ---
+ print(f"🎨 Running Inference: {final_prompt}")
+ output = self.sd_pipe(
+ prompt=final_prompt,
+ negative_prompt=negative_prompt,
+ image=img_drawn, # Input is the SKETCH
+ mask_image=mask_image, # Mask is where sketch differs
+ num_inference_steps=50,
+ strength=0.85,
+ guidance_scale=8.5,
+ ).images[0]
+
+ return {"status": "success", "image": self._to_base64(output)}
+
+ @modal.fastapi_endpoint(method="POST")
+ def inpainting_api(self, item: dict):
+ """Fal.ai Flux Lora Fill"""
+ import fal_client
+ import requests
+ import uuid
+ import numpy as np
+ import scipy.ndimage
+ from PIL import Image, ImageFilter
+
+ prompt = item.get("prompt", "A high quality image")
+ img_b64 = item.get("image")
+ mask_b64 = item.get("mask_image")
+
+ if not img_b64 or not mask_b64:
+ return {"status": "error", "message": "Missing image or mask"}
+
+ # 1. Decode & Resize (Flux supports higher res)
+ raw_clean = self._decode_base64(img_b64).convert("RGB")
+ raw_drawn = self._decode_base64(mask_b64).convert("RGB")
+
+ img_clean = self._resize_to_limit(raw_clean, max_dim=1024)
+ img_drawn = raw_drawn.resize(img_clean.size)
+
+ # 2. Robust Mask Generation
+ clean_blur = np.array(
+ img_clean.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
+ )
+ drawn_blur = np.array(
+ img_drawn.filter(ImageFilter.GaussianBlur(2)), dtype=np.int16
+ )
+
+ diff_arr = np.abs(drawn_blur - clean_blur)
+ mask_arr = np.max(diff_arr, axis=2)
+ mask_binary = mask_arr > 30
+
+ mask_filled = scipy.ndimage.binary_fill_holes(mask_binary)
+ mask = Image.fromarray((mask_filled * 255).astype(np.uint8))
+ mask = mask.filter(ImageFilter.MaxFilter(9))
+
+ # 3. Save to temp files for upload
+ temp_id = str(uuid.uuid4())
+ clean_path = f"/tmp/clean_{temp_id}.png"
+ mask_path = f"/tmp/mask_{temp_id}.png"
+
+ img_clean.save(clean_path)
+ mask.save(mask_path)
+
+ try:
+ print("🚀 Uploading to Fal.ai...")
+ image_url = fal_client.upload_file(clean_path)
+ mask_url = fal_client.upload_file(mask_path)
+
+ print(f"⚡ Running Flux Dev Fill for: {prompt}")
+ handler = fal_client.submit(
+ "fal-ai/flux-lora-fill",
+ arguments={
+ "prompt": prompt,
+ "image_url": image_url,
+ "mask_url": mask_url,
+ "guidance_scale": 30,
+ "num_inference_steps": 28,
+ "enable_safety_checker": False,
+ },
+ )
+ result = handler.get()
+
+ if "images" in result:
+ output_url = result["images"][0]["url"]
+ response = requests.get(output_url)
+ result_img = Image.open(io.BytesIO(response.content))
+ return {"status": "success", "image": self._to_base64(result_img)}
+ else:
+ return {"status": "error", "message": "Fal.ai returned no images"}
+
+ except Exception as e:
+ print(f"❌ Fal.ai Error: {e}")
+ return {"status": "error", "message": str(e)}
+ finally:
+ if os.path.exists(clean_path):
+ os.remove(clean_path)
+ if os.path.exists(mask_path):
+ os.remove(mask_path)
+
+ @modal.fastapi_endpoint(method="POST")
+ def sketch_api(self, item: dict):
+ """Sketch Text-to-Image (Flux)"""
+ import fal_client
+ import requests
+ from PIL import Image
+
+ prompt = item.get("prompt")
+ option = item.get("option", 1)
+
+ if not prompt:
+ return {"status": "error", "message": "Missing prompt"}
+
+ enhanced_prompt = (
+ f"{prompt}, sharp focus, high definition, 4k, vector art, crisp lines"
+ )
+
+ if int(option) == 1:
+ model_id = "fal-ai/flux/schnell"
+ arguments = {
+ "image_size": "square_hd",
+ "num_inference_steps": 4,
+ "enable_safety_checker": False,
+ "prompt": enhanced_prompt,
+ }
+ else:
+ model_id = "fal-ai/flux/dev"
+ arguments = {
+ "image_size": "square_hd",
+ "num_inference_steps": 28,
+ "guidance_scale": 3.5,
+ "safety_tolerance": "2",
+ "enable_safety_checker": False,
+ "prompt": enhanced_prompt,
+ }
+
+ try:
+ print(f"🚀 Running Sketch Gen ({model_id})...")
+ handler = fal_client.submit(model_id, arguments=arguments)
+ result = handler.get()
+
+ if "images" in result and len(result["images"]) > 0:
+ image_url = result["images"][0]["url"]
+ response = requests.get(image_url)
+ if response.status_code == 200:
+ img = Image.open(io.BytesIO(response.content)).convert("RGB")
+ return {"status": "success", "image": self._to_base64(img)}
+
+ return {"status": "error", "message": "Fal.ai returned no images"}
+ except Exception as e:
+ print(f"❌ Sketch API Error: {e}")
+ return {"status": "error", "message": str(e)}
+
+ @modal.fastapi_endpoint(method="POST")
+ def asset(self, item: dict):
+ import torch
+ import numpy as np
+ from PIL import Image
+
+ if not self.birefnet:
+ return {"status": "error", "message": "BiRefNet not loaded"}
+
+ img_b64 = item.get("image")
+ image = self._decode_base64(img_b64)
+ orig_w, orig_h = image.size
+
+ input_tensor = (
+ self.transform_birefnet(image).unsqueeze(0).to(self.device).half()
+ )
+ with torch.no_grad():
+ preds = self.birefnet(input_tensor)[-1].sigmoid()
+
+ res = torch.nn.functional.interpolate(
+ preds, size=(orig_h, orig_w), mode="bilinear", align_corners=True
+ )
+ mask_np = res.squeeze().cpu().numpy()
+ mask_img = Image.fromarray((mask_np * 255).astype(np.uint8))
+
+ image.putalpha(mask_img)
+ return {"status": "success", "image": self._to_base64(image)}
+
+ @modal.fastapi_endpoint(method="POST")
+ def describe(self, item: dict):
+ import torch
+
+ img_b64 = item.get("image")
+ prompt = item.get("prompt", "<DETAILED_CAPTION>")
+
+ image = self._decode_base64(img_b64)
+
+ inputs = self.florence_processor(text=prompt, images=image, return_tensors="pt")
+ inputs["pixel_values"] = inputs["pixel_values"].to(self.device, torch.float16)
+ inputs["input_ids"] = inputs["input_ids"].to(self.device)
+
+ # --- FIX: Explicitly disable caching to prevent beam search crash ---
+ generated_ids = self.florence_model.generate(
+ input_ids=inputs["input_ids"],
+ pixel_values=inputs["pixel_values"],
+ max_new_tokens=1024,
+ num_beams=3,
+ use_cache=False, # <--- CRITICAL FIX for Florence-2 on newer Transformers
+ )
+
+ text = self.florence_processor.batch_decode(
+ generated_ids, skip_special_tokens=False
+ )[0]
+ clean_text = (
+ text.replace("</s>", "").replace("<s>", "").replace(prompt, "").strip()
+ )
+
+ return {"status": "success", "output": clean_text}
+
+ # --- HELPERS ---
+ def _decode_base64(self, b64_str):
+ from PIL import Image
+
+ if "," in b64_str:
+ b64_str = b64_str.split(",")[1]
+ return Image.open(io.BytesIO(base64.b64decode(b64_str))).convert("RGB")
+
+ def _to_base64(self, img):
+ buffered = io.BytesIO()
+ img.save(buffered, format="PNG")
+ return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+ def _resize_to_limit(self, img, max_dim=1024, multiple=8):
+ from PIL import Image
+
+ w, h = img.size
+ ratio = min(max_dim / w, max_dim / h)
+ new_w = int(w * ratio)
+ new_h = int(h * ratio)
+ new_w = new_w - (new_w % multiple)
+ new_h = new_h - (new_h % multiple)
+ if new_w < multiple:
+ new_w = multiple
+ if new_h < multiple:
+ new_h = multiple
+ return img.resize((new_w, new_h), Image.LANCZOS)
diff --git a/flask/requirements.txt b/flask/requirements.txt
@@ -0,0 +1,20 @@
+accelerate
+bitsandbytes
+diffusers
+einops
+fal-client
+flask
+flask-cors
+kornia
+modal
+numpy
+opencv-python
+pillow
+requests
+safetensors
+scipy
+timm
+toml
+torch
+torchvision
+transformers