Update app.py

app.py

@@ -12,43 +12,41 @@ from mistralai import Mistral
 
 DEFAULT_KEY = os.getenv("MISTRAL_API_KEY")
 
-def get_mistral_client(alt_key: str = None):
-    key = alt_key.strip() if alt_key and alt_key.strip() else DEFAULT_KEY
+def get_client(alt_key: str = None):
+    key = (alt_key or "").strip() or DEFAULT_KEY
     return Mistral(api_key=key)
 
-def is_remote(src: str):
-    return src.startswith("http://") or src.startswith("https://")
+def is_remote(s: str):
+    return s.startswith("http://") or s.startswith("https://")
 
-def fetch_bytes(source: str):
-    if is_remote(source):
-        r = requests.get(source, timeout=30)
+def fetch_bytes(src: str):
+    if is_remote(src):
+        r = requests.get(src, timeout=30)
         r.raise_for_status()
         return r.content
-    with open(source, "rb") as f:
+    with open(src, "rb") as f:
         return f.read()
 
-def try_ffmpeg_extract_frame(media_path: str, out_path: str):
+def try_ffmpeg_extract_frame(in_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
-    ]
+    cmd = [ffmpeg, "-y", "-i", in_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)}
+def ezgif_convert(media_bytes: bytes, filename: str = "input"):
+    files = {"new-image": (filename, 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)
+    key = m.group(1)
+    conv = requests.post("https://s.ezgif.com/gif-to-jpg", data={"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:
@@ -58,7 +56,7 @@ def ezgif_convert_to_jpg_bytes(media_bytes: bytes, filename_hint: str = "input")
     r2.raise_for_status()
     return r2.content
 
-def convert_media_to_jpeg_bytes(media_bytes: bytes, filename_hint: str = "input") -> bytes:
+def convert_to_jpeg_bytes(media_bytes: bytes, filename_hint: str = "input"):
     try:
         img = Image.open(BytesIO(media_bytes))
         if img.mode != "RGB":
@@ -71,17 +69,17 @@ def convert_media_to_jpeg_bytes(media_bytes: bytes, filename_hint: str = "input"
         return buf.getvalue()
     except Exception:
         with tempfile.TemporaryDirectory() as td:
-            tmp_in = os.path.join(td, filename_hint)
-            with open(tmp_in, "wb") as f:
+            in_path = os.path.join(td, filename_hint)
+            with open(in_path, "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:
+            out_path = os.path.join(td, "frame.jpg")
+            if try_ffmpeg_extract_frame(in_path, out_path) and os.path.exists(out_path):
+                with open(out_path, "rb") as f:
                     return f.read()
-            return ezgif_convert_to_jpg_bytes(media_bytes, filename_hint=filename_hint)
+            return ezgif_convert(media_bytes, filename_hint)
 
-def image_bytes_to_base64_jpeg(image_bytes: bytes):
-    return base64.b64encode(image_bytes).decode("utf-8")
+def to_b64_jpeg(img_bytes: bytes):
+    return base64.b64encode(img_bytes).decode("utf-8")
 
 def build_prompt(use_auto: bool, custom: str):
     if custom and custom.strip():
@@ -92,15 +90,15 @@ def build_prompt(use_auto: bool, custom: str):
                 "poses, facial expressions, clothing, background elements, small distinguishing details, mood, and composition.")
     return "Provide a detailed description of this image."
 
-def generate_description_stream(image_source: str, use_auto: bool, custom_prompt: str, alt_key: str):
+def generate_stream(image_src: str, use_auto: bool, custom_prompt: str, alt_key: str):
     try:
-        raw = fetch_bytes(image_source)
-        jpg = convert_media_to_jpeg_bytes(raw, filename_hint=os.path.basename(image_source) or "input")
+        raw = fetch_bytes(image_src)
+        jpg = convert_to_jpeg_bytes(raw, filename_hint=os.path.basename(image_src) or "input")
     except Exception as e:
         yield f"Error processing media: {e}"
         return
 
-    b64 = image_bytes_to_base64_jpeg(jpg)
+    b64 = to_b64_jpeg(jpg)
     prompt = build_prompt(use_auto, custom_prompt)
     model = "pixtral-12b-2409"
     messages = [{
@@ -112,7 +110,7 @@ def generate_description_stream(image_source: str, use_auto: bool, custom_prompt
         "stream": False
     }]
 
-    client = get_mistral_client(alt_key)
+    client = get_client(alt_key)
     try:
         partial = ""
         for chunk in client.chat.stream(model=model, messages=messages):
@@ -127,44 +125,33 @@ with gr.Blocks() as demo:
 
     with gr.Row():
         with gr.Column(scale=1):
-            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):
+            alt_key = gr.Textbox(label="Alternate Mistral API Key (optional)", type="password")
+            url_input = gr.Textbox(label="Image URL", placeholder="https://...", value="")
             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)
-            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:
-            return upload_file.name
-        if url_val and url_val.strip():
-            return url_val.strip()
-        return ""
-
-    def update_preview_and_return_source(url_val, upload_file):
-        source = prepare_source(url_val, upload_file)
-        if not source:
+        with gr.Column(scale=1):
+            preview = gr.Image(label="Preview", type="pil")
+            out = gr.Textbox(label="Generated Detailed Description", lines=20)
+
+    def load_preview_and_return_url(url):
+        if not url:
             return None, ""
         try:
-            if is_remote(source):
-                r = requests.get(source, timeout=30)
-                r.raise_for_status()
-                return Image.open(BytesIO(r.content)).convert("RGB"), source
-            return Image.open(source).convert("RGB"), source
+            r = requests.get(url, timeout=30)
+            r.raise_for_status()
+            img = Image.open(BytesIO(r.content)).convert("RGB")
+            return img, url
         except Exception:
-            return None, source
+            return None, url
 
-    def start_generation(source, use_auto, custom_p, alt_k):
-        if not source:
-            return "No image provided."
-        # stream generator wrapper
-        for chunk in generate_description_stream(source, use_auto, custom_p, alt_k):
+    def start_gen(url, use_auto, custom_p, alt_k):
+        if not url:
+            return "No URL provided."
+        for chunk in generate_stream(url, use_auto, custom_p, alt_k):
             yield chunk
 
-    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])
+    url_input.change(fn=load_preview_and_return_url, inputs=[url_input], outputs=[preview, url_input])
+    # Generate when Submit (press Enter in URL or add separate button if desired)
+    url_input.submit(fn=start_gen, inputs=[url_input, auto_toggle, custom, alt_key], outputs=[out])
 
 demo.launch()