Update app.py

app.py

@@ -2,329 +2,268 @@ import os
 import random
 import uuid
 import json
+import time
+import asyncio
+from threading import Thread
 
 import gradio as gr
-import numpy as np
-from PIL import Image
 import spaces
 import torch
+import numpy as np
+from PIL import Image
+import cv2
 
-from diffusers import DiffusionPipeline
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler # EulerAncestralDiscreteScheduler not explicitly used but imported
-from typing import Tuple
-
-bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
-bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
-default_negative = os.getenv("default_negative","")
-
-def check_text(prompt, negative=""):
-    for i in bad_words:
-        if i in prompt:
-            return True
-    for i in bad_words_negative:
-        if i in negative:
-            return True
-    return False
-
-style_list = [
-    {
-        "name": "Photo",
-        "prompt": "cinematic photo {prompt}. 35mm photograph, film, bokeh, professional, 4k, highly detailed",
-        "negative_prompt": "drawing, painting, crayon, sketch, graphite, impressionist, noisy, blurry, soft, deformed, ugly",
-    },
-    {
-        "name": "Cinematic",
-        "prompt": "cinematic still {prompt}. emotional, harmonious, vignette, highly detailed, high budget, bokeh, cinemascope, moody, epic, gorgeous, film grain, grainy",
-        "negative_prompt": "anime, cartoon, graphic, text, painting, crayon, graphite, abstract, glitch, deformed, mutated, ugly, disfigured",
-    },
-    {
-        "name": "Anime",
-        "prompt": "anime artwork {prompt}. anime style, key visual, vibrant, studio anime, highly detailed",
-        "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast",
-    },
-    {
-        "name": "3D Model",
-        "prompt": "professional 3d model {prompt}. octane render, highly detailed, volumetric, dramatic lighting",
-        "negative_prompt": "ugly, deformed, noisy, low poly, blurry, painting",
-    },
-    {
-        "name": "(No style)",
-        "prompt": "{prompt}",
-        "negative_prompt": "",
-    },
-]
-
-DESCRIPTION = """## 
-
-
-"""
-
-styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
-STYLE_NAMES = list(styles.keys())
-DEFAULT_STYLE_NAME = "Photo"
-
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    if not negative:
-        negative = ""
-    return p.replace("{prompt}", positive), n + negative
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
+from transformers import (
+    Qwen2_5_VLForConditionalGeneration,
+    AutoProcessor,
+    TextIteratorStreamer,
+)
+from transformers.image_utils import load_image
 
-MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
-MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "2048"))
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
+# Constants for text generation
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-NUM_IMAGES_PER_PROMPT = 1
-
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "SG161222/RealVisXL_V5.0_Lightning",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        add_watermarker=False,
-        variant="fp16"
-    )
-    pipe2 = StableDiffusionXLPipeline.from_pretrained(
-        "SG161222/RealVisXL_V4.0_Lightning",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-        add_watermarker=False,
-        variant="fp16"
-    )
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-        pipe2.enable_model_cpu_offload()
+# Load Cosmos-Reason1-7B
+MODEL_ID_M = "nvidia/Cosmos-Reason1-7B"
+processor_m = AutoProcessor.from_pretrained(MODEL_ID_M, trust_remote_code=True)
+model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_M,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+# Load DocScope
+MODEL_ID_X = "prithivMLmods/docscopeOCR-7B-050425-exp"
+processor_x = AutoProcessor.from_pretrained(MODEL_ID_X, trust_remote_code=True)
+model_x = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_X,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+# Load Relaxed
+MODEL_ID_Z = "Ertugrul/Qwen2.5-VL-7B-Captioner-Relaxed"
+processor_z = AutoProcessor.from_pretrained(MODEL_ID_Z, trust_remote_code=True)
+model_z = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    MODEL_ID_Z,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to(device).eval()
+
+def downsample_video(video_path):
+    """
+    Downsamples the video to evenly spaced frames.
+    Each frame is returned as a PIL image along with its timestamp.
+    """
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
+    for i in frame_indices:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, image = vidcap.read()
+        if success:
+            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
+
+@spaces.GPU
+def generate_image(model_name: str, text: str, image: Image.Image,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.6,
+                   top_p: float = 0.9,
+                   top_k: int = 50,
+                   repetition_penalty: float = 1.2):
+    """
+    Generates responses using the selected model for image input.
+    """
+    if model_name == "Cosmos-Reason1-7B":
+        processor = processor_m
+        model = model_m
+    elif model_name == "docscopeOCR-7B-050425-exp":
+        processor = processor_x
+        model = model_x
+    elif model_name == "Captioner-7B":
+        processor = processor_z
+        model = model_z
     else:
