Instagram-бутиктерге арналған виртуалды киім ауыстыратын бөлме: Клиент фотосы бойынша Imagen 3 / Vertex AI-да Image-to-Image генерациясы

import os
import logging
from typing import Optional
import vertexai
from vertexai.preview.vision_models import ImageGenerationModel, Image as VertexImage

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Инициализация Vertex AI
PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT", "ozat-cloud-kz")
LOCATION = os.environ.get("VERTEX_LOCATION", "us-central1")
vertexai.init(project=PROJECT_ID, location=LOCATION)

# Загрузка модели Imagen 3
# Используем latest версию для генерации и inpainting
generation_model = ImageGenerationModel.from_pretrained("imagen-3.0-generate-001")

def generate_virtual_tryon(
    client_image_bytes: bytes, 
    mask_bytes: bytes, 
    garment_prompt: str
) -> Optional[bytes]:
    """
    Пайплайн Virtual Try-On (VTON) на базе Imagen 3 Edit API (Inpainting).
    :param client_image_bytes: Исходное фото клиента.
    :param mask_bytes: Маска старой одежды (белый цвет - зона замены, черный - фон).
    :param garment_prompt: Детальное описание новой вещи.
    """
    try:
        base_img = VertexImage(client_image_bytes)
        mask_img = VertexImage(mask_bytes)
        
        logger.info(f"Запуск Imagen 3 Inpainting с промптом: {garment_prompt}")
        
        # Inpainting-insert позволяет модели гармонично "вписать" новую текстуру 
        # с учетом освещения и теней базового фото
        response = generation_model.edit_image(
            base_image=base_img,
            mask=mask_img,
            prompt=garment_prompt,
            edit_mode="inpainting-insert",
            guidance_scale=21, # Высокий guidance для точного следования промпту
            number_of_images=1
        )
        
        if not response.images:
            logger.error("Imagen 3 не вернул изображений (возможно сработал фильтр безопасности).")
            return None
            
        return response.images[0]._image_bytes

    except Exception as e:
        logger.error(f"Ошибка вызова Vertex AI: {str(e)}")
        return None

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.responses import Response
import os
from vton_service import generate_virtual_tryon
from masking_service import generate_auto_mask # Гипотетический сервис (например, SAM)

app = FastAPI(title="OZAT VTON Bot API")

@app.post("/api/v1/try-on")
async def virtual_try_on_endpoint(
    photo: UploadFile = File(...),
    garment_description: str = Form(...)
):
    """
    Эндпоинт для Telegram-бота: принимает фото клиентки и текстовое описание одежды.
    Возвращает сгенерированное фото (JPEG).
    """
    if not photo.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Только изображения")
        
    client_image_bytes = await photo.read()
    
    # 1. Автоматическая генерация маски (сегментация текущей одежды)
    # В production здесь вызывается легковесная модель типа U-Net или Segment Anything (SAM)
    mask_bytes = await generate_auto_mask(client_image_bytes)
    
    if not mask_bytes:
        raise HTTPException(status_code=422, detail="Не удалось выделить силуэт одежды")

    # 2. Вызов Vertex AI Imagen 3
    result_bytes = generate_virtual_tryon(
        client_image_bytes=client_image_bytes,
        mask_bytes=mask_bytes,
        garment_prompt=f"A highly photorealistic full-body shot. The person is wearing {garment_description}. Perfect fit, studio lighting, highly detailed fabric texture, no artifacts."
    )

    if not result_bytes:
        raise HTTPException(status_code=500, detail="Ошибка генерации Imagen 3")

    return Response(content=result_bytes, media_type="image/jpeg")

terraform {
  required_providers {
    google = {
      source  = "hashicorp/google"
      version = "~> 5.0"
    }
  }
}

variable "project_id" { type = string }
variable "region" { default = "europe-west1" } # Cloud Run
variable "vertex_region" { default = "us-central1" } # Imagen 3 Availability

provider "google" {
  project = var.project_id
  region  = var.region
}

# Активация API Vertex AI
resource "google_project_service" "aiplatform" {
  service = "aiplatform.googleapis.com"
  disable_on_destroy = false
}

# Сервисный аккаунт для Cloud Run с доступом к Vertex AI
resource "google_service_account" "vton_sa" {
  account_id   = "vton-backend-sa"
  display_name = "SA for VTON Cloud Run"
}

resource "google_project_iam_member" "vertex_user" {
  project = var.project_id
  role    = "roles/aiplatform.user"
  member  = "serviceAccount:${google_service_account.vton_sa.email}"
}

# Cloud Run микросервис
resource "google_cloud_run_v2_service" "vton_backend" {
  name     = "vton-bot-backend"
  location = var.region
  ingress  = "INGRESS_TRAFFIC_ALL"

  template {
    service_account = google_service_account.vton_sa.email
    scaling {
      max_instance_count = 5
    }
    containers {
      image = "${var.region}-docker.pkg.dev/${var.project_id}/apps/vton-backend:latest"
      resources {
        limits = {
          memory = "1024Mi"
          cpu    = "1"
        }
      }
      env {
        name  = "VERTEX_LOCATION"
        value = var.vertex_region
      }
    }
  }
  depends_on = [google_project_iam_member.vertex_user]
}

# 1. Настройка проекта
gcloud auth login
gcloud config set project ozat-cloud-kz

# 2. Включение Vertex AI API
gcloud services enable aiplatform.googleapis.com

# 3. Сборка Docker-образа
gcloud builds submit --tag europe-west1-docker.pkg.dev/ozat-cloud-kz/apps/vton-backend:latest

# 4. Применение Terraform
terraform init
terraform apply -var="project_id=ozat-cloud-kz" -auto-approve