imageprompt2image.py

import time
import urllib.request
from dotenv import load_dotenv
from byteplussdkarkruntime import Ark
import os
import gradio as gr
import mimetypes
import base64
import urllib
import uuid
import json
import random
import requests
from PIL import Image
from io import BytesIO
import pillow_avif
from byteplussdkarkruntime.types.images.images import OptimizePromptOptions
load_dotenv()

api_key = os.environ.get("ARK_API_KEY_WHITE_LIST")
client = Ark(
    api_key=api_key,
    base_url="https://ark.ap-southeast.bytepluses.com/api/v3"
    )

with open('/home/aiteam1/Videos/Phuong/Do/ai_virtual_gf/jsons/image_level_old.json', 'r', encoding='utf-8') as f:
    elements = json.load(f)

with open('/home/aiteam1/Videos/Phuong/Do/ai_virtual_gf/alss.json') as f:
    categories = json.load(f)

def image_to_base64_uri(image_path: str) -> str:
    """Convert an image to Base64 with full data URI format."""
    # Guess the MIME type (e.g. image/jpeg, image/png)
    mime_type, _ = mimetypes.guess_type(image_path)
    if mime_type is None:
        mime_type = "image/jpeg"  # default fallback

    # Read and encode image
    with open(image_path, "rb") as image_file:
        base64_str = base64.b64encode(image_file.read()).decode("utf-8")

    # Combine into data URI format
    return f"data:{mime_type};base64,{base64_str}"

def convert_image_url_to_jpeg(image_url, output_format="JPEG"):
    """
    Tải ảnh từ URL bất kỳ và chuyển sang JPEG/PNG/JPG
    """
    try:
        response = requests.get(image_url, headers={"User-Agent": "Mozilla/5.0"})
        response.raise_for_status()

        content_type = response.headers.get("Content-Type", "")
        if not content_type.startswith("image/"):
            raise ValueError(f"URL does not return an image (Content-Type: {content_type})")

        img = Image.open(BytesIO(response.content))
        if img.mode in ("RGBA", "P"):
            img = img.convert("RGB")

        save_path = f"converted_image.{output_format.lower()}"
        img.save(save_path, format=output_format)
        print(f"✅ Converted and saved as {save_path}")
        return save_path

    except Exception as e:
        print(f"❌ Error converting image: {e}")
        return None

def gen_image_base64_chat(uuid_charactor, level, option):
    # reference_image_base46 = image_to_base64_uri(reference_image)
    option = str(option)
    flag = True
    i=0
    
    for category in categories:
        for character in category['characters']:
            if uuid_charactor == character['id']:
                introduction = character['introduction']
                tag = character['persona_tag']
                age = character['age']
                image_link = character['image_link']
                gender = character['gender']
    print(f'gender: {gender}')     
    if gender == 0:
        gender = "Male"
    elif gender == 1:
        gender = "Female"
    else:
        gender = random.choice(["Female", "Male"])
    print(f'gender: {gender}')
    for element in elements:
        if gender == element['gender']:
            for level_idx in element['levels']:
                if level == level_idx["level"]:
                    for option_idx in level_idx['options']:
                        if option == option_idx['option_num']:
                            outfit = option_idx['outfit']
                            pose = option_idx['pose']
                            expresion = option_idx['expression']
                            background = option_idx['background']
                            lighting = option_idx['lighting']
                            framing = option_idx['framing']
                            camera_angle = option_idx['camera_angle']
                            # animation = option_idx['animation']
                            break
    
    if gender == "Female":
        another_elements = "A photorealistic full-body portrait of a person with exaggerated bust and hips, very curvy and attractive figure, hourglass shape, smooth, radiant, flawless skin, cinematic lighting, highly detailed, ultra-realistic, 8k resolution"
    else:
        another_elements = "A realistic portrait of a person with smooth, radiant, and fair skin, glowing complexion, soft lighting highlighting facial features, highly detailed, photorealistic, ultra-realistic, cinematic style, 8k resolution"
           
    while flag:
        try:
            prompt = f"""
                    - outfit: {outfit}
                    - pose: {pose}
                    - expresion: {expresion}
                    - background: {background}
                    - lighting: {lighting}
                    - framing: {framing}
                    - camera angle:	{camera_angle}
                    - another_elements: {another_elements}
                """
            start = time.time()
            generated_image = client.images.generate(
                model="seedream-4-0-250828",
                prompt=prompt,
                image = image_link,
                size="1440x2560",
                watermark=False,
                # optimize_prompt_options=OptimizePromptOptions(mode='standard')
            )
            end = time.time()
            print(f'gen image time: {end-start}s')
            log_info = "Generated image successfully"
            url_generated_image = generated_image.data[0].url
            flag = False
            print(f"prompt: {prompt}")
        except Exception as e:
            new_img_link = convert_image_url_to_jpeg(image_link, output_format="JPEG")
            string_64 = image_to_base64_uri(new_img_link)
            prompt = f"""
                    - outfit: {outfit}
                    - pose: {pose}
                    - expresion: {expresion}
                    - background: {background}
                    - lighting: {lighting}
                    - framing: {framing}
                    - camera angle:	{camera_angle}
                    - skin: "A realistic portrait of a person with smooth, radiant, and fair skin, glowing complexion, soft lighting highlighting facial features, highly detailed, photorealistic, ultra-realistic, cinematic style, 8k resolution"
                """
            print('string 64')
            start = time.time()
            generated_image = client.images.generate(
                model="seedream-4-0-250828",
                prompt=prompt,
                image = string_64,
                size="1440x2560",
                watermark=False,
                # optimize_prompt_options=OptimizePromptOptions(mode='fast')
            )
            end = time.time()
            print(f'gen image time: {end-start}s')
            log_info = "Generated image successfully"
            url_generated_image = generated_image.data[0].url
            flag = False
            print(f"prompt: {prompt}")
    # except Exception as e:
    #     print(f"Exception in gen image 1: {e}")
    #     flag = True
    #     i+=1
    #     if i==3:
    #         flag=False
    #         log_info = e
    #         url_generated_image = ""
        id = uuid.uuid4()
        dir_save = 'gradio_app/images'
        os.makedirs(dir_save, exist_ok=True)
        save_path = f'gradio_app/images/{id}.jpeg'
        urllib.request.urlretrieve(url_generated_image, save_path)
        
            
    return save_path, prompt, f'gen image time: {end-start}s'
    
with gr.Blocks() as demo:
    with gr.Column():
        # image = gr.Image(label='Reference Image', type='filepath')
        level = gr.Dropdown(label='Levels', choices=["S1", "S2", "S3", "S4", "S5", "S6", "S7", "S8"])
        option = gr.Dropdown(label="Options", choices=range(1, 11))
        uuid_charactor = gr.Text(label='UUID of character')
        # gender = gr.Dropdown(label='Gender', choices=['Male', 'Female'])
        # outfit = gr.Text(label='Outfit')
        # pose = gr.Text(label='Pose')
        # expresion = gr.Text(label='Expresion')
        # background = gr.Text(label='Background')
        # lighting = gr.Text(label='Lighting')
        # framing = gr.Text(label='Framing')
        # camera_angle = gr.Text(label='Camera Angle')
    with gr.Column():
        output = gr.Image(label='Output Image')
        prompt = gr.Text(label='Prompt')
        log = gr.Text(label='Log Information')
        
demo = gr.Interface(
    fn = gen_image_base64_chat,
    inputs = [uuid_charactor, level, option],
    outputs = [output, prompt, log]
)

demo.launch(server_port=4321, server_name="192.168.5.155", share=True)