import gradio as gr
import requests
import time
import requests
import base64




token = '5UAYO8UWHNQKT3UUS9H8V360L76MD72DRIUY9QC2'


    

##############################################################
#################################################

def SD_call(prompt, image_prompt, weight, age, color, hair_color,hair_length,hair_texture,skin_details,eye_colors,NSFW):

    positive = "clothes"
    negative = "naked, nsfw, porn"
    serverless_api_id = '3g77weiulabzuk'
    # Define the URL you want to send the request to
    url = f"https://api.runpod.ai/v2/{serverless_api_id}/runsync"

    # Define your custom headers
    headers = {
        "Authorization": f"Bearer {token}",
        "Accept": "application/json",
        "Content-Type": "application/json"
    }

    if NSFW == True:
        positive = "naked, nsfw"
        negative = "clothes"

    if prompt.strip():
        total_prompt = prompt
    
    else:
        color = ", ".join(color)
        skin_details = ", ".join(skin_details)
        total_prompt = f"masterpiece, best quality, 8k, (looking at viewer:1.1), gorgeous, hot, seductive, {age} years old american {color} woman, {weight} kilos woman, (eye contact:1.1), beautiful face, hyper detailed, best quality, ultra high res, {hair_length} {hair_color} {hair_texture} hair,{eye_colors} eyes, {skin_details} photorealistic, high resolution, detailed, raw photo, 1girl,{image_prompt}, amateur cellphone photography. f8.0, samsung galaxy, noise, jpeg artefacts, poor lighting, low light, underexposed, high contrast "

    
    # Define your data (this could also be a JSON payload)
    print("SD_processing")
    data = {
        "input": {
            "api": {
                "method": "POST",
                "endpoint": "/sdapi/v1/txt2img"
            },
            "payload": {
                "override_settings": {
                    "sd_model_checkpoint": "CyberRealistic",
                    "sd_vae": ""
                },
                "override_settings_restore_afterwards": True,
                "refiner_checkpoint": "",
                "refiner_switch_at": 0.8,
                "prompt": f"{total_prompt}, {positive}",
                "negative_prompt": f"EasyNegative, fat, paintings, sketches, lowres, ((monochrome)), ((grayscale)), bad anatomy, text, error, cropped, signature, watermark, username, blurry, bad feet, poorly drawn face, bad proportions, gross proportions, ng_deepnegative_v1_75t, badhandsv5-neg, {negative}",
                "seed": -1,
                "batch_size": 1,
                "steps": 30,
                "cfg_scale": 7,
                "width": 520,
                "height": 520,
                "sampler_name": "DPM++ SDE Karras",
                "sampler_index": "DPM++ SDE Karras",
                "restore_faces": False
            }
        }
    }




    # Send the POST request with headers and data
    response = requests.post(url, headers=headers, json=data)

    # Check the response
    if response.status_code == 200:
        response_data = response.json()
        msg_id = response_data['id']
        print("Message ID:", msg_id)

        # Poll the status until it's not 'IN_QUEUE'
        while response_data['status'] == 'IN_QUEUE':
            time.sleep(5)  # Wait for 5 seconds before checking again
            response = requests.get(f"{url}/{msg_id}", headers=headers)
            
            try:
                response_data = response.json()
            except Exception as e:
                print("Error decoding JSON:", e)
                print("Response content:", response.text)
                break  # Exit the loop on JSON decoding error

        # Check if the response contains images
        if 'images' in response_data.get('output', {}):
            base64_image = response_data['output']['images'][0]
            image_bytes = base64.b64decode(base64_image)
            
            # Save the image to a file
            image_path = f"output_image_{msg_id}.png"
            with open(image_path, "wb") as img_file:
                img_file.write(image_bytes)

            print(f"Image downloaded successfully: {image_path}")
            
            return image_path

        else:
            return "No images found in the response."
            
    else:
        # Print error message
        return f"Error: {response.status_code} - {response.text}"
                        




def greet(prompt, image_prompt, weight, age, color, hair_color,hair_length,hair_texture,skin_details,eye_colors,NSFW):
    image_path = SD_call(prompt, image_prompt, weight, age, color, hair_color,hair_length,hair_texture,skin_details,eye_colors,NSFW)
    return "Image generated successfully", image_path


demo = gr.Interface(
    fn=greet,
    inputs=[
        gr.Textbox(label="Personal prompt", lines=3),
        gr.Textbox(label="Girl_prompt", lines=3),
        gr.Slider(label="Weight", value=55, minimum=40, maximum=150),
        gr.Slider(label="Age", value=22, minimum=18, maximum=75),
        gr.CheckboxGroup(choices=["asian", "white", "black", "latina", "middle eastern","indigenous", "Mixed"],label="Color",type="value"),
        gr.Radio(["black", "brown", "brunette", "dark brown", "light brown", "blonde", "dirty blonde", "platinum blonde", "red", "auburn", "ginger", "strawberry blonde", "gray", "silver", "white", "blue", "green", "purple", "pink", "rainbow", "multicolored"],label="Hair Color",type="value"),
        gr.Radio(["short", "long", "mi-long"],label="Hair length", type="value"),
        gr.Radio(["straight", "curvy", "wavy"],label="Hair texture", type="value"),
        gr.CheckboxGroup(choices=["((tattoos))", "((birthmark))", "freckles", "((scars))"],label="Skin details", type="value"),
        gr.Radio(["short", "hazel", "green", "blue", "gray", "amber", "black", "red", "violet"],label="Eye Color", type="value"),
        gr.Checkbox(label="NSFW", info="👀👀👀")
    ],
    flagging_options=["blurry", "incorrect", "other"],
    outputs=[gr.Textbox(label="Answer", lines=3), gr.Image(label="Generated Image", type="filepath")],
)

demo.launch(share=True)