import gradio as gr
import spaces
import random

import torch
from diffusers import FluxPipeline

pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
pipeline.load_lora_weights("pepper13/fluxfw")
# pipeline.to("cuda")

@spaces.GPU(duration=70)
def generate(prompt):
    pipe_end = pipeline(prompt=prompt, width=512, height=512, num_inference_steps=24, guidance_scale=7)
    image = pipe_end.images[0]
    return image

with open("main.css", "r") as link:
    with gr.Blocks(css=link) as interface:
        with gr.Column(elem_classes="interface-container"):
            prompt = gr.Textbox(
                label="Prompt",
                info="Describe the image you want to generate.",
                placeholder="e.g., Keanu Reeves holding a neon sign reading 'Hello, world!', 32k HDR, paparazzi",
                # lines=1,
                elem_classes="text-box"
            )
            
            generate_button = gr.Button("Generate Image", elem_classes="btn")
            output = gr.Image(elem_classes="image-output")
            
            generate_button.click(
                fn=generate,
                inputs=[prompt],
                outputs=[output]
            )

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