add: stablediffusion pipeline (#1)

- add: stablediffusion pipeline (fb43b6746d6a7f5bf3c64b6b63ebe6e125393e52)


Co-authored-by: Mehdi Merabet <[email protected]>

app.py

@@ -1,41 +1,30 @@
 import gradio as gr
 from PIL import Image
 import torch
-from transformers import AutoModel, AutoTokenizer
+from diffusers import StableDiffusionPipeline
 
-# Charger le modèle et le tokenizer depuis Hugging Face
+# Charger le pipeline de diffusion depuis Hugging Face
 model_name = "Yaquv/rickthenpc"
-model = AutoModel.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+try:
+    pipe = StableDiffusionPipeline.from_pretrained(model_name)
+    pipe = pipe.to(device)
+except Exception as e:
+    print(f"Erreur lors du chargement du modèle : {e}")
+    pipe = None
 
 # Fonction de génération et de post-traitement
 def generate_image(prompt):
     """
     Génère une image à partir du prompt en utilisant le modèle Hugging Face.
     """
-    try:
-        # Préparer les entrées pour le modèle
-        inputs = tokenizer(prompt, return_tensors="pt")
-
-        # Si GPU disponible, envoyer le modèle et les entrées sur le GPU
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        model.to(device)
-        inputs = {key: val.to(device) for key, val in inputs.items()}
+    if pipe is None:
+        return "Le modèle n'a pas pu être chargé."
 
-        # Générer la sortie
-        outputs = model(**inputs)
-
-        # Si la sortie est déjà une image (ex : tableau NumPy), la convertir en image
-        if isinstance(outputs, torch.Tensor):
-            image_data = outputs.detach().cpu().numpy()
-        else:
-            raise ValueError("Le modèle a retourné un format inattendu.")
-
-        # Assurer que la sortie est une image au format (H, W, C)
-        if len(image_data.shape) == 3:  # (Hauteur, Largeur, Canaux)
-            image = Image.fromarray((image_data * 255).astype("uint8"))
-        else:
-            raise ValueError("Les données générées ne sont pas au format image attendu.")
+    try:
+        # Générer l'image
+        image = pipe(prompt).images[0]
 
         return image
 
@@ -54,4 +43,4 @@ iface = gr.Interface(
 
 # Lancer l'application
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()