-        pipe.to(device)
-        pipe2.to(device)
-        print("Loaded on Device!")
-
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
-
-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
-
-@spaces.GPU(duration=30)
-@torch.no_grad()
-def generate(
-    prompt: str,
-    negative_prompt: str = "",
-    use_negative_prompt: bool = False,
-    style: str = DEFAULT_STYLE_NAME,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    randomize_seed: bool = False,
-    use_resolution_binning: bool = True, # This parameter is not exposed in the UI by default
-    progress=gr.Progress(track_tqdm=True),
-):
-    if check_text(prompt, negative_prompt):
-        raise ValueError("Prompt contains restricted words.")
-
-    prompt, negative_prompt_from_style = apply_style(style, prompt, "") # Apply style positive first
-    
-    # Combine negative prompts
-    if use_negative_prompt:
-        final_negative_prompt = negative_prompt_from_style + " " + negative_prompt + " " + default_negative
+        yield "Invalid model selected."
+        return
+
+    if image is None:
+        yield "Please upload an image."
+        return
+
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": text},
+        ]
+    }]
+    prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(
+        text=[prompt_full],
+        images=[image],
+        return_tensors="pt",
+        padding=True,
+        truncation=False,
+        max_length=MAX_INPUT_TOKEN_LENGTH
+    ).to(device)
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        time.sleep(0.01)
+        yield buffer
+
+@spaces.GPU
+def generate_video(model_name: str, text: str, video_path: str,
+                   max_new_tokens: int = 1024,
+                   temperature: float = 0.6,
+                   top_p: float = 0.9,
+                   top_k: int = 50,
+                   repetition_penalty: float = 1.2):
+    """
+    Generates responses using the selected model for video input.
+    """
+    if model_name == "Cosmos-Reason1-7B":
+        processor = processor_m
+        model = model_m
+    elif model_name == "docscopeOCR-7B-050425-exp":
+        processor = processor_x
+        model = model_x
+    elif model_name == "Captioner-7B":
+        processor = processor_z
+        model = model_z
     else:
