app.py

import requests
import base64
import numpy as np
from PIL import Image
from io import BytesIO
import gradio as gr

# Định nghĩa API và token cho việc gọi upscale
OCTOAI_TOKEN = "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCIsImtpZCI6IjNkMjMzOTQ5In0.eyJzdWIiOiI2N2M0ODE1Yi1lY2VkLTQzOGMtODJjNi00MmJhYzQwNDYyY2MiLCJ0eXBlIjoidXNlckFjY2Vzc1Rva2VuIiwidGVuYW50SWQiOiJjOWQ4YjZiNi1mZWRlLTQzYmEtYmI3OS00Y2EwNWEwMDRmZTAiLCJ1c2VySWQiOiI2NDIxOTBlZi1hZTMxLTRjYTctYjUzZS0yZmIzNDhjZmQ5MjUiLCJhcHBsaWNhdGlvbklkIjoiYTkyNmZlYmQtMjFlYS00ODdiLTg1ZjUtMzQ5NDA5N2VjODMzIiwicm9sZXMiOlsiRkVUQ0gtUk9MRVMtQlktQVBJIl0sInBlcm1pc3Npb25zIjpbIkZFVENILVBFUk1JU1NJT05TLUJZLUFQSSJdLCJhdWQiOiIzZDIzMzk0OS1hMmZiLTRhYjAtYjdlYy00NmY2MjU1YzUxMGUiLCJpc3MiOiJodHRwczovL2lkZW50aXR5Lm9jdG8uYWkiLCJpYXQiOjE3MjYzNjkyODd9.ZHL9Vd31oeaBC0aDr9N0WzcuxTn96mLbX2SxDYzcf5tZjWfX2oc3OFSrIh1nxN_CWgTlP4ISMsGXp5ees3hqW48G2KsB5HlyywpVdL2J6ChqSRqWDmPV4aVo4sz9It_sc7FvtFRd1VKA7kGffXwiRN3yuIGsxegmxHzNWZWebOA1ArrcfmraojZ_QK5qDu-bF4HeqtpJ7c2Lkyh4rY2S5UKsIjxdRKrmb0QaSjjv5NDzd7zIaq0x1SIaxj2GcHIEhZ3Ex5hAMi3gfKkn5gRJewtNd1kZLBK2OPm-6dEP6BFfNKJ4jSoMvL7bR8C16KExe1MiQlLbFutPsjKoGRDvJw"
url = "https://image.octoai.run/upscaling"
headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {OCTOAI_TOKEN}"
}

# Độ phân giải tùy chọn
resolution_options = {
    "1K": 1024,
    "2K": 2048,
    "4K": 4096,
    "8K": 8192
}

# Hàm gửi yêu cầu đến API và nhận kết quả
def upscale(image, resolution):
    model = "real-esrgan-x4-v3"  # Mô hình mặc định

    # Kích thước ảnh hiện tại
    width, height = image.size

    # Độ phân giải mục tiêu
    target_resolution = resolution_options[resolution]

    # Kiểm tra xem ảnh gốc có lớn hơn kích thước 2048x2048 không
    if width > 2048 or height > 2048:
        scale_factor = 1
        print(f"Image is larger than 2048x2048. Using scale factor: {scale_factor}")
    # Kiểm tra nếu ảnh gốc lớn hơn độ phân giải mục tiêu mà người dùng đã chọn
    elif width > target_resolution or height > target_resolution:
        scale_factor = 1
        print(f"Image is larger than the selected resolution ({resolution}). Using scale factor: {scale_factor}")
    else:
        # Tính scale factor dựa trên kích thước ảnh đầu vào và độ phân giải mục tiêu
        scale_factor_width = target_resolution // width
        scale_factor_height = target_resolution // height
        scale_factor = min(scale_factor_width, scale_factor_height)  # Sử dụng scale nhỏ nhất giữa chiều rộng và chiều cao

        if scale_factor < 1:
            scale_factor = 1  # Đảm bảo scale tối thiểu là 1

        print(f"Calculated scale factor: {scale_factor}")

    # Convert PIL image to base64
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
    
    payload = {
        "model": model,
        "scale": scale_factor,  # Gửi scale tính toán hoặc mặc định là 1
        "init_image": img_str,
        "output_image_encoding": "png"
    }

    response = requests.post(url, json=payload, headers=headers)
    
    if response.status_code == 200:
        result = response.json()
        image_b64 = result.get("image_b64", None)
        
        if image_b64:
            # Convert base64 string back to an image
            image_data = base64.b64decode(image_b64)
            image = Image.open(BytesIO(image_data))
            return image
    else:
        return None

# Tạo giao diện Gradio
css = ".output-image, .input-image, .image-preview {height: 480px !important} "

gr.Interface(
    fn=upscale,
    inputs=[
        gr.Image(type="pil", label="Input Image"),  # Đầu vào ảnh từ người dùng
        gr.Radio(list(resolution_options.keys()), type="value", label="Chọn độ phân giải"),  # Chọn độ phân giải 1K, 2K, 4K, 8K
    ],
    outputs="image",  # Trả về ảnh đầu ra
    title="Image Upscaling 🦆",
    description="Tăng độ phân giải hình ảnh dựa trên mô hình real-esrgan-x4-v3. Nếu ảnh lớn hơn 2048x2048 hoặc lớn hơn độ phân giải mục tiêu, không upscale.",
    allow_flagging="never",
    css=css
).launch()