import gradio as gr
import torch
from PIL import Image
from src.pipeline_pe_clone import FluxPipeline
import spaces
# Load default image from assets as an example
default_image = Image.open("assets/1.png")

@spaces.GPU()
def generate_image(model_path, image, height, width, prompt, guidance_scale, num_steps, lora_name):
    # Load the model
    pipeline = FluxPipeline.from_pretrained(
        model_path,
        torch_dtype=torch.bfloat16,
    ).to('cuda')

    # Load and fuse base LoRA weights
    pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name="pretrain.safetensors")
    pipeline.fuse_lora()
    pipeline.unload_lora_weights()

    # Load selected LoRA effect if not using the pretrained base model
    if lora_name != 'pretrained':
        pipeline.load_lora_weights("nicolaus-huang/PhotoDoodle", weight_name=f"{lora_name}.safetensors")
    
    # Prepare the input image
    condition_image = image.resize((height, width)).convert("RGB")
    
    # Generate the output image
    result = pipeline(
        prompt=prompt,
        condition_image=condition_image,
        height=height,
        width=width,
        guidance_scale=guidance_scale,
        num_inference_steps=num_steps,
        max_sequence_length=512
    ).images[0]
    
    return result

# Create Gradio interface with sliders for numeric inputs
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Model Path", value="black-forest-labs/FLUX.1-dev"),
        gr.Image(label="Input Image", type="pil", value=default_image),
        gr.Slider(label="Height", value=768, minimum=256, maximum=1024, step=64),
        gr.Slider(label="Width", value=512, minimum=256, maximum=1024, step=64),
        gr.Textbox(label="Prompt", value="add a halo and wings for the cat by sksmagiceffects"),
        gr.Slider(label="Guidance Scale", value=3.5, minimum=1.0, maximum=10.0, step=0.1),
        gr.Slider(label="Number of Steps", value=20, minimum=1, maximum=100, step=1),
        gr.Dropdown(
            label="LoRA Name", 
            choices=["pretrained", "sksmagiceffects", "sksmonstercalledlulu", 
                     "skspaintingeffects", "sksedgeeffect", "skscatooneffect"],
            value="sksmagiceffects"
        )
    ],
    outputs=gr.Image(label="Output Image", type="pil"),
    title="FLUX Image Generation with LoRA"
)

if __name__ == "__main__":
    iface.launch()