import gradio as gr
import requests
import json
time

# Configuration
token = "YOUR_TOKEN"  # Replace with your token if needed
api_base_url = "https://api.comfyui.com"  # Replace with the actual API base URL

# Load the workflow JSON file
def load_workflow_json(file_path):
    with open(file_path, "r") as file:
        return json.load(file)

# Load the workflow
workflow = load_workflow_json("workflow.json")

# Function to queue a prompt for processing
def queue_prompt(prompt):
    response = requests.post(
        f"{api_base_url}/prompt", 
        headers={"Authorization": f"Bearer {token}"}, 
        json={"prompt": prompt}
    )
    return response.json()

# Function to get the status of a queued prompt
def get_prompt_status(prompt_id):
    response = requests.get(
        f"{api_base_url}/prompt/{prompt_id}", 
        headers={"Authorization": f"Bearer {token}"}
    )
    return response.json()

# Function to generate an image
def generate_image(positive_prompt, negative_prompt, cfg_scale, steps):
    # Update the workflow with the input parameters
    workflow["positive_prompt"] = positive_prompt
    workflow["negative_prompt"] = negative_prompt
    workflow["cfg_scale"] = cfg_scale
    workflow["steps"] = steps
    
    # Queue the prompt
    prompt_response = queue_prompt(workflow)
    prompt_id = prompt_response.get("id")
    
    # Wait for the prompt to be processed
    max_attempts = 30  # Reduced from 60
    for _ in range(max_attempts):
        status_response = get_prompt_status(prompt_id)
        status = status_response.get("status")
        if status == "completed":
            # Retrieve the generated image
            image_url = status_response.get("result")
            return image_url
        time.sleep(0.5)  # Reduced delay from 1 second
    
    return "Image generation timed out."

# Gradio interface
demo = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.inputs.Textbox(label="Positive Prompt"),
        gr.inputs.Textbox(label="Negative Prompt"),
        gr.inputs.Slider(minimum=1, maximum=20, default=7, label="CFG Scale"),
        gr.inputs.Slider(minimum=10, maximum=100, default=50, label="Steps")
    ],
    outputs=gr.outputs.Image(label="Generated Image")
)

demo.launch()