from PIL import Image
import gradio as gr
from diffusers import DiffusionPipeline
import torch
import numpy as np
import roop.globals
from roop.core import (
    start,
    decode_execution_providers,
    suggest_max_memory,
    suggest_execution_threads,
)
from roop.processors.frame.core import get_frame_processors_modules
from roop.utilities import normalize_output_path
import os
from PIL import Image


# base_model="stabilityai/stable-diffusion-xl-base-1.0"

# # base = DiffusionPipeline.from_pretrained(base_model,
# #         torch_dtype=torch.float16, use_safetensors=True
# # ).to("cuda")
# base = DiffusionPipeline.from_pretrained(base_model,
#         torch_dtype=torch.float16, use_safetensors=True
# )
# # lora_model_path = "./pytorch_weights.safetensors"

# # base.load_lora_weights(lora_model_path)

def swap_face(source_file, target_file,doFaceEnhancer):

    source_path = "input.jpg"
    target_path = "target.jpg"

    # source_image = Image.fromarray(source_file)
    # source_image.save(source_path)
    # target_image = Image.fromarray(target_file)
    # target_image.save(target_path)
    
    # Open and save the source image
    source_image = Image.open(source_file)
    source_image.save(source_path)

    # Open and save the target image
    target_image = Image.open(target_file)
    target_image.save(target_path)

    print("source_path: ", source_path)
    print("target_path: ", target_path)

    roop.globals.source_path = source_path
    roop.globals.target_path = target_path
    output_path = "output.jpg"
    roop.globals.output_path = normalize_output_path(
        roop.globals.source_path, roop.globals.target_path, output_path
    )
    if doFaceEnhancer == True:
        roop.globals.frame_processors = ["face_swapper","face_enhancer"]
    else:
        roop.globals.frame_processors = ["face_swapper"]
    roop.globals.headless = True
    roop.globals.keep_fps = True
    roop.globals.keep_audio = True
    roop.globals.keep_frames = False
    roop.globals.many_faces = False
    roop.globals.video_encoder = "libx264"
    roop.globals.video_quality = 18
    roop.globals.max_memory = suggest_max_memory()
    roop.globals.execution_providers = decode_execution_providers(["cpu"])
    roop.globals.execution_threads = suggest_execution_threads()

    print(
        "start process",
        roop.globals.source_path,
        roop.globals.target_path,
        roop.globals.output_path,
    )

    for frame_processor in get_frame_processors_modules(
        roop.globals.frame_processors
    ):
        if not frame_processor.pre_check():
            return

    start()
    return output_path



def generate_image(prompt, height=1024, width=1024):
    
    # image = base(
    #     prompt=prompt,
    #     negative_prompt="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime), text, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
    #     num_inference_steps=40,
    #     guidance_scale=3,
    #     height=int(height),
    #     width=int(width),
    #     num_images_per_prompt=1,
    # )
    

    # for i in range(len(image.images)):
    #     # Save the generated images
    #     original_image_path = f"./output{i}.png"
    #     image.images[i].save(original_image_path)

    #     # The source image for face swapping
    #     source_image_path = "./roop/IMG_0991.png"

    #     # Run the face swap and optionally face enhancement
    #     swapped_image_path = swap_face(source_image_path, original_image_path, False)

    #     # Load the swapped image and add to the output list
    #     swapped_image = Image.open(swapped_image_path)

    #     return swapped_image
    return None
# examples = [
#     ["xkx man as superman. Uhd, 8k", 1024, 1024],
#     ["photo of xkx man on a beach. Amazing scenery, waves, sun.  Uhd, 8k", 1024, 1536],
# ]
# Set up the Gradio interface
interface = gr.Interface(
    fn=generate_image,
inputs=[
        gr.Text(label="Prompt"),
        gr.Number(label="Height"),
        gr.Number(label="Width")
    ],
    outputs=gr.Image(label="Image"),
    # examples=examples
            
)

interface.launch()