import os import random import uuid import json import time import asyncio from threading import Thread import gradio as gr import spaces import torch import numpy as np from PIL import Image import edge_tts import cv2 from transformers import ( AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, Qwen2VLForConditionalGeneration, AutoProcessor, ) from transformers.image_utils import load_image from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler 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() # Updated TTS voices list (all voices) TTS_VOICES = [ "af-ZA-AdriNeural", "af-ZA-WillemNeural", "am-ET-AmehaNeural", "am-ET-MekdesNeural", "ar-AE-FatimaNeural", "ar-AE-HamdanNeural", "ar-BH-LailaNeural", "ar-BH-MajedNeural", "ar-DZ-AminaNeural", "ar-DZ-IsmaelNeural", "ar-EG-SalmaNeural", "ar-EG-OmarNeural", "ar-IQ-LanaNeural", "ar-IQ-BassamNeural", "ar-JO-SanaNeural", "ar-JO-TaimNeural", "ar-KW-NouraNeural", "ar-KW-FahedNeural", "ar-LB-LaylaNeural", "ar-LB-RamiNeural", "ar-LY-ImanNeural", "ar-LY-OmarNeural", "ar-MA-MounaNeural", "ar-MA-JamalNeural", "ar-OM-AyshaNeural", "ar-OM-AbdullahNeural", "ar-QA-AmalNeural", "ar-QA-MoazNeural", "ar-SA-ZariyahNeural", "ar-SA-HamedNeural", "ar-SY-AmanyNeural", "ar-SY-LaithNeural", "ar-TN-ReemNeural", "ar-TN-SeifNeural", "ar-YE-MaryamNeural", "ar-YE-SalehNeural", "az-AZ-BabekNeural", "az-AZ-BanuNeural", "bg-BG-BorislavNeural", "bg-BG-KalinaNeural", "bn-BD-NabanitaNeural", "bn-BD-PradeepNeural", "bn-IN-TanishaNeural", "bn-IN-SwapanNeural", "bs-BA-GoranNeural", "bs-BA-VesnaNeural", "ca-ES-JoanaNeural", "ca-ES-AlbaNeural", "ca-ES-EnricNeural", "cs-CZ-AntoninNeural", "cs-CZ-VlastaNeural", "cy-GB-NiaNeural", "cy-GB-AledNeural", "da-DK-ChristelNeural", "da-DK-JeppeNeural", "de-AT-IngridNeural", "de-AT-JonasNeural", "de-CH-LeniNeural", "de-CH-JanNeural", "de-DE-KatjaNeural", "de-DE-ConradNeural", "el-GR-AthinaNeural", "el-GR-NestorasNeural", "en-AU-AnnetteNeural", "en-AU-MichaelNeural", "en-CA-ClaraNeural", "en-CA-LiamNeural", "en-GB-SoniaNeural", "en-GB-RyanNeural", "en-GH-EsiNeural", "en-GH-KwameNeural", "en-HK-YanNeural", "en-HK-TrevorNeural", "en-IE-EmilyNeural", "en-IE-ConnorNeural", "en-IN-NeerjaNeural", "en-IN-PrabhasNeural", "en-KE-ChantelleNeural", "en-KE-ChilembaNeural", "en-NG-EzinneNeural", "en-NG-AbechiNeural", "en-NZ-MollyNeural", "en-NZ-MitchellNeural", "en-PH-RosaNeural", "en-PH-JamesNeural", "en-SG-LunaNeural", "en-SG-WayneNeural", "en-TZ-ImaniNeural", "en-TZ-DaudiNeural", "en-US-JennyNeural", "en-US-GuyNeural", "en-ZA-LeahNeural", "en-ZA-LukeNeural", "es-AR-ElenaNeural", "es-AR-TomasNeural", "es-BO-SofiaNeural", "es-BO-MarceloNeural", "es-CL-CatalinaNeural", "es-CL-LorenzoNeural", "es-CO-SalomeNeural", "es-CO-GonzaloNeural", "es-CR-MariaNeural", "es-CR-JuanNeural", "es-CU-BelkysNeural", "es-CU-ManuelNeural", "es-DO-RamonaNeural", "es-DO-EmilioNeural", "es-EC-AndreaNeural", "es-EC-LuisNeural", "es-ES-ElviraNeural", "es-ES-AlvaroNeural", "es-GQ-TeresaNeural", "es-GQ-JavierNeural", "es-GT-MartaNeural", "es-GT-AndresNeural", "es-HN-KarlaNeural", "es-HN-CarlosNeural", "es-MX-DaliaNeural", "es-MX-JorgeNeural", "es-NI-YolandaNeural", "es-NI-FedericoNeural", "es-PA-MargaritaNeural", "es-PA-RobertoNeural", "es-PE-CamilaNeural", "es-PE-AlexNeural", "es-PR-KarinaNeural", "es-PR-VictorNeural", "es-PY-TaniaNeural", "es-PY-MarioNeural", "es-SV-LorenaNeural", "es-SV-RodrigoNeural", "es-US-SaraNeural", "es-US-AlonsoNeural", "es-UY-ValentinaNeural", "es-UY-MateoNeural", "es-VE-PaolaNeural", "es-VE-SebastianNeural", "et-EE-AnuNeural", "et-EE-KertNeural", "eu-ES-AinhoaNeural", "eu-ES-AnderNeural", "fa-IR-DilaraNeural", "fa-IR-FaridNeural", "fi-FI-NooraNeural", "fi-FI-HarriNeural", "fil-PH-BlessicaNeural", "fil-PH-AngeloNeural", "fr-BE-CharlineNeural", "fr-BE-GerardNeural", "fr-CA-SylvieNeural", "fr-CA-AntoineNeural", "fr-CH-ArianeNeural", "fr-CH-GuillaumeNeural", "fr-FR-DeniseNeural", "fr-FR-HenriNeural", "ga-IE-OrlaNeural", "ga-IE-ColmNeural", "gl-ES-SoniaNeural", "gl-ES-XiaoqiangNeural", "gu-IN-DhwaniNeural", "gu-IN-NiranjanNeural", "ha-NG-AishaNeural", "ha-NG-YusufNeural", "he-IL-HilaNeural", "he-IL-AvriNeural", "hi-IN-SwaraNeural", "hi-IN-MadhurNeural", "hr-HR-GabrijelaNeural", "hr-HR-SreckoNeural", "hu-HU-NoemiNeural", "hu-HU-TamasNeural", "hy-AM-AnushNeural", "hy-AM-HaykNeural", "id-ID-ArdiNeural", "id-ID-GadisNeural", "ig-NG-AdaNeural", "ig-NG-EzeNeural", "is-IS-GudrunNeural", "is-IS-GunnarNeural", "it-IT-ElsaNeural", "it-IT-DiegoNeural", "ja-JP-NanamiNeural", "ja-JP-KeitaNeural", "jv-ID-DianNeural", "jv-ID-GustiNeural", "ka-GE-EkaNeural", # ... (truncated for brevity; include all voices as needed) ] 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 # Environment variables and parameters for Stable Diffusion XL (left in case needed in the future) MODEL_ID_SD = os.getenv("MODEL_VAL_PATH") # SDXL Model repository path via env variable MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "4096")) USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1" ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1" BATCH_SIZE = int(os.getenv("BATCH_SIZE", "1")) # For batched image generation # Load the SDXL pipeline (not used in the current configuration) sd_pipe = StableDiffusionXLPipeline.from_pretrained( MODEL_ID_SD, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, use_safetensors=True, add_watermarker=False, ).to(device) sd_pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(sd_pipe.scheduler.config) if torch.cuda.is_available(): sd_pipe.text_encoder = sd_pipe.text_encoder.half() if USE_TORCH_COMPILE: sd_pipe.compile() if ENABLE_CPU_OFFLOAD: sd_pipe.enable_model_cpu_offload() 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.""" 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 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'''
{label}
''' def downsample_video(video_path): """ Downsamples the video to 10 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(duration=60, enable_queue=True) def generate_image_fn( prompt: str, negative_prompt: str = "", use_negative_prompt: bool = False, seed: int = 1, width: int = 1024, height: int = 1024, guidance_scale: float = 3, num_inference_steps: int = 25, randomize_seed: bool = False, use_resolution_binning: bool = True, num_images: int = 1, progress=gr.Progress(track_tqdm=True), ): """(Image generation function is preserved but not called in the current configuration)""" seed = int(randomize_seed_fn(seed, randomize_seed)) generator = torch.Generator(device=device).manual_seed(seed) options = { "prompt": [prompt] * num_images, "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None, "width": width, "height": height, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps, "generator": generator, "output_type": "pil", } if use_resolution_binning: options["use_resolution_binning"] = True images = [] for i in range(0, num_images, BATCH_SIZE): batch_options = options.copy() batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE] if "negative_prompt" in batch_options and batch_options["negative_prompt"] is not None: batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE] if device.type == "cuda": with torch.autocast("cuda", dtype=torch.float16): outputs = sd_pipe(**batch_options) else: outputs = sd_pipe(**batch_options) images.extend(outputs.images) image_paths = [save_image(img) for img in images] return image_paths, seed @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, convert_to_speech: bool = False, tts_rate: float = 1.0, tts_voice: str = "en-US-JennyNeural", ): """ Generates chatbot responses with support for multimodal input and TTS conversion. When files (images or videos) are provided, Qwen2VL is used. Otherwise, the FastThink-0.5B text model is used. After generating the response, if convert_to_speech is True the text is passed to the TTS function. """ text = input_dict["text"].strip() files = input_dict.get("files", []) # Determine which branch to use: multimodal (if files provided) or text-only. if files: # Process uploaded files as images (or videos) if len(files) > 1: images = [load_image(image) for image in files] else: images = [load_image(files[0])] messages = [{ "role": "user", "content": [ *[{"type": "image", "image": image} for image in images], {"type": "text", "text": text}, ] }] prompt_full = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = processor(text=[prompt_full], 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() buffer = "" yield progress_bar_html("Processing multimodal input...") for new_text in streamer: buffer += new_text buffer = buffer.replace("<|im_end|>", "") time.sleep(0.01) yield buffer final_response = buffer else: conversation = clean_chat_history(chat_history) conversation.append({"role": "user", "content": text}) 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() outputs = [] yield progress_bar_html("Processing text...") for new_text in streamer: outputs.append(new_text) yield "".join(outputs) final_response = "".join(outputs) # Yield the final text response. yield final_response # If TTS conversion is enabled, log the message and generate speech. if convert_to_speech: print("Generate Response to Generate Speech") # Here tts_rate can be used to adjust parameters if needed. output_file = asyncio.run(text_to_speech(final_response, tts_voice)) yield gr.Audio(output_file, autoplay=True) with gr.Blocks() as demo: with gr.Sidebar(): gr.Markdown("# TTS Conversion") tts_rate_slider = gr.Slider(label="TTS Rate", minimum=0.5, maximum=2.0, step=0.1, value=1.0) tts_voice_radio = gr.Radio(choices=TTS_VOICES, label="Choose TTS Voice", value="en-US-JennyNeural") convert_to_speech_checkbox = gr.Checkbox(label="Convert to Speech", value=False) chat_interface = 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), # Pass TTS parameters to the generate function. convert_to_speech_checkbox, tts_rate_slider, tts_voice_radio, ], examples=[ ["Write the Python Program for Array Rotation"], [{"text": "Summarize the letter", "files": ["examples/1.png"]}], [{"text": "Describe the Ad", "files": ["examples/coca.mp4"]}], [{"text": "Summarize the event in video", "files": ["examples/sky.mp4"]}], [{"text": "Describe the video", "files": ["examples/Missing.mp4"]}], ["Who is Nikola Tesla, and why did he die?"], [{"text": "Extract JSON from the image", "files": ["examples/document.jpg"]}], ["What causes rainbows to form?"], ], cache_examples=False, type="messages", description="# **QwQ Edge: Multimodal (image upload uses Qwen2-VL) with TTS conversion**", fill_height=True, textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image", "video"], file_count="multiple", placeholder="Enter text or upload files"), stop_btn="Stop Generation", multimodal=True, ) if __name__ == "__main__": demo.queue(max_size=20).launch(share=True)