-        final_negative_prompt = negative_prompt_from_style + " " + default_negative
-    final_negative_prompt = final_negative_prompt.strip()
-
-
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device=device).manual_seed(seed) # Ensure generator is on the correct device
-
-    options = {
-        "prompt": prompt,
-        "negative_prompt": final_negative_prompt,
-        "width": width,
-        "height": height,
-        "guidance_scale": guidance_scale,
-        "num_inference_steps": 25, # This is hardcoded, UI slider for steps is not connected
-        "generator": generator,
-        "num_images_per_prompt": NUM_IMAGES_PER_PROMPT, # UI slider for images is not connected to this
-        # "use_resolution_binning": use_resolution_binning, # This was in original code, but not defined. Diffusers handles it.
-        "output_type": "pil",
+        yield "Invalid model selected."
+        return
+
+    if video_path is None:
+        yield "Please upload a video."
+        return
+
+    frames = downsample_video(video_path)
+    messages = [
+        {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
+        {"role": "user", "content": [{"type": "text", "text": text}]}
+    ]
+    for frame in frames:
+        image, timestamp = frame
+        messages[1]["content"].append({"type": "text", "text": f"Frame {timestamp}:"})
+        messages[1]["content"].append({"type": "image", "image": image})
+    inputs = processor.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_dict=True,
+        return_tensors="pt",
+        truncation=False,
+        max_length=MAX_INPUT_TOKEN_LENGTH
+    ).to(device)
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = {
+        **inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_new_tokens,
+        "do_sample": True,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "repetition_penalty": repetition_penalty,
     }
-    
-    # If on CPU, ensure generator is for CPU
-    if device.type == 'cpu':
-        generator = torch.Generator(device='cpu').manual_seed(seed)
-        options["generator"] = generator
-
-    images = []
-    if 'pipe' in globals(): # Check if pipes are loaded (i.e. on GPU)
-         images.extend(pipe(**options).images)
-         images.extend(pipe2(**options).images)
-    else: # Fallback for CPU or if pipes are not loaded (though the DESCRIPTION warns about CPU)
-        # This part would need a CPU-compatible pipeline if one isn't loaded.
-        # For now, it will likely error if pipe/pipe2 aren't available.
-        # Or, we can return a placeholder or raise a specific error.
-        # To prevent errors if running without GPU and models didn't load:
-        placeholder_image = Image.new('RGB', (width, height), color = 'grey')
-        draw = ImageDraw.Draw(placeholder_image)
-        draw.text((10, 10), "GPU models not loaded. Cannot generate image.", fill=(255,0,0))
-        images.append(placeholder_image)
-
-
-    image_paths = [save_image(img) for img in images]
-    return image_paths, seed
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        time.sleep(0.01)
+        yield buffer
+
+# Define examples for image and video inference
+image_examples = [
+    ["type out the messy hand-writing as accurately as you can.", "images/1.jpg"],
+    ["count the number of birds and explain the scene in detail.", "images/2.jpg"]
+]
 
-examples = [
-    "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
+video_examples = [
+    ["give the highlights of the movie scene video.", "videos/1.mp4"],
+    ["explain the advertisement in detail.", "videos/2.mp4"]
 ]
 
-css = '''
-.gradio-container {
-    max-width: 590px !important; /* Existing style */
-    margin: 0 auto !important; /* Existing style */
-}
-h1 {
-    text-align: center; /* Existing style */
+css = """
+.submit-btn {
+    background-color: #2980b9 !important;
+    color: white !important;
 }
-footer {
-    visibility: hidden; /* Existing style */
+.submit-btn:hover {
+    background-color: #3498db !important;
 }
-'''
-with gr.Blocks(css=css) as demo:
-    gr.Markdown(DESCRIPTION)
-    with gr.Row():
-        prompt = gr.Text(
-            show_label=False,
-            max_lines=1,
-            placeholder="Enter your prompt",
-            container=False,
-        )
-        run_button = gr.Button("Run", scale=0, variant="primary")
-    result = gr.Gallery(label="Result", columns=1, preview=True) # columns=1 for single image below each other if multiple
+"""
 
-    with gr.Accordion("Advanced options", open=False):
-        style_selection = gr.Dropdown( # MODIFIED: Was gr.Radio, moved into accordion
-            label="Image Style",
-            choices=STYLE_NAMES,
-            value=DEFAULT_STYLE_NAME,
-            interactive=True,
-            show_label=True,
-            container=True,
-        )
-        use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=True, visible=True)
-        negative_prompt = gr.Text(
-            label="Negative prompt",
-            max_lines=1,
-            placeholder="Enter a negative prompt (appended to style's negative)",
-            value="(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck",
-            visible=True,
-        )
-        # Note: num_inference_steps and num_images_per_prompt sliders are defined in UI
-        # but not wired to the generate function's parameters that control these aspects.
-        # Keeping them as is, per "Don't alter the remaining functionality".
-        with gr.Row():
-            num_inference_steps = gr.Slider( # This UI element is not connected to the backend
-                label="Steps (Not Connected)",
-                minimum=10,
-                maximum=60,
-                step=1,
-                value=20, # Default value in UI
-            )
-        with gr.Row():
-            num_images_per_prompt = gr.Slider( # This UI element is not connected to the backend
-                label="Images (Not Connected)",
-                minimum=1,
-                maximum=4,
-                step=1,
-                value=2, # Default value in UI (backend NUM_IMAGES_PER_PROMPT is 1, resulting in 2 total)
-            )
-        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=MAX_IMAGE_SIZE, # Use MAX_IMAGE_SIZE
-                step=8,
-                value=1024,
-            )
-            height = gr.Slider(
-                label="Height",
-                minimum=512,
-                maximum=MAX_IMAGE_SIZE, # Use MAX_IMAGE_SIZE
-                step=8,
-                value=1024,
-            )
-        with gr.Row():
-            guidance_scale = gr.Slider(
-                label="Guidance Scale",
-                minimum=0.1,
-                maximum=20.0,
-                step=0.1,
-                value=3.0,
+# Create the Gradio Interface
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    gr.Markdown("# **VisionScope-R2**")
+    with gr.Row():
+        with gr.Column():
+            with gr.Tabs():
+                with gr.TabItem("Image Inference"):
+                    image_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+                    image_upload = gr.Image(type="pil", label="Image")
+                    image_submit = gr.Button("Submit", elem_classes="submit-btn")
+                    gr.Examples(
+                        examples=image_examples,
+                        inputs=[image_query, image_upload]
+                    )
+                with gr.TabItem("Video Inference"):
+                    video_query = gr.Textbox(label="Query Input", placeholder="Enter your query here...")
+                    video_upload = gr.Video(label="Video")
+                    video_submit = gr.Button("Submit", elem_classes="submit-btn")
+                    gr.Examples(
+                        examples=video_examples,
+                        inputs=[video_query, video_upload]
+                    )
+            with gr.Accordion("Advanced options", open=False):
+                max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS)
+                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6)
+                top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9)
+                top_k = gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50)
+                repetition_penalty = gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2)
+        with gr.Column():
+            output = gr.Textbox(label="Output", interactive=False, lines=2, scale=2)
+            model_choice = gr.Radio(
+                choices=["Cosmos-Reason1-7B", "docscopeOCR-7B-050425-exp", "Captioner-7B"],
+                label="Select Model",
+                value="Cosmos-Reason1-7B"
             )
-
-    # Original style_selection gr.Row has been removed from here.
-
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=[result, seed], # seed output is good for reproducibility
-        fn=generate,
-        cache_examples=CACHE_EXAMPLES,
+            
+            gr.Markdown("**Model Info**")
+            gr.Markdown("⤷ [Cosmos-Reason1-7B](https://huggingface.co/nvidia/Cosmos-Reason1-7B): understand physical common sense and generate appropriate embodied decisions.")
+            gr.Markdown("⤷ [docscopeOCR-7B-050425-exp](https://huggingface.co/prithivMLmods/docscopeOCR-7B-050425-exp): optimized for document-level optical character recognition, long-context vision-language understanding.")
+            gr.Markdown("⤷ [Captioner-Relaxed-7B](https://huggingface.co/Ertugrul/Qwen2.5-VL-7B-Captioner-Relaxed): build with hand-curated dataset for text-to-image models, providing significantly more detailed descriptions or captions of given images.")
+            
+    image_submit.click(
+        fn=generate_image,
+        inputs=[model_choice, image_query, image_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output
     )
-
-    use_negative_prompt.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_negative_prompt,
-        outputs=negative_prompt,
-        api_name=False,
-    )
-
-    gr.on(
-        triggers=[
-            prompt.submit,
-            negative_prompt.submit, # Allow submitting negative prompt to trigger run
-            run_button.click,
-        ],
-        fn=generate,
-        inputs=[
-            prompt,
-            negative_prompt,
-            use_negative_prompt,
-            style_selection, # style_selection is correctly in inputs
-            seed,
-            width,
-            height,
-            guidance_scale,
-            randomize_seed,
-        ],
-        outputs=[result, seed],
-        api_name="run",
+    video_submit.click(
+        fn=generate_video,
+        inputs=[model_choice, video_query, video_upload, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
+        outputs=output
     )
 
 if __name__ == "__main__":
-
-    from PIL import ImageDraw # Add ImageDraw import for CPU placeholder
-
-    demo.queue(max_size=20).launch(ssr_mode=True, show_error=True, share=True)
+    demo.queue(max_size=30).launch(share=True, mcp_server=True, ssr_mode=False, show_error=True)