fix: ensure output is PIL image

app.py

@@ -3,7 +3,7 @@ from PIL import Image
 import torch
 from diffusers import StableDiffusionPipeline
 
-# Charger le pipeline de diffusion depuis Hugging Face
+# Load the diffusion pipeline from Hugging Face
 model_name = "Yaquv/rickthenpc"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -11,46 +11,54 @@ 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}")
+    print(f"Error loading the model: {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.
+    Generates an image from the given prompt using the Hugging Face model.
     """
     if pipe is None:
         raise ValueError("The model couldn't be loaded.")
 
     try:
-        # Générer l'image
+        # Generate the image
         result = pipe(prompt)
-        
-        # Vérifier que le résultat contient des images
+
+        # Check that the result contains images
         if not hasattr(result, 'images') or len(result.images) == 0:
             raise ValueError("The model couldn't generate an image.")
 
         image = result.images[0]
 
-        # S'assurer que l'image est au format PIL.Image
+        # Ensure the image is in PIL.Image format and convert to RGB
         if not isinstance(image, Image.Image):
             image = Image.fromarray(image)
 
+        image = image.convert("RGB")
+
         return image
 
     except Exception as e:
-        # Lever une exception pour que Gradio puisse la gérer
-        raise ValueError(f"Erreur lors de la génération : {str(e)}")
+        # Raise an exception for Gradio to handle
+        raise ValueError(f"Error during image generation: {str(e)}")
 
-# Interface Gradio
+# Define the Gradio Interface
 iface = gr.Interface(
     fn=generate_image,
-    inputs=gr.Textbox(label="Prompt"),
-    outputs=gr.Image(label="Generated Image"),
+    inputs=gr.Textbox(
+        label="Prompt",
+        lines=2,
+        placeholder="Enter your prompt here..."
+    ),
+    outputs=gr.Image(
+        label="Generated Image",
+        type="pil"  # Ensure the output is a PIL Image
+    ),
     title="Rick Generator",
     description="Enter a prompt to generate an image with the Rick Generator model."
 )
 
-# Lancer l'application
+# Launch the Gradio app
 if __name__ == "__main__":
     iface.launch()