from app.workers.celery_app import celery_app from app.core.database import SessionLocal from app.models.user import User from app.models.wardrobe import WardrobeItem from app.models.body import BodyProfile from app.models.tryon import TryOnTask from app.services.storage import storage_service from app.services.qdrant_client import qdrant_service from app.services.tryon_service import tryon_service from app.services.vision_service import vision_service import uuid import time import logging import random import os import io import requests from PIL import Image from app.core.config import settings # Configure logger logger = logging.getLogger("celery_tasks") @celery_app.task(name="app.workers.tasks.analyze_wardrobe_item") def analyze_wardrobe_item(item_id: str): logger.info(f"Starting analysis for wardrobe item: {item_id}") db = SessionLocal() try: item_uuid = uuid.UUID(item_id) if isinstance(item_id, str) else item_id item = db.query(WardrobeItem).filter(WardrobeItem.id == item_uuid).first() if not item: logger.error(f"Item {item_id} not found in DB.") return False # Load image bytes (local /static file or remote URL) image_bytes = None local_path = None if item.image_url.startswith("http://localhost:") or "/static/" in item.image_url: parts = item.image_url.split("/static/") if len(parts) > 1: local_path = os.path.join(settings.UPLOAD_DIR, parts[1]) try: if local_path and os.path.exists(local_path): with open(local_path, "rb") as f: image_bytes = f.read() else: response = requests.get(item.image_url, timeout=15) response.raise_for_status() image_bytes = response.content except Exception as img_err: logger.error(f"Failed to load image from {item.image_url} / {local_path}: {img_err}") # Classify the garment via Gemini Vision API (no local ML). analysis = vision_service.classify_garment(image_bytes) if image_bytes else None # Fallback to simulated metadata if the API is unavailable / failed. if analysis is None: time.sleep(0.5) cat = random.choice(["shirt", "pants", "dress", "jacket", "skirt"]) analysis = { "category": cat, "color_hex": f"#{random.randint(0, 0xFFFFFF):06x}", "style_tags": ["casual", "minimalist"] if cat in ["shirt", "pants"] else ["office", "formal"], } logger.info(f"Wardrobe item {item_id} analyzed with simulated metadata.") else: logger.info(f"Wardrobe item {item_id} classified via Gemini Vision: {analysis['category']}.") # Store a mock 512-d vector in Qdrant (embeddings are not computed locally). vector_id = uuid.uuid4() mock_vector = [random.uniform(-0.1, 0.1) for _ in range(512)] qdrant_service.upsert_vector(vector_id, mock_vector, { "user_id": str(item.user_id), "category": analysis["category"], "style_tags": analysis["style_tags"], }) item.category = analysis["category"] item.color_hex = analysis["color_hex"] item.style_tags = analysis["style_tags"] item.qdrant_vector_id = vector_id db.commit() return True except Exception as e: db.rollback() logger.error(f"Error analyzing wardrobe item {item_id}: {str(e)}", exc_info=True) return False finally: db.close() @celery_app.task(name="app.workers.tasks.reconstruct_body_mesh") def reconstruct_body_mesh(profile_id: str): logger.info(f"Starting 3D body reconstruction for profile: {profile_id}") db = SessionLocal() try: profile_uuid = uuid.UUID(profile_id) if isinstance(profile_id, str) else profile_id profile = db.query(BodyProfile).filter(BodyProfile.id == profile_uuid).first() if not profile: logger.error(f"BodyProfile {profile_id} not found.") return False # Simulate MediaPipe Keypoint extraction & SMPL reconstruction time.sleep(3.0) # Mock keypoints and SMPL betas parameters mock_keypoints = { "shoulder_left": [0.35, 0.25, -0.05], "shoulder_right": [0.65, 0.25, -0.05], "hip_left": [0.40, 0.55, 0.0], "hip_right": [0.60, 0.55, 0.0] } mock_shape_params = [random.uniform(-1.5, 1.5) for _ in range(10)] # Mock GLB avatar file upload to R2 mock_glb_content = b"Mock GLB file bytes representing 3D avatar mesh" mesh_path = f"users/{profile.user_id}/body/avatar.glb" avatar_url = storage_service.upload_file(mock_glb_content, mesh_path, "model/gltf-binary") if avatar_url: profile.keypoints_2d = mock_keypoints profile.shape_params = mock_shape_params profile.avatar_mesh_url = avatar_url db.commit() logger.info(f"Successfully reconstructed body for profile {profile_id}. GLB url: {avatar_url}") return True else: logger.error("Failed to upload GLB mesh to R2.") return False except Exception as e: db.rollback() logger.error(f"Error reconstructing body mesh {profile_id}: {str(e)}", exc_info=True) return False finally: db.close() @celery_app.task(name="app.workers.tasks.run_tryon") def run_tryon(task_id: str): logger.info(f"Starting Virtual Try-On using provider '{settings.TRYON_PROVIDER}' for task: {task_id}") db = SessionLocal() try: task_uuid = uuid.UUID(task_id) if isinstance(task_id, str) else task_id task = db.query(TryOnTask).filter(TryOnTask.id == task_uuid).first() if not task: logger.error(f"TryOnTask {task_id} not found.") return False task.status = "PROCESSING" db.commit() # 1. Load person image person_img = None if task.person_image_url: local_person_path = None if task.person_image_url.startswith("http://localhost:") or "/static/" in task.person_image_url: parts = task.person_image_url.split("/static/") if len(parts) > 1: local_person_path = os.path.join(settings.UPLOAD_DIR, parts[1]) try: if local_person_path and os.path.exists(local_person_path): person_img = Image.open(local_person_path).convert("RGB") else: response = requests.get(task.person_image_url, timeout=15) person_img = Image.open(io.BytesIO(response.content)).convert("RGB") except Exception as e: logger.warning(f"Could not load custom person image from {task.person_image_url}: {e}") # Fallback to default model/hero photo if person_img is None: default_hero_paths = [ "/home/llm/MinhPV/AI_Virtual_Wardrobe/frontend/src/assets/hero.jpg", "/home/llm/MinhPV/AI_Virtual_Wardrobe/src/assets/hero.jpg", ] for path in default_hero_paths: if os.path.exists(path): try: person_img = Image.open(path).convert("RGB") break except Exception as e: logger.error(f"Error loading hero from {path}: {e}") if person_img is None: logger.warning("No hero.jpg found on disk. Creating a blank image for try-on simulation.") person_img = Image.new("RGB", (512, 512), (240, 240, 240)) # 2. Load garments from app.models.wardrobe import WardrobeItem garments = [] if task.garment_item_ids: garments = db.query(WardrobeItem).filter(WardrobeItem.id.in_(task.garment_item_ids)).all() elif task.garment_item_id: garment = db.query(WardrobeItem).filter(WardrobeItem.id == task.garment_item_id).first() if garment: garments = [garment] if not garments: raise ValueError("No valid garments found for the try-on task.") # 3. Collect garment images (bytes + category + url) for the generation service garment_inputs = [] # list[(bytes, category)] garment_urls = [] # aligned URLs, used by URL-based providers (e.g. Replicate) for garment in garments: logger.info(f"Loading garment item ID {garment.id} (category: {garment.category})") garment_bytes = None if garment.image_url: local_garment_path = None if garment.image_url.startswith("http://localhost:") or "/static/" in garment.image_url: parts = garment.image_url.split("/static/") if len(parts) > 1: local_garment_path = os.path.join(settings.UPLOAD_DIR, parts[1]) try: if local_garment_path and os.path.exists(local_garment_path): with open(local_garment_path, "rb") as f: garment_bytes = f.read() else: response = requests.get(garment.image_url, timeout=15) response.raise_for_status() garment_bytes = response.content except Exception as e: logger.error(f"Could not load garment image from {garment.image_url}: {e}") raise ValueError(f"Không thể tải ảnh sản phẩm {garment.image_url}: {e}") if garment_bytes is None: raise ValueError(f"Không có ảnh sản phẩm hợp lệ cho item {garment.id}") garment_inputs.append((garment_bytes, (garment.category or "clothing").lower())) garment_urls.append(garment.image_url) # 4. Encode the person image and generate the try-on via the configured provider person_bio = io.BytesIO() person_img.convert("RGB").save(person_bio, format="JPEG", quality=95) person_bytes = person_bio.getvalue() result_image_bytes = tryon_service.generate( person_bytes=person_bytes, garments=garment_inputs, person_url=task.person_image_url, garment_urls=garment_urls, ) if not result_image_bytes: raise ValueError("Dịch vụ tạo ảnh thử đồ không trả về ảnh.") # Normalize provider output to JPEG so the .jpg path and content-type stay consistent out_img = Image.open(io.BytesIO(result_image_bytes)).convert("RGB") out_bio = io.BytesIO() out_img.save(out_bio, format="JPEG", quality=95) result_image_bytes = out_bio.getvalue() logger.info(f"Successfully generated try-on image via provider '{settings.TRYON_PROVIDER}'.") # 5. Upload final image result_path = f"users/{task.user_id}/try-on/{task.id}.jpg" result_url = storage_service.upload_file(result_image_bytes, result_path, "image/jpeg") if result_url: task.status = "SUCCESS" task.result_url = result_url db.commit() logger.info(f"Successfully completed Try-On task {task_id}. Image url: {result_url}") return True else: task.status = "FAILED" task.error_message = "Không thể lưu trữ kết quả thử đồ lên hệ thống lưu trữ local/R2." db.commit() return False except Exception as e: db.rollback() logger.error(f"Error in Virtual Try-On task {task_id}: {str(e)}", exc_info=True) if task: err = str(e) low = err.lower() if "429" in err or "quota" in low or "resource_exhausted" in low or "limit: 0" in low: friendly = ( "Đã hết quota tạo ảnh của nhà cung cấp (model sinh ảnh không khả dụng " "trên gói miễn phí). Vui lòng bật billing hoặc đổi provider rồi thử lại." ) else: friendly = f"Tạo ảnh thử đồ thất bại: {err}" task.status = "FAILED" task.error_message = friendly db.commit() return False finally: db.close()