import gradio as gr
from diffusers import DiffusionPipeline
from PIL import Image
import torch

# Load the model and LoRA weights
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
pipe = pipe.to("cuda")  # Assuming you're using a GPU
pipe.load_lora_weights("enhanceaiteam/Flux-Uncensored-V2")

# Function for text-to-image
def text_to_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Function for image-to-image
def image_to_image(init_image, prompt):
    init_image = init_image.convert("RGB")  # Ensuring image is RGB
    image = pipe(prompt, image=init_image).images[0]  # Passing the image as the initial image for editing
    return image

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Flux Uncensored V2 Demo")

    with gr.Tab("Text-to-Image"):
        with gr.Row():
            text_prompt = gr.Textbox(label="Enter your prompt")
            generated_image = gr.Image(label="Generated Image")
        generate_button = gr.Button("Generate")
        
        # Connect the button to text-to-image function
        generate_button.click(fn=text_to_image, inputs=text_prompt, outputs=generated_image)

    with gr.Tab("Image-to-Image"):
        with gr.Row():
            init_image = gr.Image(source="upload", label="Upload Initial Image", type="pil")
            edit_prompt = gr.Textbox(label="Enter your prompt")
        edited_image = gr.Image(label="Edited Image")
        edit_button = gr.Button("Generate from Image")
        
        # Connect the button to image-to-image function
        edit_button.click(fn=image_to_image, inputs=[init_image, edit_prompt], outputs=edited_image)

demo.launch()