Update app.py

app.py

@@ -1,210 +1,220 @@
-import spaces
-import gradio as gr
-import torch
-from PIL import Image
-from diffusers import DiffusionPipeline
-import random
+import os
+import re
 import uuid
-from typing import Tuple
-import numpy as np
-
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    return seed
-
-MAX_SEED = np.iinfo(np.int32).max
-
-if not torch.cuda.is_available():
-    DESCRIPTIONz += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
-
-base_model = "black-forest-labs/FLUX.1-dev"
-pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
-
-lora_repo = "strangerzonehf/3d-Station-Toon"
-trigger_word = "3d station toon"  # Leave trigger_word blank if not used.
-
-pipe.load_lora_weights(lora_repo)
-pipe.to("cuda")
-
-style_list = [
-    {
-        "name": "3840 x 2160",
-        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-    },
-    {
-        "name": "2560 x 1440",
-        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-    },
-    {
-        "name": "HD+",
-        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-    },
-    {
-        "name": "Style Zero",
-        "prompt": "{prompt}",
-    },
-]
-
-styles = {k["name"]: k["prompt"] for k in style_list}
-
-DEFAULT_STYLE_NAME = "3840 x 2160"
-STYLE_NAMES = list(styles.keys())
+import json
+import time
+import random
+import asyncio
+import cv2
+from datetime import datetime, timedelta
+from threading import Thread
 
-def apply_style(style_name: str, positive: str) -> str:
-    return styles.get(style_name, styles[DEFAULT_STYLE_NAME]).replace("{prompt}", positive)
+import gradio as gr
+import numpy as np
+from PIL import Image
+from huggingface_hub import hf_hub_download
+from vllm import LLM
+from vllm.sampling_params import SamplingParams
+
+# -----------------------------------------------------------------------------
+# Helper functions
+# -----------------------------------------------------------------------------
+
+def progress_bar_html(label: str) -> str:
+    """Return an HTML snippet for a progress bar."""
+    return f'''
+<div style="display: flex; align-items: center;">
+    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
+    <div style="width: 110px; height: 5px; background-color: #F0FFF0; border-radius: 2px; overflow: hidden;">
+        <div style="width: 100%; height: 100%; background-color: #00FF00; animation: loading 1.5s linear infinite;"></div>
+    </div>
+</div>
+<style>
+@keyframes loading {{
+    0% {{ transform: translateX(-100%); }}
+    100% {{ transform: translateX(100%); }}
+}}
+</style>
+    '''
+
+def downsample_video(video_path: str, num_frames: int = 10):
+    """
+    Downsample a video to extract a set number of evenly spaced frames.
+    Returns a list of tuples (PIL.Image, timestamp in seconds).
+    """
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    if total_frames <= 0 or fps <= 0:
+        vidcap.release()
+        return frames
+    # Get evenly spaced frame indices.
+    frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
+    for i in frame_indices:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, image = vidcap.read()
+        if success:
+            # Convert BGR to RGB and then to a PIL Image.
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            pil_image = Image.fromarray(image)
+            timestamp = round(i / fps, 2)
+            frames.append((pil_image, timestamp))
+    vidcap.release()
+    return frames
+
+def load_system_prompt(repo_id: str, filename: str) -> str:
+    """
+    Load the system prompt from the given Hugging Face Hub repo file,
+    and format it with the model name and current dates.
+    """
+    file_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    with open(file_path, "r") as file:
+        system_prompt = file.read()
+    today = datetime.today().strftime("%Y-%m-%d")
+    yesterday = (datetime.today() - timedelta(days=1)).strftime("%Y-%m-%d")
+    model_name = repo_id.split("/")[-1]
+    return system_prompt.format(name=model_name, today=today, yesterday=yesterday)
+
+# -----------------------------------------------------------------------------
+# Global Settings and Model Initialization
+# -----------------------------------------------------------------------------
+
+# Model details (adjust as needed)
+MODEL_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"
+# Load the system prompt from HF Hub (make sure SYSTEM_PROMPT.txt exists in the repo)
+SYSTEM_PROMPT = load_system_prompt(MODEL_ID, "SYSTEM_PROMPT.txt")
+# If you prefer a hardcoded system prompt, you can use:
+# SYSTEM_PROMPT = "You are a conversational agent that always answers straight to the point, and ends with an ASCII cat."
+
+# Initialize the Mistral LLM via vllm.
+# Note: Running this model on GPU may require very high VRAM.
+llm = LLM(model=MODEL_ID, tokenizer_mode="mistral")
+
+# -----------------------------------------------------------------------------
+# Main Generation Function
+# -----------------------------------------------------------------------------
 
