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

# List of available models
models = [
    "HHM29/finetuning_dream_fin",
    "KappaNeuro/needlepoint",
    "Norod78/ClaymationX_LoRA",
    "KappaNeuro/movie-poster",
    "digiplay/MixTape_RocknRoll_v3punk_bake_fp16",
    "digiplay/BeautifulFantasyRealMix_diffusers",
    "Yntec/pineappleAnimeMix",
    "Yntec/DucHaiten-Retro-Diffusers",
    "joachimsallstrom/aether-pixel-lora-for-sdxl",
    "runwayml/stable-diffusion-v1-5",
    "stabilityai/stable-diffusion-xl-base-1.0",
    "CompVis/stable-diffusion-v1-4",
]

def generate_image(model_name, image, prompt, length, temperature, n_samples, use_image2image=False):
    data = {
        "image_prompt": image,
        "prompt": prompt,
        "length": length,
        "temperature": temperature,
        "n_samples": n_samples,
        "model": model_name,
    }

    if use_image2image:
        data["use_image2image"] = True
        data["image2image_prompt"] = image  # Provide the target image for image2image

    response = requests.post("https://api.stable-diffusion.ml/generate", json=data)
    response_json = response.json()

    if response.status_code == 200:
        results = response_json["generated_images"]
        generated_image = np.frombuffer(BytesIO(results[0]["image"]).read(), dtype=np.uint8)
        generated_image = generated_image.reshape(results[0]["metadata"]["height"], results[0]["metadata"]["width"], 3)
        return Image.fromarray(generated_image)
    else:
        return None

def app(model=gr.inputs.Selector(options=models), 
        image=gr.inputs.Image(shape=(None, None)), 
        prompt=gr.inputs.Textbox(default="an image generated with"),
        length=gr.inputs.Slider(1, 20, step=1, default=8),
        temperature=gr.inputs.Slider(0.5, 1.5, step=0.1, default=1),
        n_samples=gr.inputs.Slider(1, 5, step=1, default=1),
        use_image2image=gr.inputs.Boolean(default=False)):
    generated_image = generate_image(model, 
                                  image=image.data if image else None, 
                                  prompt=prompt, 
                                  length=int(length), 
                                  temperature=float(temperature), 
                                  n_samples=int(n_samples),
                                  use_image2image=use_image2image)
    return gr.outputs.Image(as_pil=True)(generated_image) if generated_image else None

if __name__ == "__main__":
    title = "Image Generation App"
    description = "Select a model and customize your image generation or image2image settings!"
    gradio.launch(app, port=8000, title=title, description=description)