import gradio as gr
from diffusers import DiffusionPipeline
import random

# Load your model with LoRA weights
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev")
pipe.load_lora_weights("fyp1/pattern_generation")

def generate_images(prompt, num_images=3):
    # Generate multiple images based on the prompt
    images = []
    
    for _ in range(num_images):
        # Generate a random seed for each image to ensure diversity
        seed = random.randint(0, 100000)
        image = pipe(prompt, parameters={"seed": seed}).images[0]
        images.append(image)

    return images

# Define the Gradio interface
iface = gr.Interface(
    fn=generate_images,
    inputs=[
        gr.Textbox(label="Enter your prompt", lines=2, placeholder="A Kashmiri shawl-inspired pattern..."),
        gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Number of images to generate")
    ],
    outputs=[
        gr.Gallery(label="Generated Images").style(height=300)
    ],
    title="Pattern Generation",
    description="Generate multiple unique patterns based on your prompt."
)

iface.launch()