import json
import os
import gradio as gr
import requests
import time
import numpy as np
import random

from PIL  import Image

URL = "http://127.0.0.1:8188/prompt"
OUTPUT_DIR = "E:\\SD\\ComfyUI-aki-v1.3\\ComfyUI-aki-v1.3\\output"

def get_latest_image(folder):
    files = os.listdir(folder)
    image_files = [f for f in files if f.lower().endswith(('.png','.jpg','.jpeg'))]
    image_files.sort(key=lambda x: os.path.getmtime(os.path.join(folder,x)))
    latest_image = os.path.join(folder,image_files[-1]) if image_files else None
    return latest_image

def update_seed():
    seed_value = random.randint(1, 1000000)
    json_file_path = "workflow_api.json"  
    with open(json_file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    
    data['10']['inputs']['noise_seed'] = seed_value
    
    with open(json_file_path, 'w', encoding='utf-8') as file:
        json.dump(data, file, indent=4)
    
    return seed_value

def start_queue(prompt_workflow):
    p = {"prompt":prompt_workflow}
    data = json.dumps(p).encode('utf-8')
    requests.post(URL,data=data)

def generate_image(prompt_text,prompt_text2,cfg_count):
    with open("workflow_api.json","r",encoding="utf-8") as file_json:
        prompt = json.load(file_json)
        prompt["6"]["inputs"]["text"] = f"digital artwork of a {prompt_text}"
        prompt["7"]["inputs"]["text"] = f"digital artwork of a {prompt_text2}"
        prompt["10"]["inputs"]["cfg"] = cfg_count

    previous_image = get_latest_image(OUTPUT_DIR)
    random_seed = update_seed()
    print(f"New seed value generated:{random_seed}")
    start_queue(prompt)
    
    while True:
        latest_image = get_latest_image(OUTPUT_DIR)
        if latest_image != previous_image:
            return latest_image
        
        time.sleep(1)
demo = gr.Interface(fn=generate_image,inputs=["text","text","text"],outputs=["image"])

demo.launch(share=True)