# Installer les bibliothèques nécessaires
!pip install diffusers[torch] transformers gradio

# Importer les bibliothèques nécessaires
from diffusers import DiffusionPipeline
import torch
import gradio as gr

# Vérifier si un GPU est disponible et l'utiliser si possible
device = "cuda" if torch.cuda.is_available() else "cpu"

# Charger le modèle principal
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev").to(device)

# Charger les poids LoRA
pipe.load_lora_weights("gokaygokay/Flux-White-Background-LoRA")

# Fonction pour générer l'image
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Créer l'interface Gradio
interface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Générateur d'images FLUX",
    description="Entrez un prompt pour générer une image avec le modèle FLUX.",
)

# Lancer l'interface
interface.launch()