Update app.py

app.py

@@ -4,6 +4,7 @@ import uuid
 import json
 import time
 import asyncio
+import re
 from threading import Thread
 
 import gradio as gr
@@ -12,106 +13,22 @@ import torch
 import numpy as np
 from PIL import Image
 import edge_tts
+import subprocess
 
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-    Qwen2_5_VLForConditionalGeneration,
-    AutoProcessor,
+# Install flash-attn with our environment flag (if needed)
+subprocess.run(
+    'pip install flash-attn --no-build-isolation',
+    env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"},
+    shell=True
 )
-from transformers.image_utils import load_image
-from diffusers import DiffusionPipeline
-
-DESCRIPTION = "# Flux.1 Realism 🥖"
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
-
-css = '''
-h1 {
-  text-align: center;
-  display: block;
-}
-
-#duplicate-button {
-  margin: auto;
-  color: #fff;
-  background: #1565c0;
-  border-radius: 100vh;
-}
-'''
-
-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")
-
-# Load text-only model and tokenizer
-model_id = "prithivMLmods/FastThink-0.5B-Tiny"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-)
-model.eval()
-
-TTS_VOICES = [
-    "en-US-JennyNeural",  # @tts1
-    "en-US-GuyNeural",    # @tts2
-]
-
-# Load multimodal Qwen model & processor
-MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct" 
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-model_m = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to("cuda").eval()
-
-async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
-    """Convert text to speech using Edge TTS and save as MP3"""
-    communicate = edge_tts.Communicate(text, voice)
-    await communicate.save(output_file)
-    return output_file
-
-def clean_chat_history(chat_history):
-    """
-    Filter out any chat entries whose "content" is not a string.
-    This helps prevent errors when concatenating previous messages.
-    """
-    cleaned = []
-    for msg in chat_history:
-        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
-            cleaned.append(msg)
-    return cleaned
-
-def progress_bar_html(label: str) -> str:
-    """
-    Returns an HTML snippet for a thin progress bar with a label.
-    The progress bar is styled as a dark red animated 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: #FFC0CB; border-radius: 2px; overflow: hidden;">
-        <div style="width: 100%; height: 100%; background-color: #FF69B4; animation: loading 1.5s linear infinite;"></div>
-    </div>
-</div>
-<style>
-@keyframes loading {{
-    0% {{ transform: translateX(-100%); }}
-    100% {{ transform: translateX(100%); }}
-}}
-</style>
-    '''
 
+# -------------------------------
 # FLUX.1 IMAGE GENERATION SETUP
+# -------------------------------
 MAX_SEED = np.iinfo(np.int32).max
 
 def save_image(img: Image.Image) -> str:
-    """Save a PIL image with a unique filename and return the path."""
+    """Save a PIL image with a unique filename and return its path."""
     unique_name = str(uuid.uuid4()) + ".png"
     img.save(unique_name)
     return unique_name
@@ -121,7 +38,9 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-# Initialize Flux.1 pipeline
+# Load Flux.1 pipeline and LoRA weights
+from diffusers import DiffusionPipeline
+
 base_model = "black-forest-labs/FLUX.1-dev"
 pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
 lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
@@ -148,14 +67,14 @@ style_list = [
         "prompt": "{prompt}",
     },
 ]
-styles = {k["name"]: k["prompt"] for k in style_list}
+styles = {s["name"]: s["prompt"] for s in style_list}
 DEFAULT_STYLE_NAME = "3840 x 2160"
 STYLE_NAMES = list(styles.keys())
 
 def apply_style(style_name: str, positive: str) -> str:
     return styles.get(style_name, styles[DEFAULT_STYLE_NAME]).replace("{prompt}", positive)
 
-@spaces.GPU
+@spaces.GPU(duration=60, enable_queue=True)
 def generate_image_flux(
     prompt: str,
     seed: int = 0,
@@ -166,7 +85,7 @@ def generate_image_flux(
     style_name: str = DEFAULT_STYLE_NAME,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate images using the Flux.1 pipeline with style prompts."""
