Update app.py

app.py

@@ -30,23 +30,46 @@ def get_client(key: str = None):
 def is_remote(src: str) -> bool:
     return bool(src) and src.startswith(("http://", "https://"))
 
-def fetch_bytes(src: str) -> bytes:
+def fetch_bytes(src: str, stream_threshold=20 * 1024 * 1024) -> bytes:
+    """Fetch bytes. For remote large files stream to a temp file then read."""
     if is_remote(src):
-        r = requests.get(src, timeout=60)
-        r.raise_for_status()
-        return r.content
+        with requests.get(src, timeout=60, stream=True) as r:
+            r.raise_for_status()
+            content_length = r.headers.get("content-length")
+            if content_length and int(content_length) > stream_threshold:
+                # stream to temp file
+                fd, path = tempfile.mkstemp()
+                os.close(fd)
+                with open(path, "wb") as f:
+                    for chunk in r.iter_content(8192):
+                        if chunk:
+                            f.write(chunk)
+                with open(path, "rb") as f:
+                    data = f.read()
+                try:
+                    os.remove(path)
+                except Exception:
+                    pass
+                return data
+            else:
+                return r.content
     with open(src, "rb") as f:
         return f.read()
 
 def convert_to_jpeg_bytes(media_bytes: bytes, base_h=480) -> bytes:
     img = Image.open(BytesIO(media_bytes))
+    try:
+        # take first frame for animated formats
+        img.seek(0)
+    except Exception:
+        pass
     if img.mode != "RGB":
         img = img.convert("RGB")
     h = base_h
     w = max(1, int(img.width * (h / img.height)))
     img = img.resize((w, h), Image.LANCZOS)
     buf = BytesIO()
-    img.save(buf, format="JPEG", quality=90)
+    img.save(buf, format="JPEG", quality=85)
     return buf.getvalue()
 
 def b64_jpeg(img_bytes: bytes) -> str:
@@ -59,10 +82,21 @@ def save_bytes_to_temp(b: bytes, suffix: str):
         f.write(b)
     return path
 
+IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
+VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv")
+
+def ext_from_src(src: str) -> str:
+    _, ext = os.path.splitext((src or "").split("?")[0])
+    return ext.lower()
+
 def choose_model_for_src(src: str):
-    lower = (src or "").lower()
-    video_exts = (".mp4", ".mov", ".webm", ".mkv")
-    if is_remote(src) or any(lower.endswith(ext) for ext in video_exts):
+    ext = ext_from_src(src)
+    if ext in VIDEO_EXTS:
+        return DEFAULT_VIDEO_MODEL
+    if ext in IMAGE_EXTS:
+        return DEFAULT_IMAGE_MODEL
+    # fallback: remote URLs likely videos, local files inspect existence
+    if is_remote(src):
         return DEFAULT_VIDEO_MODEL
     return DEFAULT_IMAGE_MODEL
 
@@ -88,17 +122,20 @@ def extract_delta(chunk):
     if not data:
         return None
     try:
-        return data.choices[0].delta.content
-    except Exception:
-        pass
-    try:
+        # try common shapes and coerce to string
         c = data.choices[0].delta
         if isinstance(c, dict):
-            return c.get("content")
+            txt = c.get("content") or c.get("text")
+            return str(txt) if txt is not None else None
     except Exception:
         pass
     try:
-        return data.choices[0].message.content
+        msg = data.choices[0].message
+        if isinstance(msg, dict):
+            txt = msg.get("content")
+        else:
+            txt = getattr(msg, "content", None)
+        return str(txt) if txt is not None else None
     except Exception:
         pass
     try:
@@ -110,63 +147,60 @@ def extract_delta(chunk):
 def generate_final_text(src: str, custom_prompt: str, api_key: str):
     client = get_client(api_key)
     prompt = (custom_prompt.strip() if custom_prompt and custom_prompt.strip() else "Please provide a detailed visual review.")
-    image_exts = (".jpg", ".jpeg", ".png", ".webp", ".gif")
     lower = (src or "").lower()
