virtual_tryon.py

from PIL import Image
import requests
import io
import base64
import jwt
import time
import logging
import sys
import asyncio
from requests.exceptions import RequestException

# Set up logging
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('virtual_tryon.log'),
        logging.StreamHandler(sys.stdout)
    ]
)
logger = logging.getLogger(__name__)

# Constants
VALID_CLOTH_TYPES = ["upper", "lower", "full"]
VALID_IMAGE_SIZES = ["256x256", "512x512", "768x768"]
DEFAULT_IMAGE_SIZE = "512x512"
DEFAULT_NUM_STEPS = 30
DEFAULT_GUIDANCE_SCALE = 7.5
DEFAULT_SEED = 42
API_BASE_URL = "https://api.klingai.com"

def generate_api_token(access_key, secret_key):
    """Generate JWT token for API authentication"""
    try:
        current_time = int(time.time())
        payload = {
            "iss": access_key,
            "exp": current_time + 1800,  # 30 minutes expiration
            "nbf": current_time
        }
        
        logger.debug(f"Generating token with payload: {payload}")
        token = jwt.encode(payload, secret_key, algorithm="HS256")
        logger.debug("Token generated successfully")
        return token
        
    except Exception as e:
        logger.error(f"Error generating token: {str(e)}")
        raise

def encode_image_to_base64(image):
    """Convert PIL Image to base64 string"""
    try:
        if isinstance(image, Image.Image):
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            base64_string = base64.b64encode(buffered.getvalue()).decode('utf-8')
            logger.debug(f"Image encoded to base64 successfully. Length: {len(base64_string)}")
            return base64_string
        logger.error("Input is not a PIL Image")
        return None
    except Exception as e:
        logger.error(f"Error encoding image to base64: {str(e)}")
        return None

async def check_task_status(task_id, access_key, secret_key):
    """Check the status of a task"""
    max_attempts = 3
    wait_interval = 20
    attempt = 1

    while attempt <= max_attempts:
        await asyncio.sleep(wait_interval)
        logger.info(f"Checking task status (Attempt {attempt}/{max_attempts})...")

        try:
            # Generate new token for status check
            token = generate_api_token(access_key, secret_key)
            
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {token}"
            }
            
            # Status check endpoint
            url = f"{API_BASE_URL}/v1/images/kolors-virtual-try-on/{task_id}"
            
            response = requests.get(url, headers=headers, verify=False)
            logger.debug(f"Status check response: {response.text}")
            
            result = response.json()
            if response.status_code == 200 and result.get('code') == 0:
                data = result.get('data', {})
                task_status = data.get('task_status', '').lower()
                
                if task_status in ['completed', 'succeed']:
                    images = data.get('task_result', {}).get('images', [])
                    if images:
                        image_url = images[0].get('url')
                        return None, image_url
                    else:
                        return "No images found in the task result.", None
                elif task_status in ['failed', 'error']:
                    error_message = data.get('task_status_msg', 'Task failed.')
                    return f"Task failed: {error_message}", None
                else:
                    logger.info(f"Task status: {task_status}. Waiting for next attempt...")
            else:
                error_message = result.get('message', 'Unknown error occurred.')
                logger.error(f"Error fetching task status: {error_message}")

        except Exception as e:
            logger.error(f"Error checking task status: {str(e)}")
        
        attempt += 1
    
    return "Task did not complete within the expected time.", None

async def apply_virtual_tryon_async(
    person_image,
    garment_image,
    access_key,
    secret_key
):
    """Apply virtual try-on using Kling API asynchronously"""
    try:
        logger.info("Starting virtual try-on process")
        
        # Generate API token
        jwt_token = generate_api_token(access_key, secret_key)
        if not jwt_token:
            return None, "Failed to generate JWT token"

        # Ensure token is string
        if isinstance(jwt_token, bytes):
            jwt_token = jwt_token.decode('utf-8')
        
        # Prepare images
        logger.debug("Preparing images")
        person_base64 = encode_image_to_base64(person_image)
        garment_base64 = encode_image_to_base64(garment_image)
        
        if not person_base64 or not garment_base64:
            logger.error("Failed to convert images to base64")
            return None, "Error converting images to base64"

        # Prepare request
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {jwt_token}"
        }
        
        # Payload structure
        payload = {
            "model_name": "kolors-virtual-try-on-v1",
            "human_image": person_base64,
            "cloth_image": garment_base64
        }

        # Submit task
        url = f"{API_BASE_URL}/v1/images/kolors-virtual-try-on"
        logger.debug(f"Making API request to {url}")
        
        response = requests.post(url, headers=headers, json=payload, verify=False)
        result = response.json()

        if response.status_code == 200 and result.get('code') == 0:
            task_id = result.get('data', {}).get('task_id')
            if not task_id:
                return None, "No task ID received"

            logger.info(f"Task submitted successfully. Task ID: {task_id}")
            
            # Check task status
            error_message, image_url = await check_task_status(task_id, access_key, secret_key)
            
            if error_message:
                return None, error_message
            
            # Download result image
            try:
                image_response = requests.get(image_url)
                if image_response.status_code == 200:
                    return Image.open(io.BytesIO(image_response.content)), "Success"
                else:
                    return None, f"Failed to download result image: {image_response.status_code}"
            except Exception as e:
                return None, f"Error downloading result image: {str(e)}"
        else:
            error_msg = result.get('message', 'Unknown error')
            logger.error(f"API Error: {error_msg}")
            return None, f"API Error: {error_msg}"

    except Exception as e:
        logger.error(f"Unexpected Error: {str(e)}")
        return None, f"Error: {str(e)}"

def apply_virtual_tryon(
    person_image,
    garment_image,
    access_key,
    secret_key,
    cloth_type="upper",
    image_size="512x512",
    num_steps=DEFAULT_NUM_STEPS,
    guidance_scale=DEFAULT_GUIDANCE_SCALE,
    seed=DEFAULT_SEED
):
    """Synchronous wrapper for async virtual try-on function"""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(
            apply_virtual_tryon_async(
                person_image,
                garment_image,
                access_key,
                secret_key
            )
        )
    finally:
        loop.close()