import gradio as gr
import requests
import base64
from PIL import Image, ImageDraw
import io
import random

url = "http://73.255.78.150:7909/sdapi/v1/img2img"
session_seed, last_session_seed = None, None

def generate_image(prompt, steps, reset_seed, increment_seed, image_file=None):
    global session_seed, last_session_seed
    
    if reset_seed:
        last_session_seed, session_seed = session_seed, -1 if session_seed != -1 else None
    elif not reset_seed and last_session_seed:
        session_seed = last_session_seed
    elif session_seed is None:
        session_seed = random.randint(0, 10000)

    if increment_seed and session_seed != -1:
        session_seed += 1

    init_image_value = ""
    if image_file:
        image_data = image_file.read()
        encoded_image = base64.b64encode(image_data).decode("utf-8")
        init_image_value = "data:image/jpeg;base64," + encoded_image

    payload = {
        "init_image": init_image_value,
        "prompt": prompt,
        "steps": steps,
        # ... rest of the predefined parameters ...
    }

    response = requests.post(url, json=payload)
    response_data = response.json()

    try:
        return Image.open(io.BytesIO(base64.b64decode(response_data['images'][0])))
    except KeyError:
        error_image = Image.new('RGB', (512, 512), color=(73, 109, 137))
        ImageDraw.Draw(error_image).text((10, 10), f"Error: {response_data.get('error', 'Unknown error')}", fill=(255, 255, 0))
        return error_image

iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(placeholder="Imagine something...", label="Prompt"),
        gr.Slider(minimum=15, maximum=50, step=15, default=5, label="Steps"),
        gr.Checkbox(label="Randomize Composition"),
        gr.Checkbox(label="New Composition"),
        gr.File(label="Upload Image")
    ],
    outputs=gr.Image(label="Generated Image"),
    live=False
)

iface.launch()