-@spaces.GPU(duration=60, enable_queue=True)
 def generate(
-    prompt: str,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    randomize_seed: bool = False,
-    style_name: str = DEFAULT_STYLE_NAME,
-    progress=gr.Progress(track_tqdm=True),
+    input_dict: dict,
+    chat_history: list,
+    max_new_tokens: int = 512,
+    temperature: float = 0.15,
+    top_p: float = 0.9,
+    top_k: int = 50,
 ):
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-
-    positive_prompt = apply_style(style_name, prompt)
+    """
+    The main generation function for the Mistral chatbot.
+    It supports:
+      - Text-only inference.
+      - Image inference (attaches image file paths).
+      - Video inference (extracts and attaches sampled video frames).
+    """
+    text = input_dict["text"]
+    files = input_dict.get("files", [])
+    # Prepare the conversation with a system prompt.
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT}
+    ]
     
-    if trigger_word:
-        positive_prompt = f"{trigger_word} {positive_prompt}"
-
-    images = pipe(
-        prompt=positive_prompt,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=30,
-        num_images_per_prompt=1,
-        output_type="pil",
-    ).images
-    image_paths = [save_image(img) for img in images]
-    print(image_paths)
-    return image_paths, seed
-
-examples = [
-    "Super Realism, High-resolution photograph, woman, UHD, photorealistic, shot on a Sony A7III --chaos 20 --ar 1:2 --style raw --stylize 250",
-    "Woman in a red jacket, snowy, in the style of hyper-realistic portraiture, caninecore, mountainous vistas, timeless beauty, palewave, iconic, distinctive noses --ar 72:101 --stylize 750 --v 6",
-    "Super Realism, Headshot of handsome young man, wearing dark gray sweater with buttons and big shawl collar, brown hair and short beard, serious look on his face, black background, soft studio lighting, portrait photography --ar 85:128 --v 6.0 --style",
-    "Super-realism, Purple Dreamy, a medium-angle shot of a young woman with long brown hair, wearing a pair of eye-level glasses, stands in front of a backdrop of purple and white lights. The womans eyes are closed, her lips are slightly parted, as if she is looking up at the sky. Her hair is cascading over her shoulders, framing her face. She is wearing a sleeveless top, adorned with tiny white dots, and a gold chain necklace around her neck. Her left earrings are dangling from her ears, adding a pop of color to the scene."
-]
-
-css = '''
-.gradio-container{max-width: 888px !important}
-h1{text-align:center}
-footer {
-    visibility: hidden
-}
-.submit-btn {
-    background-color: #e34949 !important;
-    color: white !important;
-}
-.submit-btn:hover {
-    background-color: #ff3b3b !important;
-}
-'''
-
-with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
-    with gr.Row():
-        with gr.Column(scale=1):
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-            run_button = gr.Button("Generate as ( 768 x 1024 )🤗", scale=0, elem_classes="submit-btn")
-            
-            with gr.Accordion("Advanced options", open=True, visible=True):
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=0,
-                    visible=True
-                )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-                
-                with gr.Row(visible=True):
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=512,
-                        maximum=2048,
-                        step=64,
-                        value=768,
-                    )
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=512,
-                        maximum=2048,
-                        step=64,
-                        value=1024,
-                    )
-                
-                with gr.Row():
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=0.1,
-                        maximum=20.0,
-                        step=0.1,
-                        value=3.0,
-                    )
-                    num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
-                        minimum=1,
-                        maximum=40,
-                        step=1,
-                        value=30,
-                    )
-
-                style_selection = gr.Radio(
-                    show_label=True,
-                    container=True,
-                    interactive=True,
-                    choices=STYLE_NAMES,
-                    value=DEFAULT_STYLE_NAME,
-                    label="Quality Style",
-                )
-        
-        with gr.Column(scale=2):
-            result = gr.Gallery(label="Result", columns=1, show_label=False)
-            
-            gr.Examples(
-                examples=examples,
-                inputs=prompt,
-                outputs=[result, seed],
-                fn=generate,
-                cache_examples=False,
-            )
-
-    gr.on(
-        triggers=[
-            prompt.submit,
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            randomize_seed,
-            style_selection,
-        ],
-        outputs=[result, seed],
-        api_name="run",
+    # Check if any file is provided
+    video_extensions = (".mp4", ".mov", ".avi", ".mkv", ".webm")
+    if files:
+        # If any file is a video, use video inference branch.
+        if any(str(f).lower().endswith(video_extensions) for f in files):
+            # Remove any @video-infer tag if present.
+            prompt_clean = re.sub(r"@video-infer", "", text, flags=re.IGNORECASE).strip().strip('"')
+            video_path = files[0]  # currently process the first video file
+            frames = downsample_video(video_path)
+            # Build a list that contains the prompt plus each frame information.
+            user_content = [{"type": "text", "text": prompt_clean}]
+            for frame in frames:
+                image, timestamp = frame
+                # Save the frame to a temporary file.
+                image_path = f"video_frame_{uuid.uuid4().hex}.png"
+                image.save(image_path)
+                user_content.append({"type": "text", "text": f"Frame at {timestamp} seconds:"})
+                user_content.append({"type": "image_path", "image_path": image_path})
+            messages.append({"role": "user", "content": user_content})
+        else:
+            # Assume provided files are images.
+            prompt_clean = re.sub(r"@mistral", "", text, flags=re.IGNORECASE).strip().strip('"')
+            user_content = [{"type": "text", "text": prompt_clean}]
+            for file in files:
+                try:
+                    image = Image.open(file)
+                    image_path = f"image_{uuid.uuid4().hex}.png"
+                    image.save(image_path)
+                    user_content.append({"type": "image_path", "image_path": image_path})
+                except Exception as e:
+                    user_content.append({"type": "text", "text": f"Could not open file {file}"})
+            messages.append({"role": "user", "content": user_content})
+    else:
+        # Text-only branch.
+        messages.append({"role": "user", "content": [{"type": "text", "text": text}]})
+    
+    # Show a progress bar before generating.
+    yield progress_bar_html("Processing with Mistral")
+    
+    # Set up sampling parameters.
+    sampling_params = SamplingParams(
+        max_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k
     )
+    # Run the chat (synchronously) using vllm.
+    outputs = llm.chat(messages, sampling_params=sampling_params)
+    final_response = outputs[0].outputs[0].text
+
+    # Simulate streaming output by chunking the result.
+    buffer = ""
+    chunk_size = 20  # number of characters per chunk
+    for i in range(0, len(final_response), chunk_size):
+        buffer = final_response[: i + chunk_size]
+        yield buffer
+        time.sleep(0.05)
+    return
+
+# -----------------------------------------------------------------------------
+# Gradio Interface Setup
+# -----------------------------------------------------------------------------
+
+demo = gr.ChatInterface(
+    fn=generate,
+    additional_inputs=[
+        gr.Slider(label="Max new tokens", minimum=1, maximum=1024, step=1, value=512),
+        gr.Slider(label="Temperature", minimum=0.05, maximum=2.0, step=0.05, value=0.15),
+        gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
+        gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
+    ],
+    examples=[
+        # Example with text only.
+        ["Explain the significance of today in the context of current events."],
+        # Example with image files (ensure you have valid image paths).
+        [{
+            "text": "Describe what you see in the image.",
+            "files": ["examples/3.jpg"]
+        }],
+        # Example with video file (ensure you have a valid video file).
+        [{
+            "text": "@video-infer Summarize the events shown in the video.",
+            "files": ["examples/sample_video.mp4"]
+        }],
+    ],
+    cache_examples=False,
+    type="messages",
+    description="# **Mistral Multimodal Chatbot** \nSupports text, image (by reference) and video inference. Use @video-infer in your query when providing a video.",
+    fill_height=True,
+    textbox=gr.MultimodalTextbox(
+        label="Query Input",
+        file_types=["image", "video"],
+        file_count="multiple",
+        placeholder="Enter your query here. Tag with @video-infer if using a video file."
+    ),
+    stop_btn="Stop Generation",
+    examples_per_page=3,
+)
 
 if __name__ == "__main__":
-    demo.queue(max_size=40).launch()
+    demo.queue(max_size=20).launch(share=True)