+    """Generate an image using the Flux.1 pipeline with a chosen style."""
     seed = int(randomize_seed_fn(seed, randomize_seed))
     positive_prompt = apply_style(style_name, prompt)
     if trigger_word:
@@ -183,29 +102,72 @@ def generate_image_flux(
     image_paths = [save_image(img) for img in images]
     return image_paths, seed
 
-# CHAT GENERATION FUNCTION (TEXT & MULTIMODAL)
+# -------------------------------
+# SMOLVLM2 SETUP (Default Text/Multimodal Model)
+# -------------------------------
+from transformers import AutoProcessor, AutoModelForImageTextToText, TextIteratorStreamer
+# Load the SmolVLM2 processor and model with flash attention enabled.
+smol_processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct")
+smol_model = AutoModelForImageTextToText.from_pretrained(
+    "HuggingFaceTB/SmolVLM2-2.2B-Instruct", 
+    _attn_implementation="flash_attention_2",
+    torch_dtype=torch.bfloat16
+).to("cuda:0")
+
+# -------------------------------
+# UTILITY FUNCTIONS
+# -------------------------------
+def progress_bar_html(label: str) -> str:
+    """
+    Returns an HTML snippet for an animated progress bar with a given label.
+    """
+    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: #FFC0CB; border-radius: 2px; overflow: hidden;">
+        <div style="width: 100%; height: 100%; background-color: #FF69B4; animation: loading 1.5s linear infinite;"></div>
+    </div>
+</div>
+<style>
+@keyframes loading {{
+    0% {{ transform: translateX(-100%); }}
+    100% {{ transform: translateX(100%); }}
+}}
+</style>
+    '''
+
+# TTS voices (if using TTS commands)
+TTS_VOICES = [
+    "en-US-JennyNeural",  # @tts1
+    "en-US-GuyNeural",    # @tts2
+]
+
+async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
+    """Convert text to speech using Edge TTS and save the output as MP3."""
+    communicate = edge_tts.Communicate(text, voice)
+    await communicate.save(output_file)
+    return output_file
+
+# -------------------------------
+# CHAT / MULTIMODAL GENERATION FUNCTION
+# -------------------------------
 @spaces.GPU
 def generate(
     input_dict: dict,
     chat_history: list[dict],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
+    max_tokens: int = 200,
 ):
     """
-    Generates chatbot responses with support for multimodal input, TTS, and image generation.
+    Generates chatbot responses using SmolVLM2 by default—with support for multimodal inputs and TTS.
     Special commands:
-      - "@tts1" or "@tts2": triggers text-to-speech.
       - "@image": triggers image generation using the Flux.1 pipeline.
+      - "@tts1" or "@tts2": triggers text-to-speech after generation.
     """
     text = input_dict["text"]
     files = input_dict.get("files", [])
 
-    # If the text begins with "@image", use Flux for image generation.
+    # If the query starts with "@image", use Flux.1 to generate an image.
     if text.strip().lower().startswith("@image"):
-        # Remove the "@image" tag and use the remainder as the prompt.
         prompt = text[len("@image"):].strip()
         yield progress_bar_html("Hold Tight Generating Flux.1 Image")
         image_paths, used_seed = generate_image_flux(
@@ -219,109 +181,144 @@ def generate(
             progress=gr.Progress(track_tqdm=True),
         )
         yield gr.Image(image_paths[0])
-        return  # Exit early after image generation.
+        return
 
-    # Check if a TTS command is issued.
+    # Handle TTS commands if present.
     tts_prefix = "@tts"
     is_tts = any(text.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 3))
-    voice_index = next((i for i in range(1, 3) if text.strip().lower().startswith(f"{tts_prefix}{i}")), None)
-    
-    if is_tts and voice_index:
-        voice = TTS_VOICES[voice_index - 1]
-        text = text.replace(f"{tts_prefix}{voice_index}", "").strip()
-        # Clear previous chat history for a fresh TTS request.
-        conversation = [{"role": "user", "content": text}]
-    else:
-        voice = None
-        # Remove any stray @tts tags and build the conversation history.
-        text = text.replace(tts_prefix, "").strip()
-        conversation = clean_chat_history(chat_history)
-        conversation.append({"role": "user", "content": text})
+    voice = None
+    if is_tts:
+        voice_index = next((i for i in range(1, 3) if text.strip().lower().startswith(f"{tts_prefix}{i}")), None)
+        if voice_index:
+            voice = TTS_VOICES[voice_index - 1]
+            text = text.replace(f"{tts_prefix}{voice_index}", "").strip()
 
-    # Handle multimodal input if files are provided.
-    if files:
-        if len(files) > 1:
-            images = [load_image(image) for image in files]
-        elif len(files) == 1:
-            images = [load_image(files[0])]
-        else:
-            images = []
-        messages = [{
-            "role": "user",
-            "content": [
-                *[{"type": "image", "image": image} for image in images],
-                {"type": "text", "text": text},
-            ]
-        }]
-        prompt = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[prompt], images=images, return_tensors="pt", padding=True).to("cuda")
-        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_m.generate, kwargs=generation_kwargs)
-        thread.start()
+    # Now use SmolVLM2 for chat/multimodal text generation.
+    yield "Processing with SmolVLM2"
 
-        buffer = ""
-        yield progress_bar_html("Thinking...")
-        for new_text in streamer:
-            buffer += new_text
-            buffer = buffer.replace("<|im_end|>", "")
-            time.sleep(0.01)
-            yield buffer
+    # Build conversation messages based on input and history.
+    user_content = []
+    media_queue = []
+    if chat_history == []:
+        text = text.strip()
+        for file in files:
+            if file.endswith((".png", ".jpg", ".jpeg", ".gif", ".bmp")):
+                media_queue.append({"type": "image", "path": file})
+            elif file.endswith((".mp4", ".mov", ".avi", ".mkv", ".flv")):
+                media_queue.append({"type": "video", "path": file})
+        if "<image>" in text or "<video>" in text:
+            parts = re.split(r'(<image>|<video>)', text)
+            for part in parts:
+                if part == "<image>" and media_queue:
+                    user_content.append(media_queue.pop(0))
+                elif part == "<video>" and media_queue:
+                    user_content.append(media_queue.pop(0))
+                elif part.strip():
+                    user_content.append({"type": "text", "text": part.strip()})
+        else:
+            user_content.append({"type": "text", "text": text})
+            for media in media_queue:
+                user_content.append(media)
+        resulting_messages = [{"role": "user", "content": user_content}]
     else:
-        input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
-        if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-            input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-            gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-        input_ids = input_ids.to(model.device)
-        streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {
-            "input_ids": input_ids,
-            "streamer": streamer,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "top_p": top_p,
-            "top_k": top_k,
-            "temperature": temperature,
-            "num_beams": 1,
-            "repetition_penalty": repetition_penalty,
-        }
-        t = Thread(target=model.generate, kwargs=generation_kwargs)
-        t.start()
+        resulting_messages = []
+        user_content = []
+        media_queue = []
+        for hist in chat_history:
+            if hist["role"] == "user" and isinstance(hist["content"], tuple):
+                file_name = hist["content"][0]
+                if file_name.endswith((".png", ".jpg", ".jpeg")):
+                    media_queue.append({"type": "image", "path": file_name})
+                elif file_name.endswith(".mp4"):
+                    media_queue.append({"type": "video", "path": file_name})
+        for hist in chat_history:
+            if hist["role"] == "user" and isinstance(hist["content"], str):
+                txt = hist["content"]
+                parts = re.split(r'(<image>|<video>)', txt)
+                for part in parts:
+                    if part == "<image>" and media_queue:
+                        user_content.append(media_queue.pop(0))
+                    elif part == "<video>" and media_queue:
+                        user_content.append(media_queue.pop(0))
+                    elif part.strip():
+                        user_content.append({"type": "text", "text": part.strip()})
+            elif hist["role"] == "assistant":
+                resulting_messages.append({
+                    "role": "user",
+                    "content": user_content
+                })
+                resulting_messages.append({
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": hist["content"]}]
+                })
+                user_content = []
+        if not resulting_messages:
+            resulting_messages = [{"role": "user", "content": user_content}]
+
+    if text == "" and not files:
+        yield "Please input a query and optionally image(s)."
+        return
+    if text == "" and files:
+        yield "Please input a text query along with the image(s)."
+        return
 
-        outputs = []
-        yield progress_bar_html("Thinking...")
-        for new_text in streamer:
-            outputs.append(new_text)
-            yield "".join(outputs)
+    print("resulting_messages", resulting_messages)
+    inputs = smol_processor.apply_chat_template(
+        resulting_messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    )
+    inputs = inputs.to(smol_model.device)
 
-        final_response = "".join(outputs)
-        yield final_response
+    streamer = TextIteratorStreamer(smol_processor, skip_prompt=True, skip_special_tokens=True)
+    generation_args = dict(inputs, streamer=streamer, max_new_tokens=max_tokens)
+    thread = Thread(target=smol_model.generate, kwargs=generation_args)
+    thread.start()
 
-        # If TTS was requested, convert the final response to speech.
-        if is_tts and voice:
-            output_file = asyncio.run(text_to_speech(final_response, voice))
-            yield gr.Audio(output_file, autoplay=True)
+    yield "..."
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        time.sleep(0.01)
+        yield buffer
 
+    if is_tts and voice:
+        final_response = buffer
+        output_file = asyncio.run(text_to_speech(final_response, voice))
+        yield gr.Audio(output_file, autoplay=True)
+
+# -------------------------------
 # GRADIO CHAT INTERFACE
+# -------------------------------
+DESCRIPTION = "# Flux.1 Realism 🥖 + SmolVLM2 Chat"
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>⚠️Running on CPU, this may not work as expected.</p>"
+
+css = '''
+h1 {
+  text-align: center;
+  display: block;
+}
+#duplicate-button {
+  margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
+}
+'''
+
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
-        gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
-        gr.Slider(label="Top-p (nucleus sampling)", 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),
-        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
+        gr.Slider(minimum=100, maximum=500, step=50, value=200, label="Max Tokens"),
     ],
     examples=[
-        [{"text": "@image Chocolate dripping from a donut against a yellow background, in the style of hyper-realistic 8K"}],
-        [{"text": "@image Super Realism, High-resolution photograph, woman, UHD, photorealistic, shot on a Sony A7III --chaos 20 --ar 1:2 --style raw --stylize 250"}],
-        [{"text": "Extract JSON from the image", "files": ["examples/document.jpg"]}],
-        [{"text": "@image 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"}],
-        [{"text": "@image 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."}],
-        [{"text": "Python Program for Array Rotation"}],
-        [{"text": "@tts1 Who is Nikola Tesla, and why did he die?"}],
-        [{"text": "Summarize the letter", "files": ["examples/1.png"]}],
-        [{"text": "@tts2 What causes rainbows to form?"}],
+        [{"text": "@image A futuristic cityscape at dusk in hyper-realistic 8K"}],
+        [{"text": "Describe this image.", "files": ["example_images/mosque.jpg"]}],
+        [{"text": "What does this document say?", "files": ["example_images/document.jpg"]}],
+        [{"text": "@tts1 Explain the weather patterns shown in this diagram.", "files": ["example_images/examples_weather_events.png"]}],
     ],
     cache_examples=False,
     type="messages",
@@ -330,9 +327,9 @@ demo = gr.ChatInterface(
     fill_height=True,
     textbox=gr.MultimodalTextbox(
         label="Query Input", 
-        file_types=["image"], 
+        file_types=["image", ".mp4"], 
         file_count="multiple",  
-        placeholder="‎ @image-flux.1 image gen, @tts1, @tts2-voices, default [text, vision]"
+        placeholder="Type text and/or upload media. Use '@image' for Flux.1 image gen, '@tts1' or '@tts2' for TTS."
     ),
     stop_btn="Stop Generation",
     multimodal=True,