-    is_image = lower.endswith(image_exts) or (not is_remote(src) and os.path.isfile(src) and lower.endswith(image_exts))
+    ext = ext_from_src(src)
+    is_image = ext in IMAGE_EXTS or (not is_remote(src) and os.path.isfile(src) and ext in IMAGE_EXTS)
     parts = []
 
     def stream_and_collect(model, messages):
         try:
+            # prefer streaming if available; fall back to non-streaming
             stream_gen = None
             try:
                 stream_gen = client.chat.stream(model=model, messages=messages)
             except Exception:
                 stream_gen = None
-
             if stream_gen:
                 for chunk in stream_gen:
                     d = extract_delta(chunk)
                     if d:
                         parts.append(d)
-            else:
-                try:
-                    res = client.chat.complete(model=model, messages=messages, stream=False)
-                except Exception:
-                    res = None
-                if res is None:
-                    parts.append("[Model returned no response]")
-                    return
+                return
+            # fallback complete
+            res = client.chat.complete(model=model, messages=messages, stream=False)
+            # try to extract text
+            try:
+                choices = getattr(res, "choices", None) or res.get("choices", [])
+            except Exception:
+                choices = []
+            if choices:
                 try:
-                    choices = getattr(res, "choices", None) or res.get("choices", [])
-                except Exception:
-                    choices = []
-                if choices:
-                    try:
-                        msg = choices[0].message
-                        if isinstance(msg, dict):
-                            content = msg.get("content")
+                    msg = choices[0].message
+                    if isinstance(msg, dict):
+                        content = msg.get("content")
+                    else:
+                        content = getattr(msg, "content", None)
+                    if content:
+                        if isinstance(content, str):
+                            parts.append(content)
                         else:
-                            content = getattr(msg, "content", None)
-                        if content:
-                            if isinstance(content, str):
-                                parts.append(content)
-                            else:
-                                if isinstance(content, list):
-                                    for c in content:
-                                        if isinstance(c, dict) and c.get("type") == "text":
-                                            parts.append(c.get("text", ""))
-                                elif isinstance(content, dict):
-                                    text = content.get("text") or content.get("content")
-                                    if text:
-                                        parts.append(text)
-                    except Exception:
-                        parts.append(str(res))
-                else:
+                            # handle list/dict shaped content
+                            if isinstance(content, list):
+                                for c in content:
+                                    if isinstance(c, dict) and c.get("type") == "text":
+                                        parts.append(c.get("text", ""))
+                            elif isinstance(content, dict):
+                                text = content.get("text") or content.get("content")
+                                if text:
+                                    parts.append(text)
+                except Exception:
                     parts.append(str(res))
+            else:
+                parts.append(str(res))
         except Exception as e:
-            parts.append(f"\n[Model error: {e}]")
+            parts.append(f"[Model error: {e}]")
 
-    # Image handling
+    # Image path
     if is_image:
         try:
             raw = fetch_bytes(src)
@@ -175,14 +209,14 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         except Exception as e:
             return f"Error processing image: {e}"
         msgs = build_messages_for_image(prompt, b64)
-        stream_and_collect(DEFAULT_IMAGE_MODEL, msgs)
-        return "".join(parts)
+        stream_and_collect(choose_model_for_src(src), msgs)
+        return "".join(parts).strip()
 
-    # Remote video
+    # Remote video: send URL as plain text (avoid embedding non-text chunk types)
     if is_remote(src):
         msgs = build_messages_for_text(prompt, f"Video URL: {src}")
         stream_and_collect(choose_model_for_src(src), msgs)
-        return "".join(parts)
+        return "".join(parts).strip()
 
     # Local video fallback: extract a frame with ffmpeg
     tmp_media = None
@@ -193,26 +227,39 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         tmp_media = save_bytes_to_temp(media_bytes, suffix=ext)
         ffmpeg = shutil.which("ffmpeg")
         if ffmpeg:
