Update app.py

app.py

@@ -1,34 +1,68 @@
 import os
 import re
+import shutil
 import base64
+import tempfile
+import subprocess
 from io import BytesIO
 from PIL import Image
 import requests
 import gradio as gr
 from mistralai import Mistral
 
-Mistralclient = Mistral(api_key=os.getenv("MISTRAL_API_KEY"))
+DEFAULT_KEY = os.getenv("MISTRAL_API_KEY")
 
-def is_video_or_gif(url_or_path: str):
-    lower = url_or_path.lower()
-    return any(lower.endswith(ext) for ext in (".gif", ".mp4", ".mov", ".webm", ".avi"))
+def get_mistral_client(alt_key: str = None):
+    key = alt_key.strip() if alt_key and alt_key.strip() else DEFAULT_KEY
+    return Mistral(api_key=key)
+
+def is_remote(src: str):
+    return src.startswith("http://") or src.startswith("https://")
 
 def fetch_bytes(source: str):
-    if source.startswith("http://") or source.startswith("https://"):
+    if is_remote(source):
         r = requests.get(source, timeout=30)
         r.raise_for_status()
         return r.content
-    else:
-        with open(source, "rb") as f:
-            return f.read()
+    with open(source, "rb") as f:
+        return f.read()
+
+def try_ffmpeg_extract_frame(media_path: str, out_path: str):
+    ffmpeg = shutil.which("ffmpeg")
+    if not ffmpeg:
+        return False
+    cmd = [
+        ffmpeg, "-y", "-i", media_path, "-vf", "scale=-2:512", "-frames:v", "1", out_path
+    ]
+    try:
+        subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=30)
+        return os.path.exists(out_path)
+    except Exception:
+        return False
+
+def ezgif_convert_to_jpg_bytes(media_bytes: bytes, filename_hint: str = "input"):
+    files = {"new-image": (filename_hint, media_bytes)}
+    r = requests.post("https://s.ezgif.com/upload", files=files, timeout=60)
+    r.raise_for_status()
+    m = re.search(r'name="file" value="([^"]+)"', r.text)
+    if not m:
+        raise RuntimeError("ezgif upload failed")
+    file_key = m.group(1)
+    conv = requests.post("https://s.ezgif.com/gif-to-jpg", data={"file": file_key}, timeout=60)
+    conv.raise_for_status()
+    m2 = re.search(r'<img src="(https?://s.ezgif.com/tmp/[^"]+)"', conv.text) or re.search(r'<a href="(https?://s.ezgif.com/tmp/[^"]+)"', conv.text)
+    if not m2:
+        raise RuntimeError("ezgif conversion failed")
+    jpg_url = m2.group(1)
+    r2 = requests.get(jpg_url, timeout=60)
+    r2.raise_for_status()
+    return r2.content
 
 def convert_media_to_jpeg_bytes(media_bytes: bytes, filename_hint: str = "input") -> bytes:
-    # If input already an image (jpg/png/webp), open with PIL and convert.
     try:
         img = Image.open(BytesIO(media_bytes))
         if img.mode != "RGB":
             img = img.convert("RGB")
-        # Resize height to 512 maintaining aspect ratio
         base_h = 512
         w = int(img.width * (base_h / img.height))
         img = img.resize((w, base_h), Image.LANCZOS)
