Update app.py

app.py

@@ -2,27 +2,35 @@ import gradio as gr
 import replicate
 import os
 from pytube import YouTube
+from urllib.parse import urlparse, parse_qs
 
 REPLICATE_API_TOKEN = os.environ.get("REPLICATE_API_TOKEN")
 
+def clean_youtube_url(url):
+    parsed_url = urlparse(url)
+    video_id = parse_qs(parsed_url.query).get("v", [None])[0]
+    return f"https://www.youtube.com/watch?v={video_id}" if video_id else None
+
 def generate_from_youtube(url, style, ambiance):
     try:
-        yt = YouTube(url)
+        clean_url = clean_youtube_url(url)
+        yt = YouTube(clean_url)
         title = yt.title
-        thumb_url = yt.thumbnail_url
         prompt = f"Miniature YouTube pour un clip {style}, ambiance {ambiance}, titre: {title}"
     except:
         prompt = f"Miniature YouTube {style}, ambiance {ambiance}"
 
-    output = replicate.run(
-        "stability-ai/stable-diffusion",
-        input={
-            "prompt": prompt,
-            "image_dimensions": "512x512"
-        }
-    )
-
-    return output[0]
+    try:
+        output = replicate.run(
+            "stability-ai/stable-diffusion",
+            input={
+                "prompt": prompt,
+                "image_dimensions": "512x512"
+            }
+        )
+        return output[0]
+    except Exception as e:
+        return f"Erreur génération : {str(e)}"
 
 style_choices = ["Rap", "Afro", "Techno", "R&B", "Pop", "Reggae", "Drill", "Autre"]
 
@@ -33,7 +41,7 @@ interface = gr.Interface(
         gr.Dropdown(style_choices, label="Style musical"),
         gr.Textbox(label="Ambiance / Détail artistique (ex : sombre, festif, futuriste...)"),
     ],
-    outputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil", label="Miniature générée"),
     title="YT Thumbnail Generator 🎬",
     description="Colle un lien YouTube et génère automatiquement une miniature stylisée."
 )