Spaces:
Running
Running
| import json | |
| import os | |
| import shutil | |
| import uuid | |
| from fastapi import APIRouter, File, UploadFile, Form, HTTPException, BackgroundTasks | |
| from fastapi.responses import FileResponse | |
| from ..config import settings | |
| from ..core.files import session_dir, sub_dir | |
| from ..services.listing.analyzer import extract_product_attributes | |
| from ..services.listing.prompts import generate_prompts, PROMPT_TYPES | |
| from ..services.listing.generator import ( | |
| generate_image_flux, generate_image_openai, build_scene_prompt, build_negative_prompt, | |
| _next_cf_account, CF_ACCOUNTS, _composite_ref, | |
| ) | |
| from ..services.listing.qc import qc_image | |
| from ..services.listing.textgen import generate_copy_pack, generate_shopee_text | |
| from ..services.auth.storage import is_authorized as _check_user_auth | |
| from ..services.history.auto_save import auto_save_result | |
| from ..core.bg import submit_bg | |
| from ..core.reqctx import get_ctx | |
| from ..services.analytics import track as tracking | |
| router = APIRouter(prefix="/api/listing", tags=["listing"]) | |
| import time | |
| import threading | |
| _sessions = {} | |
| def _new_sid(): | |
| ts = time.strftime("%y%m%d_%H%M%S") | |
| return f"listing_{ts}_{uuid.uuid4().hex[:6]}" | |
| def _get(sid): | |
| s = _sessions.get(sid) | |
| if not s: | |
| raise HTTPException(404, "Session không tồn tại") | |
| return s | |
| def analyze( | |
| file: UploadFile = File(...), | |
| person_image: UploadFile = File(None), | |
| scene_image: UploadFile = File(None), | |
| language: str = Form("vi"), | |
| selected_types: str = Form(""), | |
| quantity: int = Form(9), | |
| api_key: str = Form(""), | |
| image_api_key: str = Form(""), | |
| product_info: str = Form(""), | |
| ): | |
| # `def` (không async): FastAPI chạy trong threadpool -> gọi OpenAI Vision đồng bộ | |
| # (extract_product_attributes/generate_prompts) KHÔNG chặn event loop cả server. | |
| sid = _new_sid() | |
| sess_dir = sub_dir(sid, "listing") | |
| os.makedirs(sess_dir, exist_ok=True) | |
| def _save_img(upload, name): | |
| ext = os.path.splitext(upload.filename or f"{name}.jpg")[1] or ".jpg" | |
| path = os.path.join(sess_dir, f"{name}{ext}") | |
| with open(path, "wb") as f: | |
| f.write(upload.file.read()) | |
| upload.file.seek(0) | |
| return path | |
| img_path = _save_img(file, "original") | |
| person_path = _save_img(person_image, "person") if person_image and person_image.filename else None | |
| scene_path = _save_img(scene_image, "scene") if scene_image and scene_image.filename else None | |
| custom_key = api_key.strip() or None | |
| custom_image_key = image_api_key.strip() or None | |
| user_product_info = product_info.strip() or None | |
| try: | |
| attrs = extract_product_attributes(img_path, api_key=custom_key, user_product_info=user_product_info) | |
| except Exception as e: | |
| shutil.rmtree(sess_dir, ignore_errors=True) | |
| return {"ok": False, "error": str(e)} | |
| prompts = generate_prompts(attrs, language, api_key=custom_key) | |
| if not isinstance(prompts, list) or len(prompts) == 0: | |
| return {"ok": False, "error": "Không thể tạo prompts"} | |
| limited = prompts[:9] | |
| sel = [] | |
| if selected_types: | |
| try: | |
| sel = json.loads(selected_types) | |
| except json.JSONDecodeError: | |
| pass | |
| if sel: | |
| limited = [p for p in limited if p.get("type") in sel] | |
| limited = limited[:min(quantity, 9)] | |
| prompts_file = os.path.join(sess_dir, f"prompts_{sid}.txt") | |
| with open(prompts_file, "w", encoding="utf-8") as f: | |
| for i, p in enumerate(limited): | |
| f.write(f"{'='*60}\n{i+1}. {p.get('label', p.get('type', ''))} ({p.get('type', '')})\n{'='*60}\n{p.get('prompt', '')}\n\n") | |
| _sessions[sid] = { | |
| "id": sid, "status": "analyzed", "prompts": limited, "attributes": attrs, | |
| "image_path": img_path, "image_paths": {"product": img_path, "person": person_path, "scene": scene_path}, | |
| "results": [], "errors": [], "progress": {"current": 0, "total": 0}, | |
| "last_status": None, "copy_pack": None, "shopee_text": None, | |
| "created_at": time.time(), "dir": sess_dir, "original_name": file.filename, | |
| "api_key": custom_key, "image_api_key": custom_image_key, | |
| } | |
| return { | |
| "ok": True, "session_id": sid, "original_name": file.filename, | |
| "prompts": limited, "message": f"Đã tạo {len(limited)} prompts thành công", | |
| } | |
| def create_session( | |
| file: UploadFile = File(...), | |
| selected_types: str = Form(""), | |
| quantity: int = Form(9), | |
| api_key: str = Form(""), | |
| image_api_key: str = Form(""), | |
| ): | |
| sid = _new_sid() | |
| sess_dir = sub_dir(sid, "listing") | |
| os.makedirs(sess_dir, exist_ok=True) | |
| ext = os.path.splitext(file.filename or "image.jpg")[1] or ".jpg" | |
| img_path = os.path.join(sess_dir, f"original{ext}") | |
| with open(img_path, "wb") as f: | |
| shutil.copyfileobj(file.file, f) # stream xuống đĩa, không nạp cả file vào RAM | |
| sel = [] | |
| if selected_types: | |
| try: | |
| sel = json.loads(selected_types) | |
| except json.JSONDecodeError: | |
| pass | |
| types = PROMPT_TYPES[:] | |
| if sel: | |
| types = [t for t in types if t["id"] in sel] | |
| types = types[:min(quantity, 9)] | |
| empty_prompts = [{"type": t["id"], "label": t["label"], "prompt": ""} for t in types] | |
| _sessions[sid] = { | |
| "id": sid, "status": "analyzed", "prompts": empty_prompts, "attributes": {}, | |
| "image_path": img_path, "results": [], "errors": [], "progress": {"current": 0, "total": 0}, | |
| "last_status": None, "copy_pack": None, "shopee_text": None, | |
| "created_at": time.time(), "dir": sess_dir, "original_name": file.filename, | |
| "api_key": api_key.strip() or None, "image_api_key": image_api_key.strip() or None, | |
| } | |
| return {"ok": True, "session_id": sid, "original_name": file.filename, "prompts": empty_prompts} | |
| def get_session(sid: str): | |
| s = _get(sid) | |
| return { | |
| "id": s["id"], "status": s["status"], | |
| "prompts": s["prompts"], "results": s["results"], | |
| "errors": s["errors"], "progress": s["progress"], | |
| "last_status": s["last_status"], | |
| "copy_pack": s.get("copy_pack"), "shopee_text": s.get("shopee_text"), | |
| } | |
| def check_auth(body: dict): | |
| username = (body.get("username") or "").strip() | |
| if not username: | |
| return {"ok": False, "authorized": False, "error": "Missing username"} | |
| authorized = _check_user_auth(username) | |
| return {"ok": True, "authorized": authorized, "username": username} | |
| async def generate(body: dict): | |
| sid = body.get("session_id") | |
| prompt_ids = body.get("prompt_ids") or [] | |
| model = body.get("model") or "gpt-image-2" | |
| quality = body.get("quality") or "high" | |
| custom_prompts = body.get("custom_prompts") or [] | |
| api_key = body.get("api_key") or "" | |
| image_api_key = body.get("image_api_key") or "" | |
| cf_account_id = body.get("cf_account_id") or "" | |
| cf_token = body.get("cf_token") or "" | |
| username = body.get("username", "") | |
| s = _get(sid) | |
| if s["status"] == "generating": | |
| raise HTTPException(400, "Đang trong quá trình sinh ảnh") | |
| if api_key.strip(): | |
| s["api_key"] = api_key.strip() | |
| if image_api_key.strip(): | |
| s["image_api_key"] = image_api_key.strip() | |
| if cf_account_id.strip(): | |
| s["cf_account_id"] = cf_account_id.strip() | |
| if cf_token.strip(): | |
| s["cf_token"] = cf_token.strip() | |
| if username.strip(): | |
| s["username"] = username.strip() | |
| if custom_prompts: | |
| s["prompts"] = custom_prompts | |
| image_prompts = [p for p in (s["prompts"] or []) if p.get("type") and (not prompt_ids or p["type"] in prompt_ids)] | |
| if not image_prompts: | |
| raise HTTPException(400, "Không có prompt để sinh") | |
| s["status"] = "generating" | |
| s["progress"] = {"current": 0, "total": len(image_prompts)} | |
| s["results"] = [] | |
| s["errors"] = [] | |
| s["model"] = model | |
| s["quality"] = quality | |
| ctx = get_ctx() # luồng nền không giữ được ContextVar -> bắt tại đây | |
| tracking.emit("job_queued", ctx, job_id=sid, kind="listing") | |
| submit_bg(_run_generation, sid, image_prompts, ctx) | |
| return {"ok": True, "session_id": sid, "status": "started", "total": len(image_prompts)} | |
| def _run_generation(sid, image_prompts, ctx=None): | |
| ctx = ctx or {} | |
| _t0 = time.time() | |
| tracking.emit("job_started", ctx, job_id=sid, kind="listing") | |
| s = _sessions.get(sid) | |
| if not s: | |
| return | |
| model = s.get("model") or "flux-2-klein-4b" | |
| quality = s.get("quality") or "high" | |
| is_openai = model.startswith("gpt-") or model.startswith("chatgpt-") | |
| max_retries = 2 | |
| custom_key = s.get("api_key") | |
| custom_image_key = s.get("image_api_key") | |
| custom_cf_id = s.get("cf_account_id") or None | |
| custom_cf_token = s.get("cf_token") or None | |
| if s.get("attributes") and not s.get("copy_pack"): | |
| try: | |
| cp = generate_copy_pack(s["attributes"], api_key=custom_key) | |
| s["copy_pack"] = cp | |
| except Exception: | |
| s["copy_pack"] = None | |
| image_paths = s.get("image_paths") or {} | |
| prod_img = image_paths.get("product") or s.get("image_path") | |
| person_img = image_paths.get("person") | |
| scene_img = image_paths.get("scene") | |
| has_person = bool(person_img and os.path.exists(person_img)) | |
| has_scene = bool(scene_img and os.path.exists(scene_img)) | |
| USE_PERSON_TYPES = {"thumbnail", "in_use", "lifestyle", "pain_point", "before_after"} | |
| USE_SCENE_TYPES = {"display", "lifestyle"} | |
| for i, p in enumerate(image_prompts): | |
| s["progress"] = {"current": i + 1, "total": len(image_prompts)} | |
| p_type = p.get("type", "") | |
| ref_type = "product" | |
| ref_img = prod_img | |
| mask_mode = "auto" | |
| if p_type in USE_PERSON_TYPES and has_person: | |
| ref_type = "person" | |
| composite_path = os.path.join(s["dir"], f"composite_{p_type}_{i}.png") | |
| _composite_ref(person_img, prod_img, overlay_scale=0.35, position="bottom_center").save(composite_path, format="PNG") | |
| ref_img = composite_path | |
| mask_mode = "none" | |
| elif p_type in USE_SCENE_TYPES and has_scene: | |
| ref_type = "scene" | |
| composite_path = os.path.join(s["dir"], f"composite_{p_type}_{i}.png") | |
| _composite_ref(scene_img, prod_img, overlay_scale=0.3, position="center").save(composite_path, format="PNG") | |
| ref_img = composite_path | |
| base_prompt = build_scene_prompt(p.get("prompt", ""), p_type, ref_type=ref_type, has_person=has_person, has_scene=has_scene) | |
| neg_prompt = build_negative_prompt(p_type) | |
| last_error = None | |
| success = False | |
| cf_account_tried = 0 | |
| use_custom_cf = not is_openai and custom_cf_id and custom_cf_token | |
| cf_account = None | |
| max_cf_attempts = len(CF_ACCOUNTS) + (1 if use_custom_cf else 0) | |
| if not is_openai: | |
| if use_custom_cf: | |
| cf_account = {"id": custom_cf_id, "token": custom_cf_token, "custom": True} | |
| else: | |
| cf_account = _next_cf_account() | |
| attempt = 0 | |
| while attempt < max_retries: | |
| attempt += 1 | |
| try: | |
| seed = abs(hash(f"{sid}_{i}_{attempt}")) % 2147483647 | |
| msg = f"[{i + 1}/{len(image_prompts)}] \"{p.get('label', p.get('type', ''))}\" attempt {attempt}/{max_retries}..." | |
| s["last_status"] = {"level": "info", "message": msg, "time": time.time()} | |
| output_path = os.path.join(s["dir"], f"{p['type']}_{sid}.png") | |
| if is_openai: | |
| generate_image_openai(base_prompt, output_path, model, ref_img, quality, api_key=custom_image_key or custom_key, mask_mode=mask_mode) | |
| else: | |
| cf_label = cf_account.get("custom", False) and "CF custom" or f"CF #{cf_account_tried + 1}" | |
| s["last_status"] = {"level": "info", "message": f" → dùng {cf_label}", "time": time.time()} | |
| cf_opts = {"account_id": cf_account["id"], "api_token": cf_account["token"]} if cf_account else {} | |
| generate_image_flux( | |
| base_prompt, output_path, ref_img, | |
| negative_prompt=neg_prompt, seed=seed, strength=0.45, | |
| model=model, **cf_opts, | |
| ) | |
| qc_result = qc_image(output_path, s.get("attributes", {}), base_prompt, api_key=custom_key) | |
| if not qc_result.get("product_match", True): | |
| s["last_status"] = {"level": "warning", "message": f"⚠️ QC: {qc_result.get('issues', ['unknown'])} — vẫn chấp nhận", "time": time.time()} | |
| s["results"].append({ | |
| "type": p["type"], "label": p.get("label", p["type"]), | |
| "prompt": base_prompt, "images": [output_path], "status": "success", | |
| }) | |
| s["last_status"] = {"level": "success", "message": f"✅ Đã tạo \"{p.get('label', p['type'])}\"", "time": time.time()} | |
| success = True | |
| break | |
| except Exception as e: | |
| last_error = e | |
| err_msg = str(e) | |
| is_quota = "daily free allocation" in err_msg or "4006" in err_msg | |
| s["last_status"] = {"level": "error", "message": f" \u274c {err_msg[:200]}", "time": time.time()} | |
| if not is_openai: | |
| cf_account_tried += 1 | |
| use_custom_cf = cf_account and cf_account.get("custom") | |
| if use_custom_cf and CF_ACCOUNTS: | |
| cf_account = _next_cf_account() | |
| seed = None | |
| attempt -= 1 | |
| continue | |
| if cf_account_tried < len(CF_ACCOUNTS): | |
| cf_account = _next_cf_account() | |
| seed = None | |
| attempt -= 1 | |
| continue | |
| try: | |
| fb_path = os.path.join(s["dir"], f"{p['type']}_{sid}_fallback.png") | |
| generate_image_openai(base_prompt, fb_path, model, ref_img, quality, api_key=custom_image_key or custom_key, mask_mode=mask_mode) | |
| s["results"].append({ | |
| "type": p["type"], "label": p.get("label", p["type"]), | |
| "prompt": base_prompt, "images": [fb_path], "status": "success", | |
| }) | |
| success = True | |
| break | |
| except Exception: | |
| pass | |
| if attempt < max_retries: | |
| time.sleep(0.8 * attempt) | |
| if not success: | |
| s["results"].append({ | |
| "type": p["type"], "label": p.get("label", p["type"]), | |
| "prompt": base_prompt, "images": [], "status": "error", | |
| "error": str(last_error) if last_error else "Unknown error", | |
| }) | |
| s["errors"].append({"type": p["type"], "error": str(last_error) if last_error else "Unknown error"}) | |
| s["status"] = "completed" | |
| tracking.emit("job_done", ctx, latency_ms=int((time.time() - _t0) * 1000), | |
| job_id=sid, kind="listing", phase="analyze_generate", | |
| items=len(s.get("results") or [])) | |
| s["last_status"] = {"level": "success", "message": "✅ Hoàn thành tất cả! Đang sinh text Shopee...", "time": time.time()} | |
| if s.get("attributes"): | |
| try: | |
| text_result = generate_shopee_text(s["attributes"], s["dir"], sid, api_key=custom_key) | |
| s["shopee_text"] = text_result | |
| except Exception: | |
| pass | |
| s["last_status"] = {"level": "success", "message": "✅ Hoàn thành tất cả (ảnh + Copy Pack + text Shopee)!", "time": time.time()} | |
| _username = s.get("username", "") | |
| if _username: | |
| _files = [] | |
| for r in s.get("results") or []: | |
| for img in r.get("images") or []: | |
| if os.path.exists(img): | |
| _files.append({"path": img, "name": os.path.basename(img)}) | |
| if _files: | |
| auto_save_result(_username, "listing", sid, _files, | |
| title=f"Listing ảnh ({len(_files)} ảnh)", | |
| summary=f"{len(_files)} ảnh · {s.get('original_name', '')}") | |
| def get_results(sid: str): | |
| s = _get(sid) | |
| files = [] | |
| for r in s.get("results") or []: | |
| for img in r.get("images") or []: | |
| if os.path.exists(img): | |
| files.append({"type": r["type"], "label": r.get("label", r["type"]), "file": f"/api/listing/file/{sid}/{os.path.basename(img)}", "filename": os.path.basename(img)}) | |
| return {"session": s, "files": files} | |
| def serve_file(sid: str, filename: str): | |
| s = _get(sid) | |
| file_path = os.path.join(s["dir"], filename) | |
| if not os.path.exists(file_path): | |
| file_path = os.path.join(s["dir"], "..", filename) | |
| if not os.path.exists(file_path): | |
| raise HTTPException(404, "File không tồn tại") | |
| return FileResponse(file_path) | |
| def download_all(sid: str, background_tasks: BackgroundTasks): | |
| import tempfile | |
| import zipfile | |
| s = _get(sid) | |
| tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".zip") | |
| zip_path = tmp.name | |
| tmp.close() | |
| with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf: | |
| for r in s.get("results") or []: | |
| for img in r.get("images") or []: | |
| if os.path.exists(img): | |
| arcname = f"{r['type']}_{os.path.basename(img)}" | |
| zf.write(img, arcname) | |
| prompts_file = os.path.join(s["dir"], f"prompts_{sid}.txt") | |
| if os.path.exists(prompts_file): | |
| zf.write(prompts_file, f"prompts_{sid}.txt") | |
| text_file = os.path.join(s["dir"], f"listing_text_{sid}.txt") | |
| if os.path.exists(text_file): | |
| zf.write(text_file, f"listing_text_{sid}.txt") | |
| background_tasks.add_task(os.unlink, zip_path) | |
| return FileResponse(zip_path, filename=f"listing_images_{sid}.zip", media_type="application/zip") | |
| async def edit_image( | |
| session_id: str = Form(...), | |
| type: str = Form(...), | |
| modification: str = Form(""), | |
| model: str = Form("gpt-image-2"), | |
| reference_image: UploadFile = File(None), | |
| ): | |
| sid = session_id | |
| img_type = type | |
| modification = modification.strip() | |
| s = _get(sid) | |
| if not modification: | |
| raise HTTPException(400, "Thiếu yêu cầu chỉnh sửa") | |
| result = None | |
| for r in s.get("results") or []: | |
| if r["type"] == img_type: | |
| result = r | |
| break | |
| if not result or not result.get("images"): | |
| raise HTTPException(400, "Không tìm thấy ảnh để chỉnh sửa") | |
| source_img = result["images"][-1] | |
| if not os.path.exists(source_img): | |
| raise HTTPException(400, "File ảnh gốc không tồn tại") | |
| input_img = source_img | |
| ref_path = None | |
| if reference_image and reference_image.filename: | |
| ext = os.path.splitext(reference_image.filename)[1] or ".jpg" | |
| ref_path = os.path.join(s["dir"], f"ref_{img_type}_{sid}{ext}") | |
| with open(ref_path, "wb") as f: | |
| f.write(await reference_image.read()) | |
| input_img = ref_path | |
| orig_prompt = "" | |
| for p in s.get("prompts") or []: | |
| if p.get("type") == img_type: | |
| orig_prompt = p.get("prompt", "") | |
| break | |
| if not orig_prompt: | |
| orig_prompt = result.get("prompt", "") | |
| image_paths = s.get("image_paths") or {} | |
| product_ref = image_paths.get("product") or s.get("image_path") | |
| is_openai = model.startswith("gpt-") or model.startswith("chatgpt-") | |
| neg_prompt = build_negative_prompt(img_type) | |
| s["last_status"] = {"level": "info", "message": f"✏️ Đang edit \"{result.get('label', img_type)}\": {modification}", "time": time.time()} | |
| try: | |
| edit_ts = int(time.time() * 1000) | |
| edit_path = os.path.join(s["dir"], f"{img_type}_{sid}_edit_{edit_ts}.png") | |
| preserve = "Giữ NGUYÊN hình dáng, màu sắc, bao bì, nhãn mác của sản phẩm như ảnh gốc. KHÔNG thay đổi thiết kế sản phẩm." | |
| if is_openai: | |
| modified_prompt = f"{orig_prompt}\n\n{preserve}\n\nYêu cầu chỉnh sửa: {modification}" | |
| custom_key = s.get("api_key") | |
| custom_image_key = s.get("image_api_key") | |
| generate_image_openai(modified_prompt, edit_path, model, product_ref, "low", api_key=custom_image_key or custom_key, mask_mode="auto") | |
| else: | |
| modified_prompt = f"{orig_prompt}\n\n{preserve}\n\nYêu cầu chỉnh sửa: {modification}" if not ref_path else modification | |
| cf_account_id = s.get("cf_account_id") or None | |
| cf_token = s.get("cf_token") or None | |
| cf_opts = {} | |
| if cf_account_id and cf_token: | |
| cf_opts = {"account_id": cf_account_id, "api_token": cf_token} | |
| else: | |
| from ..services.listing.generator import _next_cf_account | |
| cf_acc = _next_cf_account() | |
| cf_opts = {"account_id": cf_acc["id"], "api_token": cf_acc["token"]} | |
| generate_image_flux( | |
| modified_prompt, edit_path, input_img, | |
| negative_prompt=neg_prompt, seed=abs(hash(f"{sid}_edit_{img_type}_{edit_ts}")) % 2147483647, | |
| strength=0.6, model=model, **cf_opts, | |
| ) | |
| result["images"] = [edit_path] | |
| s["last_status"] = {"level": "success", "message": f"✅ Đã edit xong \"{result.get('label', img_type)}\"", "time": time.time()} | |
| return {"ok": True, "type": img_type, "filename": os.path.basename(edit_path), "message": "Đã chỉnh sửa ảnh thành công"} | |
| except Exception as e: | |
| s["last_status"] = {"level": "error", "message": f"❌ Edit thất bại: {str(e)[:200]}", "time": time.time()} | |
| return {"ok": False, "error": str(e)} | |
| def delete_session(sid: str): | |
| s = _get(sid) | |
| shutil.rmtree(s["dir"], ignore_errors=True) | |
| _sessions.pop(sid, None) | |
| return {"ok": True, "message": "Đã xoá session"} | |