Реализация AI-генерации аватаров по фотографии

AI-генерация аватаров из фотографии

AI-аватары создают стилизованные версии портрета пользователя: anime, oil painting, fantasy, pixar, cyberpunk. Коммерчески популярный сервис — «загрузи 10 фото → получи 50 аватаров в разных стилях».

Архитектура сервиса аватаров

Пользователь загружает 10–20 фото
    ↓ Предобработка (crop лица, качество, фильтрация)
    ↓ DreamBooth LoRA обучение (~15–30 мин, GPU)
    ↓ Генерация в N стилях (batch inference)
    ↓ Постобработка (GFPGAN face enhance)
    ↓ Готовые аватары пользователю

Обучение персональной LoRA

import subprocess
import asyncio
from pathlib import Path

async def train_personal_avatar_lora(
    user_id: str,
    user_photos: list[bytes],
    gpu_id: int = 0
) -> str:
    work_dir = Path(f"/tmp/avatar/{user_id}")
    work_dir.mkdir(parents=True, exist_ok=True)

    # Сохраняем и предобрабатываем фото
    photos_dir = work_dir / "photos"
    photos_dir.mkdir(exist_ok=True)

    for i, photo_bytes in enumerate(user_photos):
        from PIL import Image
        import io
        img = Image.open(io.BytesIO(photo_bytes)).convert("RGB")
        # Crop по лицу через insightface
        face_crop = crop_face(img)
        if face_crop:
            face_crop.save(photos_dir / f"{i:03d}.jpg", quality=95)

    # Автогенерация подписей
    for img_path in photos_dir.glob("*.jpg"):
        caption = f"photo of {user_id} person, portrait"
        txt_path = img_path.with_suffix(".txt")
        txt_path.write_text(caption)

    # Обучение LoRA
    output_dir = work_dir / "lora"
    output_dir.mkdir(exist_ok=True)

    proc = await asyncio.create_subprocess_exec(
        "accelerate", "launch", "train_network.py",
        "--pretrained_model_name_or_path", "stabilityai/stable-diffusion-xl-base-1.0",
        "--dataset_config", str(work_dir / "dataset.toml"),
        "--output_dir", str(output_dir),
        "--output_name", f"avatar_{user_id}",
        "--network_module", "networks.lora",
        "--network_dim", "32",
        "--network_alpha", "16",
        "--learning_rate", "1e-4",
        "--max_train_steps", "600",
        "--train_batch_size", "1",
        "--mixed_precision", "fp16",
        f"--cuda_ids={gpu_id}",
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    await proc.wait()

    return str(output_dir / f"avatar_{user_id}.safetensors")

Генерация аватаров в стилях

from diffusers import StableDiffusionXLPipeline
import torch

AVATAR_STYLES = {
    "anime": "anime portrait, Studio Ghibli style, cel shading, soft colors",
    "oil_painting": "oil painting portrait, classical style, museum quality, dramatic lighting",
    "cyberpunk": "cyberpunk portrait, neon lights, futuristic, digital art",
    "fantasy": "fantasy portrait, epic illustration, magical background, detailed",
    "pixar": "pixar 3D animation style, cute, cartoon, colorful",
    "sketch": "pencil sketch portrait, detailed, artistic, black and white",
    "watercolor": "watercolor portrait, soft edges, pastel colors, artistic",
    "professional": "professional headshot, business attire, clean background, LinkedIn style",
}

async def generate_avatar_set(
    user_id: str,
    lora_path: str,
    styles: list[str] = None
) -> dict[str, bytes]:
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16
    ).to("cuda")

    pipe.load_lora_weights(lora_path)

    target_styles = styles or list(AVATAR_STYLES.keys())
    results = {}

    for style_name in target_styles:
        style_desc = AVATAR_STYLES[style_name]
        prompt = f"portrait of {user_id} person, {style_desc}, high quality, detailed face"
        negative = "deformed, ugly, low quality, blurry, multiple faces"

        image = pipe(
            prompt=prompt,
            negative_prompt=negative,
            guidance_scale=7.5,
            num_inference_steps=30
        ).images[0]

        # Face enhancement
        img_np = face_enhance(image)

        import io
        buf = io.BytesIO()
        img_np.save(buf, format="PNG")
        results[style_name] = buf.getvalue()

    pipe.unload_lora_weights()
    return results

Celery задача обработки

from celery import Celery

celery_app = Celery("avatars", broker="redis://localhost:6379/0")

@celery_app.task(name="generate_avatars", bind=True, max_retries=2)
def generate_avatars_task(self, user_id: str, photo_paths: list[str]) -> dict:
    try:
        photos = [open(p, "rb").read() for p in photo_paths]
        lora_path = asyncio.run(train_personal_avatar_lora(user_id, photos))
        avatars = asyncio.run(generate_avatar_set(user_id, lora_path))
        urls = {style: upload_to_cdn(f"{user_id}_{style}.png", img) for style, img in avatars.items()}

        notify_user(user_id, urls)
        return {"status": "done", "urls": urls}
    except Exception as exc:
        raise self.retry(exc=exc, countdown=60)

Время обработки одного заказа: обучение LoRA (~600 шагов) — 15–25 минут. Генерация 8 стилей × 4 варианта = 32 изображения — 10–15 минут. Итого: 30–40 минут на GPU (RTX 3090). Сроки разработки сервиса аватаров с очередью и веб-интерфейсом — 4–6 недель.