Update app.py

app.py

@@ -1,3 +1,57 @@
 import gradio as gr
+from PIL import Image
+import torch
+from transformers import AutoModel, AutoTokenizer
 
-gr.load("models/Yaquv/rickthenpc").launch()
+# Charger le modèle et le tokenizer depuis Hugging Face
+model_name = "Yaquv/rickthenpc"
+model = AutoModel.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# 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()}
+
+        # 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.")
+
+        return image
+
+    except Exception as e:
+        # Gestion des erreurs avec un retour informatif
+        return f"Erreur lors de la génération : {str(e)}"
+
+# Interface Gradio
+iface = gr.Interface(
+    fn=generate_image,
+    inputs=gr.Textbox(label="Prompt"),
+    outputs=gr.Image(label="Image générée"),
+    title="Générateur d'images RickTheNPC",
+    description="Entrez un prompt et générez une image avec le modèle RickTheNPC."
+)
+
+# Lancer l'application
+if __name__ == "__main__":
+    iface.launch()