@@ -36,55 +70,38 @@ def convert_media_to_jpeg_bytes(media_bytes: bytes, filename_hint: str = "input"
         img.save(buf, format="JPEG", quality=90)
         return buf.getvalue()
     except Exception:
-        # For GIF/MP4 etc., use ezgif conversion: upload the file then fetch a single frame JPG.
-        # This is a simple flow relying on ezgif.org endpoints; replace with your own converter if desired.
-        files = {"new-image": (f"{filename_hint}", media_bytes)}
-        # Upload to ezgif conversion page to get a conversion token
-        r = requests.post("https://s.ezgif.com/upload", files=files, timeout=60)
-        r.raise_for_status()
-        # extract the file key from returned HTML
-        m = re.search(r'name="file" value="([^"]+)"', r.text)
-        if not m:
-            raise RuntimeError("Conversion upload failed")
-        file_key = m.group(1)
-        # Request convert to JPG using ezgif convert endpoint
-        conv = requests.post("https://s.ezgif.com/gif-to-jpg", data={"file": file_key}, timeout=60)
-        conv.raise_for_status()
-        # Find the first resulting JPG link
-        m2 = re.search(r'<img src="(https?://s.ezgif.com/tmp/[^"]+)"', conv.text)
-        if not m2:
-            # try alternate pattern
-            m2 = re.search(r'<a href="(https?://s.ezgif.com/tmp/[^"]+)"', conv.text)
-        if not m2:
-            raise RuntimeError("Conversion failed to produce JPG")
-        jpg_url = m2.group(1)
-        r2 = requests.get(jpg_url, timeout=60)
-        r2.raise_for_status()
-        return r2.content
+        with tempfile.TemporaryDirectory() as td:
+            tmp_in = os.path.join(td, filename_hint)
+            with open(tmp_in, "wb") as f:
+                f.write(media_bytes)
+            tmp_out = os.path.join(td, "frame.jpg")
+            if try_ffmpeg_extract_frame(tmp_in, tmp_out) and os.path.exists(tmp_out):
+                with open(tmp_out, "rb") as f:
+                    return f.read()
+            return ezgif_convert_to_jpg_bytes(media_bytes, filename_hint=filename_hint)
 
 def image_bytes_to_base64_jpeg(image_bytes: bytes):
     return base64.b64encode(image_bytes).decode("utf-8")
 
-def build_prompt(auto_prompt: bool, custom_prompt: str):
-    if custom_prompt and custom_prompt.strip():
-        return custom_prompt.strip()
-    if auto_prompt:
-        return ("Provide a highly detailed, vivid, and exhaustive visual description of the image. "
-                "Include objects, textures, colors, lighting, perspective, environment, clothing, "
-                "facial expressions, poses, mood, camera lens and settings, and any notable small details. "
-                "Write in descriptive, actionable phrases suitable for use as image-generation prompts.")
+def build_prompt(use_auto: bool, custom: str):
+    if custom and custom.strip():
+        return custom.strip()
+    if use_auto:
+        return ("Provide an exhaustive, highly detailed visual description suitable as an image-generation prompt. "
+                "Include objects, materials, textures, colors, lighting, viewpoint, camera lens and settings, "
+                "poses, facial expressions, clothing, background elements, small distinguishing details, mood, and composition.")
     return "Provide a detailed description of this image."
 
-def generate_description(image_source: str, use_auto_prompt: bool, custom_prompt: str):
+def generate_description_stream(image_source: str, use_auto: bool, custom_prompt: str, alt_key: str):
     try:
         raw = fetch_bytes(image_source)
-        jpg_bytes = convert_media_to_jpeg_bytes(raw, filename_hint=image_source.split("/")[-1])
+        jpg = convert_media_to_jpeg_bytes(raw, filename_hint=os.path.basename(image_source) or "input")
     except Exception as e:
-        return f"Error processing image/media: {e}"
-
-    b64 = image_bytes_to_base64_jpeg(jpg_bytes)
-    prompt = build_prompt(use_auto_prompt, custom_prompt)
+        yield f"Error processing media: {e}"
+        return
 
+    b64 = image_bytes_to_base64_jpeg(jpg)
+    prompt = build_prompt(use_auto, custom_prompt)
     model = "pixtral-12b-2409"
     messages = [{
         "role": "user",
@@ -95,57 +112,59 @@ def generate_description(image_source: str, use_auto_prompt: bool, custom_prompt
         "stream": False
     }]
 
+    client = get_mistral_client(alt_key)
     try:
         partial = ""
-        for chunk in Mistralclient.chat.stream(model=model, messages=messages):
+        for chunk in client.chat.stream(model=model, messages=messages):
             if chunk.data.choices[0].delta.content is not None:
                 partial += chunk.data.choices[0].delta.content
                 yield partial
     except Exception as e:
-        yield f"Error from model: {e}"
+        yield f"Model error: {e}"
 
 with gr.Blocks() as demo:
     gr.Markdown("Image To Flux Prompt")
 
     with gr.Row():
         with gr.Column(scale=1):
-            url_input = gr.Textbox(label="Image URL or local path", placeholder="https://... or upload below")
-            upload = gr.File(label="Or upload image/video (gif/mp4 allowed)", file_types=["image","video"])
+            url_input = gr.Textbox(label="Image URL or local path", placeholder="https://... or leave blank if uploading")
+            upload = gr.File(label="Or upload image/video (gif/mp4 allowed)", file_types=None)
             preview = gr.Image(label="Preview", type="pil")
             submit = gr.Button("Submit")
         with gr.Column(scale=1):
             custom = gr.Textbox(label="Custom prompt (optional)", lines=4, placeholder="Leave blank to use auto prompt")
             auto_toggle = gr.Checkbox(label="Use detailed auto prompt", value=True)
-            out = gr.Textbox(label="Generated Detailed Description", lines=12)
+            alt_key = gr.Textbox(label="Alternate Mistral API Key (optional)", type="password")
+            out = gr.Textbox(label="Generated Detailed Description", lines=16)
 
     def prepare_source(url_val, upload_file):
-        if upload_file is not None:
+        if upload_file:
             return upload_file.name
         if url_val and url_val.strip():
             return url_val.strip()
         return ""
 
-    def update_preview(source):
+    def update_preview_and_return_source(url_val, upload_file):
+        source = prepare_source(url_val, upload_file)
         if not source:
-            return None
+            return None, ""
         try:
-            if source.startswith("http://") or source.startswith("https://"):
+            if is_remote(source):
                 r = requests.get(source, timeout=30)
                 r.raise_for_status()
-                return Image.open(BytesIO(r.content)).convert("RGB")
-            else:
-                return Image.open(source).convert("RGB")
+                return Image.open(BytesIO(r.content)).convert("RGB"), source
+            return Image.open(source).convert("RGB"), source
         except Exception:
-            return None
+            return None, source
 
-    submit.click(fn=lambda url, f: prepare_source(url, f), inputs=[url_input, upload], outputs=[preview], _js=None)
-    # On submit, generate description stream. Use a small wrapper to pass other inputs.
-    def run_and_stream(src, use_auto, custom_p):
-        if not src:
+    def start_generation(source, use_auto, custom_p, alt_k):
+        if not source:
             return "No image provided."
-        for chunk in generate_description(src, use_auto, custom_p):
+        # stream generator wrapper
+        for chunk in generate_description_stream(source, use_auto, custom_p, alt_k):
             yield chunk
 
-    submit.click(fn=run_and_stream, inputs=[url_input, auto_toggle, custom], outputs=[out])
+    submit.click(fn=update_preview_and_return_source, inputs=[url_input, upload], outputs=[preview, url_input])
+    submit.click(fn=start_generation, inputs=[url_input, auto_toggle, custom, alt_key], outputs=[out])
 
 demo.launch()