import torch
from diffusers import StableDiffusionInstructPix2PixPipeline
import gradio as gr
from PIL import Image
import random

# Load the InstructPix2Pix model
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cpu")

# Initialize a random seed
seed = random.randint(0, 10000)

# Function to reset the seed
def change_style():
    global seed
    seed = random.randint(0, 10000)
    return f"Seed changed to: {seed}"

# Furniture adding function
def add_furniture(image, style, color, room_type):
    # Construct the furniture prompt
    prompt = f"Add {style} style furniture with a {color} tone to this {room_type}."
    
    # Apply the edit using InstructPix2Pix
    edited_image = pipe(prompt=prompt, image=image, num_inference_steps=50, guidance_scale=7.5, generator=torch.manual_seed(seed)).images[0]
    return edited_image

# General image editing function
def edit_image(image, instruction):
    # Apply the edit using InstructPix2Pix
    edited_image = pipe(prompt=instruction, image=image, num_inference_steps=50, guidance_scale=7.5, generator=torch.manual_seed(seed)).images[0]
    return edited_image

# Gradio interface for furniture adding
def furniture_interface():
    with gr.Blocks() as demo_furniture:
        gr.Markdown("## Furniture Adding App")

        # Image upload
        image_input = gr.Image(type="pil", label="Upload an Image")

        # Dropdown for furniture style
        style_input = gr.Dropdown(["Modern", "Classic", "Minimalist", "Vintage"], label="Choose Furniture Style")

        # Dropdown for color
        color_input = gr.Dropdown(["Blue", "Green", "Red", "White", "Black"], label="Choose Furniture Color")

        # Dropdown for room type
        room_type_input = gr.Dropdown(["Living Room", "Bedroom", "Office", "Dining Room"], label="Room Type")

        # Display output image
        result_image = gr.Image(label="Edited Image")

        # Button to apply the furniture transformation
        submit_button = gr.Button("Add Furniture")

        # Button to change the seed (style)
        change_style_button = gr.Button("Change the Style")

        # Output for seed change message
        seed_output = gr.Textbox(label="Seed Info", interactive=False)

        # Define action on button click
        submit_button.click(fn=add_furniture, inputs=[image_input, style_input, color_input, room_type_input], outputs=result_image)
        change_style_button.click(fn=change_style, outputs=seed_output)

    return demo_furniture

# Gradio interface for general image editing
def general_editing_interface():
    with gr.Blocks() as demo_general:
        gr.Markdown("## General Image Editing App")

        # Image upload
        image_input = gr.Image(type="pil", label="Upload an Image")

        # Textbox for instruction
        instruction_input = gr.Textbox(label="Enter the Instruction", placeholder="Describe the changes (e.g., 'Add sunglasses to the person')")

        # Display output image
        result_image = gr.Image(label="Edited Image")

        # Button to apply the transformation
        submit_button = gr.Button("Apply Edit")

        # Button to change the seed (style)
        change_style_button = gr.Button("Change the Style")

        # Output for seed change message
        seed_output = gr.Textbox(label="Seed Info", interactive=False)

        # Define action on button click
        submit_button.click(fn=edit_image, inputs=[image_input, instruction_input], outputs=result_image)
        change_style_button.click(fn=change_style, outputs=seed_output)

    return demo_general

# Launch both Gradio apps
furniture_app = furniture_interface()
general_editing_app = general_editing_interface()

with gr.Blocks() as combined_demo:
    gr.Markdown("## Select the Application")

    with gr.Tab("Furniture Adding App"):
        furniture_app.render()

    with gr.Tab("General Image Editing App"):
        general_editing_app.render()

# Launch the combined Gradio app
combined_demo.launch()