import gradio as gr
import spaces
import random

import torch
from diffusers import FluxPipeline
from huggingface_hub.utils import RepositoryNotFoundError

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

with open("main.css", "r") as css:
    link = css

@spaces.GPU(duration=70)
def generate(prompt):
    return pipeline(
        prompt=prompt, 
        width=512, 
        height=512, 
        num_inference_steps=20, 
        generator=torch.Generator("cpu").manual_seed(random.randint(42, 69)), 
        guidance_scale=7
    ).images[0]

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.",
            value="Keanu Reeves holding a neon sign reading 'Hello, world!', 32k HDR, paparazzi",
            lines=1,
            interactive=True,
            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()