+            tmp_frame = None
             try:
                 tmp_frame_fd, tmp_frame = tempfile.mkstemp(suffix=".jpg")
                 os.close(tmp_frame_fd)
                 cmd = [ffmpeg, "-nostdin", "-y", "-i", tmp_media, "-frames:v", "1", "-q:v", "2", tmp_frame]
                 proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-                proc.communicate(timeout=20)
-                if os.path.exists(tmp_frame):
+                out, err = proc.communicate(timeout=30)
+                if proc.returncode == 0 and os.path.exists(tmp_frame) and os.path.getsize(tmp_frame) > 0:
                     with open(tmp_frame, "rb") as f:
                         frame_bytes = f.read()
                     try:
                         jpg = convert_to_jpeg_bytes(frame_bytes, base_h=480)
                         b64 = b64_jpeg(jpg)
                         msgs = build_messages_for_image(prompt, b64)
-                        stream_and_collect(DEFAULT_IMAGE_MODEL, msgs)
-                        return "".join(parts)
+                        stream_and_collect(choose_model_for_src(src), msgs)
+                        return "".join(parts).strip()
                     finally:
                         try:
                             os.remove(tmp_frame)
                         except Exception:
                             pass
+                else:
+                    # ffmpeg failed or produced no frame
+                    try:
+                        if tmp_frame and os.path.exists(tmp_frame):
+                            os.remove(tmp_frame)
+                    except Exception:
+                        pass
+            except subprocess.TimeoutExpired:
+                try:
+                    proc.kill()
+                except Exception:
+                    pass
             except Exception:
                 pass
         return "Unable to process the provided file. Provide a direct image/frame URL or a remote video URL."
@@ -223,7 +270,7 @@ def generate_final_text(src: str, custom_prompt: str, api_key: str):
         except Exception:
             pass
 
-# UI
+# Simple UI: left preview + inputs, right final text only (plain text)
 css = """
 .preview_column { min-width: 380px; }
 .preview_media img, .preview_media video { max-width: 100%; height: auto; }
@@ -231,17 +278,37 @@ css = """
 
 def load_preview(url: str):
     if not url:
-        return None, None, '<div id="preview_box" style="color:gray">No URL</div>'
+        return None, None, "No URL"
     try:
         r = requests.get(url, timeout=30, stream=True)
         r.raise_for_status()
         ctype = r.headers.get("content-type", "")
-        if (ctype and ctype.startswith("video/")) or any(url.lower().endswith(ext) for ext in (".mp4", ".mov", ".webm", ".mkv")):
-            return None, url, '<div id="preview_box" style="color:gray">Video</div>'
+        if (ctype and ctype.startswith("video/")) or any(url.lower().split("?")[0].endswith(ext) for ext in VIDEO_EXTS):
+            return None, url, "Video"
         data = r.content
         img = Image.open(BytesIO(data)).convert("RGB")
-        return img, None, '<div id="preview_box" style="color:gray">Image</div>'
+        return img, None, "Image"
     except Exception:
-        return None, None, '<div id="preview_box" style="color:red">Preview failed</div>'
+        return None, None, "Preview failed"
+
+with gr.Blocks(title="Flux", css=css) as demo:
+    with gr.Row():
+        with gr.Column(scale=1, elem_classes="preview_column"):
+            preview_status = gr.Textbox(value="No URL", interactive=False, label="Preview Status")
+            preview_image = gr.Image(label="Preview Image", type="pil", elem_classes="preview_media")
+            preview_video = gr.Video(label="Preview Video", elem_classes="preview_media")
+
+            url_input = gr.Textbox(label="Image or Video URL", placeholder="https://...", lines=1)
+            custom_prompt = gr.Textbox(label="Prompt (optional)", lines=2, value="")
+            with gr.Accordion("Mistral API Key (optional)", open=False):
+                api_key = gr.Textbox(label="API Key", type="password", max_lines=1)
+            submit = gr.Button("Submit")
+
+        with gr.Column(scale=2):
+            final_text = gr.Textbox(label="Final Text", lines=20)
+
+    url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video, preview_status])
+    submit.click(fn=generate_final_text, inputs=[url_input, custom_prompt, api_key], outputs=[final_text])
 
-with gr.Blocks(title
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))