import gradio as gr
import os
import sys
import subprocess
import numpy as np
from PIL import Image
import cv2

import torch
from diffusers import StableDiffusion3Pipeline
from diffusers.models.controlnet_sd3 import ControlNetSD3Model
from diffusers.utils.torch_utils import randn_tensor
from diffusers.utils import load_image

# Clone the specific branch
subprocess.run(["git", "clone", "-b", "sd3_control", "https://github.com/instantX-research/diffusers_sd3_control.git"])

# Change directory to the cloned repository and install it
os.chdir('diffusers_sd3_control')
subprocess.run(["pip", "install", "-e", "."])

# Add the path to the examples directory
sys.path.append(os.path.abspath('./examples/community'))

# Import the required pipeline
from pipeline_stable_diffusion_3_controlnet import StableDiffusion3CommonPipeline



# load pipeline
base_model = 'stabilityai/stable-diffusion-3-medium-diffusers'
pipe = StableDiffusion3CommonPipeline.from_pretrained(
    base_model, 
    controlnet_list=['InstantX/SD3-Controlnet-Canny'],
)
pipe.to('cuda:0', torch.float16)

def resize_image(input_path, output_path, target_height):
    # Open the input image
    img = Image.open(input_path)

    # Calculate the aspect ratio of the original image
    original_width, original_height = img.size
    original_aspect_ratio = original_width / original_height

    # Calculate the new width while maintaining the aspect ratio and the target height
    new_width = int(target_height * original_aspect_ratio)

    # Resize the image while maintaining the aspect ratio and fixing the height
    img = img.resize((new_width, target_height), Image.LANCZOS)

    # Save the resized image
    img.save(output_path)

    return output_path

def infer(image_in, prompt):
    prompt = 'Anime style illustration of a girl wearing a suit. A moon in sky. In the background we see a big rain approaching. text "InstantX" on image'
    n_prompt = 'NSFW, nude, naked, porn, ugly'

    image_to_canny = load_image(image_in)

    image_to_canny = np.array(image_to_canny)
    image_to_canny = cv2.Canny(image_to_canny, 100, 200)
    image_to_canny = image_to_canny[:, :, None]
    image_to_canny = np.concatenate([image_to_canny, image_to_canny, image_to_canny], axis=2)
    image_to_canny = Image.fromarray(image_to_canny)
    
    # controlnet config
    controlnet_conditioning = [
        dict(
            control_index=0,
            control_image=image_to_canny,
            control_weight=0.7,
            control_pooled_projections='zeros'
        )
    ]
    # infer
    image = pipe(
        prompt=prompt,
        negative_prompt=n_prompt,
        controlnet_conditioning=controlnet_conditioning,
        num_inference_steps=28,
        guidance_scale=7.0,
        height=1024,
        width=1024,
    ).images[0]

    return image


with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("""
        # SD3 ControlNet
        """)
        image_in = gr.Image(label="Image reference", sources=["upload"], type="filepath")
        prompt = gr.Textbox(label="Prompt")
        submit_btn = gr.Button("Submit")
        result = gr.Image(label="Result")
    
    submit_btn.click(
        fn = infer,
        inputs = [image_in, prompt],
        outputs = [result],
        show_api=False
    )
demo.queue().launch()