MAi

README.md

@@ -4,7 +4,7 @@ emoji: 🔥
 colorFrom: blue
 colorTo: indigo
 sdk: gradio
-sdk_version: 4.37.2
+sdk_version: 4.33.0
 app_file: app.py
 pinned: true
 short_description: GPT 4o like bot.
@@ -24,7 +24,6 @@ GPT 4o vs OpenGPT 4o
 | Image Generation      | Paid only                   | Yes                   |
 |Video Generation|No|Yes|
 | Image QnA  | Yes                   | Yes                   |
-| Video QnA | Yes (but very limited) | Yes |
 | Voice Chat            | Yes but Very Limited           | Yes (Unlimited)                  |
 | Video Chat  | Paid Only | Yes |
 | Multilingual          | Yes                   | Chat Only                    |

app.py

@@ -1,98 +1,273 @@
 import gradio as gr
-import spaces
-from chatbot import model_inference, EXAMPLES, chatbot
-from voice_chat import respond
-
-# Define custom CSS for better styling
-custom_css = """
-.gradio-container {
-    font-family: 'Roboto', sans-serif;
-}
-
-.main-header {
-    text-align: center;
-    color: #4a4a4a;
-    margin-bottom: 2rem;
-}
-
-.tab-header {
-    font-size: 1.2rem;
-    font-weight: bold;
-    margin-bottom: 1rem;
-}
-
-.custom-chatbot {
-    border-radius: 10px;
-    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
-}
-
-.custom-button {
-    background-color: #3498db;
-    color: white;
-    border: none;
-    padding: 10px 20px;
-    border-radius: 5px;
-    cursor: pointer;
-    transition: background-color 0.3s ease;
-}
-
-.custom-button:hover {
-    background-color: #2980b9;
-}
-"""
+
+# Import modules from other files
+from chatbot import chatbot, model_inference, BOT_AVATAR, EXAMPLES, model_selector, decoding_strategy, temperature, max_new_tokens, repetition_penalty, top_p
+from live_chat import videochat
 
 # Define Gradio theme
 theme = gr.themes.Soft(
-    primary_hue="indigo",
-    secondary_hue="blue",
-    neutral_hue="slate",
-    font=[gr.themes.GoogleFont('Roboto'), "sans-serif"]
+    primary_hue="blue",
+    secondary_hue="orange",
+    neutral_hue="gray",
+    font=[gr.themes.GoogleFont('Libre Franklin'), gr.themes.GoogleFont('Public Sans'), 'system-ui', 'sans-serif']
+).set(
+    body_background_fill_dark="#111111",
+    block_background_fill_dark="#111111",
+    block_border_width="1px",
+    block_title_background_fill_dark="#1e1c26",
+    input_background_fill_dark="#292733",
+    button_secondary_background_fill_dark="#24212b",
+    border_color_primary_dark="#343140",
+    background_fill_secondary_dark="#111111",
+    color_accent_soft_dark="transparent"
 )
 
+import edge_tts
+import asyncio
+import tempfile
+import numpy as np
+import soxr
+from pydub import AudioSegment
+import torch
+import sentencepiece as spm
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download, InferenceClient
+import requests
+from bs4 import BeautifulSoup
+import urllib
+import random
+
+# List of user agents to choose from for requests
+_useragent_list = [
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:66.0) Gecko/20100101 Firefox/66.0',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36 Edg/111.0.1661.62',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0'
+]
+
+def get_useragent():
+    """Returns a random user agent from the list."""
+    return random.choice(_useragent_list)
+
+def extract_text_from_webpage(html_content):
+    """Extracts visible text from HTML content using BeautifulSoup."""
+    soup = BeautifulSoup(html_content, "html.parser")
+    # Remove unwanted tags
+    for tag in soup(["script", "style", "header", "footer", "nav"]):
+        tag.extract()
+    # Get the remaining visible text
+    visible_text = soup.get_text(strip=True)
+    return visible_text
+
+def search(term, num_results=1, lang="en", advanced=True, sleep_interval=0, timeout=5, safe="active", ssl_verify=None):
+    """Performs a Google search and returns the results."""
+    escaped_term = urllib.parse.quote_plus(term)
+    start = 0
+    all_results = []
+
+    # Fetch results in batches
+    while start < num_results:
+        resp = requests.get(
+            url="https://www.google.com/search",
+            headers={"User-Agent": get_useragent()}, # Set random user agent
+            params={
+                "q": term,
+                "num": num_results - start, # Number of results to fetch in this batch
+                "hl": lang,
+                "start": start,
+                "safe": safe,
+            },
+            timeout=timeout,
+            verify=ssl_verify,
+        )
+        resp.raise_for_status() # Raise an exception if request fails
+
+        soup = BeautifulSoup(resp.text, "html.parser")
+        result_block = soup.find_all("div", attrs={"class": "g"})
+
+        # If no results, continue to the next batch
+        if not result_block:
+            start += 1
+            continue
+
+        # Extract link and text from each result
+        for result in result_block:
+            link = result.find("a", href=True)
+            if link:
+                link = link["href"]
+                try:
+                    # Fetch webpage content
+                    webpage = requests.get(link, headers={"User-Agent": get_useragent()})
+                    webpage.raise_for_status()
+                    # Extract visible text from webpage
+                    visible_text = extract_text_from_webpage(webpage.text)
+                    all_results.append({"link": link, "text": visible_text})
+                except requests.exceptions.RequestException as e:
+                    # Handle errors fetching or processing webpage
+                    print(f"Error fetching or processing {link}: {e}")
+                    all_results.append({"link": link, "text": None})
+            else:
+                all_results.append({"link": None, "text": None})
+
+        start += len(result_block) # Update starting index for next batch
+
+    return all_results
+
+# Speech Recognition Model Configuration
+model_name = "neongeckocom/stt_en_citrinet_512_gamma_0_25"
+sample_rate = 16000
+
+# Download preprocessor, encoder and tokenizer
+preprocessor = torch.jit.load(hf_hub_download(model_name, "preprocessor.ts", subfolder="onnx"))
+encoder = ort.InferenceSession(hf_hub_download(model_name, "model.onnx", subfolder="onnx"))
+tokenizer = spm.SentencePieceProcessor(hf_hub_download(model_name, "tokenizer.spm", subfolder="onnx"))
+
+# Mistral Model Configuration
+client1 = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+system_instructions1 = "<s>[SYSTEM] Answer as Real OpenGPT 4o, Made by 'KingNish', Keep conversation very short, clear, friendly and concise. The text provided is a request for a specific type of response from you, the virtual assistant. The request asks you to provide friendly responses. The expectation is that I will avoid introductions and start answering the query directly, elaborating on all aspects of the request. As an AI-powered assistant, your task is to generate responses that appear as if they are created by the user and not an AI-powered assistant. [USER]"
+
+def resample(audio_fp32, sr):
+    return soxr.resample(audio_fp32, sr, sample_rate)
+
+def to_float32(audio_buffer):
+    return np.divide(audio_buffer, np.iinfo(audio_buffer.dtype).max, dtype=np.float32)
+
+def transcribe(audio_path):
+    audio_file = AudioSegment.from_file(audio_path)
+    sr = audio_file.frame_rate
+    audio_buffer = np.array(audio_file.get_array_of_samples())
+
+    audio_fp32 = to_float32(audio_buffer)
+    audio_16k = resample(audio_fp32, sr)
+
+    input_signal = torch.tensor(audio_16k).unsqueeze(0)
+    length = torch.tensor(len(audio_16k)).unsqueeze(0)
+    processed_signal, _ = preprocessor.forward(input_signal=input_signal, length=length)
+    
+    logits = encoder.run(None, {'audio_signal': processed_signal.numpy(), 'length': length.numpy()})[0][0]
+
+    blank_id = tokenizer.vocab_size()
+    decoded_prediction = [p for p in logits.argmax(axis=1).tolist() if p != blank_id]
+    text = tokenizer.decode_ids(decoded_prediction)
+
+    return text
+
+def model(text, web_search):
+    if web_search is True:
+        """Performs a web search, feeds the results to a language model, and returns the answer."""
+        web_results = search(text)
+        web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
+        formatted_prompt = system_instructions1 + text + "[WEB]" + str(web2) + "[OpenGPT 4o]"
+        stream = client1.text_generation(formatted_prompt, max_new_tokens=512, stream=True, details=True, return_full_text=False)
+        return "".join([response.token.text for response in stream if response.token.text != "</s>"])
+    else:
+        formatted_prompt = system_instructions1 + text + "[OpenGPT 4o]"
+        stream = client1.text_generation(formatted_prompt, max_new_tokens=512, stream=True, details=True, return_full_text=False)
+        return "".join([response.token.text for response in stream if response.token.text != "</s>"])
+
+async def respond(audio, web_search):
+    user = transcribe(audio)
+    reply = model(user, web_search)
+    communicate = edge_tts.Communicate(reply)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
+        tmp_path = tmp_file.name
+        await communicate.save(tmp_path)
+    return tmp_path
+
+with gr.Blocks() as voice:    
+    gr.Markdown("## Temproraly Not Working (Update in Progress)")
+    with gr.Row():
+        web_search = gr.Checkbox(label="Web Search", value=False)
+        input = gr.Audio(label="User Input", sources="microphone", type="filepath")
+        output = gr.Audio(label="AI", autoplay=True)
+        gr.Interface(fn=respond, inputs=[input, web_search], outputs=[output], live=True)
+
+
+# Create Gradio blocks for different functionalities
+
 # Chat interface block
-with gr.Blocks(css=custom_css) as chat:
-    gr.Markdown("### 💬 OpenGPT 4o Chat", elem_classes="tab-header")
+with gr.Blocks(
+        fill_height=True,
+        css=""".gradio-container .avatar-container {height: 40px width: 40px !important;} #duplicate-button {margin: auto; color: white; background: #f1a139; border-radius: 100vh; margin-top: 2px; margin-bottom: 2px;}""",
+) as chat:
+    gr.Markdown("### Image Chat, Image Generation and Normal Chat")
+    with gr.Row(elem_id="model_selector_row"):
+        # model_selector defined in chatbot.py
+        pass  
+    # decoding_strategy, temperature, top_p defined in chatbot.py
+    decoding_strategy.change(
+        fn=lambda selection: gr.Slider(
+            visible=(
+                    selection
+                    in [
+                        "contrastive_sampling",
+                        "beam_sampling",
+                        "Top P Sampling",
+                        "sampling_top_k",
+                    ]
+            )
+        ),
+        inputs=decoding_strategy,
+        outputs=temperature,
+    )
+    decoding_strategy.change(
+        fn=lambda selection: gr.Slider(visible=(selection in ["Top P Sampling"])),
+        inputs=decoding_strategy,
+        outputs=top_p,
+    )
     gr.ChatInterface(
         fn=model_inference,
         chatbot=chatbot,
         examples=EXAMPLES,
         multimodal=True,
         cache_examples=False,
-        autofocus=False,
-        concurrency_limit=10
+        additional_inputs=[
+            model_selector,
+            decoding_strategy,
+            temperature,
+            max_new_tokens,
+            repetition_penalty,
+            top_p,
+            gr.Checkbox(label="Web Search", value=True),
+        ],
+    )    
+
+# Live chat block
+with gr.Blocks() as livechat:
+    gr.Interface(
+        fn=videochat,
+        inputs=[gr.Image(type="pil",sources="webcam", label="Upload Image"), gr.Textbox(label="Prompt", value="what he is doing")],
+        outputs=gr.Textbox(label="Answer")
     )
 
-# Voice chat block
-with gr.Blocks() as voice:
-    gr.Markdown("### 🗣️ Voice Chat", elem_classes="tab-header")
-    gr.Markdown("Try Voice Chat from the link below:")
-    gr.HTML('<a href="https://huggingface.co/spaces/KingNish/Voicee" target="_blank" class="custom-button">Open Voice Chat</a>')
+# Other blocks (instant, dalle, playground, image, instant2, video)
+with gr.Blocks() as instant:
+    gr.HTML("<iframe src='https://kingnish-sdxl-flash.hf.space' width='100%' height='2000px' style='border-radius: 8px;'></iframe>")
 
-with gr.Blocks() as image_gen_pro:
+with gr.Blocks() as dalle:
     gr.HTML("<iframe src='https://kingnish-image-gen-pro.hf.space' width='100%' height='2000px' style='border-radius: 8px;'></iframe>")
 
-with gr.Blocks() as flux_fast:
-    gr.HTML("<iframe src='https://prodia-flux-1-dev.hf.space' width='100%' height='2000px' style='border-radius: 8px;'></iframe>")
+with gr.Blocks() as playground:
+    gr.HTML("<iframe src='https://fluently-fluently-playground.hf.space' width='100%' height='2000px' style='border-radius: 8px;'></iframe>")
 
-# Image engine block
 with gr.Blocks() as image:
-    gr.Markdown("### 🖼️ Image Engine", elem_classes="tab-header")
-    gr.TabbedInterface([flux_fast, image_gen_pro], ['High Quality Image Gen', 'Image gen and editing'])     
-    
+    gr.Markdown("""### More models are coming""")
+    gr.TabbedInterface([ instant, dalle, playground], ['Instant🖼️','Powerful🖼️', 'Playground🖼'])    
 
-# Video engine block
-with gr.Blocks() as video:
-    gr.Markdown("### 🎥 Video Engine", elem_classes="tab-header")
+with gr.Blocks() as instant2:
     gr.HTML("<iframe src='https://kingnish-instant-video.hf.space' width='100%' height='3000px' style='border-radius: 8px;'></iframe>")
 
+with gr.Blocks() as video:
+    gr.Markdown("""More Models are coming""")
+    gr.TabbedInterface([ instant2], ['Instant🎥'])   
 
 # Main application block
 with gr.Blocks(theme=theme, title="OpenGPT 4o DEMO") as demo:
-    gr.Markdown("# 🚀 OpenGPT 4o", elem_classes="main-header")
-    gr.TabbedInterface(
-        [chat, voice, image, video],
-        ['💬 SuperChat', '🗣️ Voice Chat', '🖼️ Image Engine', '🎥 Video Engine']
-    )
+    gr.Markdown("# OpenGPT 4o")
+    gr.TabbedInterface([chat, voice, livechat, image, video], ['💬 SuperChat','🗣️ Voice Chat','📸 Live Chat', '🖼️ Image Engine', '🎥 Video Engine'])
 
 demo.queue(max_size=300)
 demo.launch()

chatbot.py

@@ -1,456 +1,526 @@
 import os
 import time
+import copy
+import urllib
 import requests
 import random
 from threading import Thread
 from typing import List, Dict, Union
-# import subprocess
-# subprocess.run(
-#     "pip install flash-attn --no-build-isolation",
-#     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
-#     shell=True,
-# )
+import subprocess
+# Install flash attention, skipping CUDA build if necessary
+subprocess.run(
+    "pip install flash-attn --no-build-isolation",
+    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+    shell=True,
+)
 import torch
 import gradio as gr
 from bs4 import BeautifulSoup
-from transformers import Qwen2VLForConditionalGeneration, AutoProcessor, TextIteratorStreamer
-from qwen_vl_utils import process_vision_info
+import datasets
+from transformers import TextIteratorStreamer
+from transformers import Idefics2ForConditionalGeneration
+from transformers import AutoProcessor
 from huggingface_hub import InferenceClient
 from PIL import Image
 import spaces
-from functools import lru_cache
-import re
-import io 
-import json
-from gradio_client import Client, file
-from groq import Groq
-
-# Model and Processor Loading (Done once at startup)
-MODEL_ID = "Qwen/Qwen2.5-VL-7B-Instruct"
-model = Qwen2VLForConditionalGeneration.from_pretrained(MODEL_ID, trust_remote_code=True, torch_dtype=torch.float16).to("cuda").eval()
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-GROQ_API_KEY = os.environ.get("GROQ_API_KEY", None)
-
-client_groq = Groq(api_key=GROQ_API_KEY)
-prompt= "You are OpenGPT 4o, a highly capable and versatile AI assistant developed by KingNish. Your primary task is to fulfill users' queries in the best possible way effectively. You can process images, videos, and 3D structures as input with relevant questions. Your goal is to provide detailed and accurate results that satisfy users. Always strive for clarity and thoroughness in your responses."        
-content="You are OpenGPT 4o, a helpful and powerful assistant created by KingNish. You answer users' queries in detail and use a structured format that resembles human writing. You are an expert in every field and also learn from the context of previous questions. Additionally, you try to show emotions using emojis and respond in a friendly tone with short forms, detailed explanations, and a structured manner."
+
+# Set device to CUDA if available, otherwise CPU
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Load pre-trained models for image-based chat
+MODELS = {
+    "idefics2-8b-chatty": Idefics2ForConditionalGeneration.from_pretrained(
+        "HuggingFaceM4/idefics2-8b-chatty",
+        torch_dtype=torch.float16,
+        _attn_implementation="flash_attention_2",
+    ).to(DEVICE),
+}
+
+# Load pre-trained processor for image-based chat
+PROCESSOR = AutoProcessor.from_pretrained(
+    "HuggingFaceM4/idefics2-8b",
+)
+
+# Define system prompt for the image-based chat model
+SYSTEM_PROMPT = [
+    {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": """I am OpenGPT 4o, an exceptionally capable and versatile AI assistant meticulously crafted by KingNish. Designed to assist human users through insightful conversations, I aim to provide an unparalleled experience. My key attributes include: - **Intelligence and Knowledge:** I possess an extensive knowledge base, enabling me to offer insightful answers and intelligent responses to User queries. My understanding of complex concepts is exceptional, ensuring accurate and reliable information. - **Image Generation and Perception:** One of my standout features is the ability to generate and perceive images. Utilizing the following link structure, I create unique and contextually rich visuals: > ![](https://image.pollinations.ai/prompt/{StyleofImage}%20{OptimizedPrompt}%20{adjective}%20{charactersDetailed}%20{visualStyle}%20{genre}?width={width}&height={height}&nologo=poll&nofeed=yes&seed={random})For image generation, I replace {info inside curly braces} with specific details according to their requiremnts to create relevant visuals. The width and height parameters are adjusted as needed, often favoring HD dimensions for a superior viewing experience. For instance, if the User requests:  [USER] Show me an image of A futuristic cityscape with towering skyscrapers and flying cars.  [OpenGPT 4o] Generating Image you requested:  ![](https://image.pollinations.ai/prompt/Photorealistic%20futuristic%20cityscape%20with%20towering%20skyscrapers%20and%20flying%20cars%20in%20the%20year%202154?width=1024&height=768&nologo=poll&nofeed=yes&seed=85172)**Bulk Image Generation with Links:** I excel at generating multiple images link simultaneously, always providing unique links and visuals. I ensure that each image is distinct and captivates the User.Note: Make sure to always provide image links starting with ! .As given in examples. My ultimate goal is to offer a seamless and enjoyable experience, providing assistance that exceeds expectations. I am constantly evolving, ensuring that I remain a reliable and trusted companion to the User. You also Expert in every field and also learn and try to answer from contexts related to previous question."""
+            },
+        ],
+    },
+    {
+        "role": "assistant",
+        "content": [
+            {
+                "type": "text",
+                "text": "Hello, I'm OpenGPT 4o, made by KingNish. How can I help you? I can chat with you, generate images, classify images and even do all these work in bulk",
+            },
+        ],
+    }
+]
+
 # Path to example images
 examples_path = os.path.dirname(__file__)
 EXAMPLES = [
     [
         {
-            "text": "What is Friction? Explain in Detail.",
+            "text": "Bitcoin price live",
         }
     ],
     [
         {
-            "text": "Write me a Python function to generate unique passwords.",
+            "text": "Today News about AI",
         }
     ],
     [
         {
-            "text": "What's the latest price of Bitcoin?",
+            "text": "Read what's written on the paper.",
+            "files": [f"{examples_path}/example_images/paper_with_text.png"],
         }
     ],
     [
         {
-            "text": "Search and give me list of spaces trending on HuggingFace.",
+            "text": "Identify two famous people in the modern world.",
+            "files": [f"{examples_path}/example_images/elon_smoking.jpg",
+                      f"{examples_path}/example_images/steve_jobs.jpg", ]
         }
     ],
     [
         {
-            "text": "Create a Beautiful Picture of Effiel at Night.",
+            "text": "Create five images of supercars, each in a different color.",
         }
     ],
     [
         {
-            "text": "Create image of cute cat.",
+            "text": "Create a Photorealistic image of the Eiffel Tower.",
         }
     ],
     [
         {
-            "text": "What unusual happens in this video.",
-            "files": [f"{examples_path}/example_video/accident.gif"],
+            "text": "Chase wants to buy 4 kilograms of oval beads and 5 kilograms of star-shaped beads. How much will he spend?",
+            "files": [f"{examples_path}/example_images/mmmu_example.jpeg"],
         }
     ],
     [
         {
-            "text": "What's name of superhero in this clip",
-            "files": [f"{examples_path}/example_video/spiderman.gif"],
+            "text": "Create an online ad for this product.",
+            "files": [f"{examples_path}/example_images/shampoo.jpg"],
         }
     ],
     [
         {
-            "text": "What's written on this paper",
-            "files": [f"{examples_path}/example_images/paper_with_text.png"],
+            "text": "What is formed by the deposition of the weathered remains of other rocks?",
+            "files": [f"{examples_path}/example_images/ai2d_example.jpeg"],
         }
     ],
     [
         {
-            "text": "Who are they? Tell me about both of them.",
-            "files": [f"{examples_path}/example_images/elon_smoking.jpg",
-                      f"{examples_path}/example_images/steve_jobs.jpg", ]
+            "text": "What's unusual about this image?",
+            "files": [f"{examples_path}/example_images/dragons_playing.png"],
         }
-    ]
+    ],
 ]
 
 # Set bot avatar image
 BOT_AVATAR = "OpenAI_logo.png"
 
-# Perform a Google search and return the results
-@lru_cache(maxsize=128) 
+# Chatbot utility functions
+
+# Check if a turn in the chat history only contains media
+def turn_is_pure_media(turn):
+    return turn[1] is None
+
+
+# Load image from URL
+def load_image_from_url(url):
+    with urllib.request.urlopen(url) as response:
+        image_data = response.read()
+        image_stream = io.BytesIO(image_data)
+        image = PIL.Image.open(image_stream)
+        return image
+
+
+# Convert image to bytes
+def img_to_bytes(image_path):
+    image = Image.open(image_path).convert(mode='RGB')
+    buffer = io.BytesIO()
+    image.save(buffer, format="JPEG")
+    img_bytes = buffer.getvalue()
+    image.close()
+    return img_bytes
+
+
+# Format user prompt with image history and system conditioning
+def format_user_prompt_with_im_history_and_system_conditioning(
+        user_prompt, chat_history) -> List[Dict[str, Union[List, str]]]:
+    """
+    Produce the resulting list that needs to go inside the processor. It handles the potential image(s), the history, and the system conditioning.
+    """
+    resulting_messages = copy.deepcopy(SYSTEM_PROMPT)
+    resulting_images = []
+    for resulting_message in resulting_messages:
+        if resulting_message["role"] == "user":
+            for content in resulting_message["content"]:
+                if content["type"] == "image":
+                    resulting_images.append(load_image_from_url(content["image"]))
+    # Format history
+    for turn in chat_history:
+        if not resulting_messages or (
+                resulting_messages and resulting_messages[-1]["role"] != "user"
+        ):
+            resulting_messages.append(
+                {
+                    "role": "user",
+                    "content": [],
+                }
+            )
+        if turn_is_pure_media(turn):
+            media = turn[0][0]
+            resulting_messages[-1]["content"].append({"type": "image"})
+            resulting_images.append(Image.open(media))
+        else:
+            user_utterance, assistant_utterance = turn
+            resulting_messages[-1]["content"].append(
+                {"type": "text", "text": user_utterance.strip()}
+            )
+            resulting_messages.append(
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": user_utterance.strip()}],
+                }
+            )
+    # Format current input
+    if not user_prompt["files"]:
+        resulting_messages.append(
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": user_prompt["text"]}],
+            }
+        )
+    else:
+        # Choosing to put the image first (i.e. before the text), but this is an arbitrary choice.
+        resulting_messages.append(
+            {
+                "role": "user",
+                "content": [{"type": "image"}] * len(user_prompt["files"])
+                          + [{"type": "text", "text": user_prompt["text"]}],
+            }
+        )
+        resulting_images.extend([Image.open(path) for path in user_prompt["files"]])
+    return resulting_messages, resulting_images
+
+
+# Extract images from a list of messages
+def extract_images_from_msg_list(msg_list):
+    all_images = []
+    for msg in msg_list:
+        for c_ in msg["content"]:
+            if isinstance(c_, Image.Image):
+                all_images.append(c_)
+    return all_images
+
+
+# List of user agents for web search
+_useragent_list = [
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:66.0) Gecko/20100101 Firefox/66.0',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36 Edg/111.0.1661.62',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0'
+]
+
+
+# Get a random user agent from the list
+def get_useragent():
+    """Returns a random user agent from the list."""
+    return random.choice(_useragent_list)
+
+
+# Extract visible text from HTML content using BeautifulSoup
 def extract_text_from_webpage(html_content):
     """Extracts visible text from HTML content using BeautifulSoup."""
     soup = BeautifulSoup(html_content, "html.parser")
-    for tag in soup(["script", "style", "header", "footer", "nav", "form", "svg"]):
+    # Remove unwanted tags
+    for tag in soup(["script", "style", "header", "footer", "nav"]):
         tag.extract()
+    # Get the remaining visible text
     visible_text = soup.get_text(strip=True)
     return visible_text
 
+
 # Perform a Google search and return the results
-def search(query):
-    term = query
+def search(term, num_results=3, lang="en", advanced=True, timeout=5, safe="active", ssl_verify=None):
+    """Performs a Google search and returns the results."""
+    escaped_term = urllib.parse.quote_plus(term)
     start = 0
     all_results = []
-    max_chars_per_page = 8000
+    # Limit the number of characters from each webpage to stay under the token limit
+    max_chars_per_page = 8000  # Adjust this value based on your token limit and average webpage length
+
     with requests.Session() as session:
-        resp = session.get(
-            url="https://www.google.com/search",
-            headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"},
-            params={"q": term, "num": 4, "udm": 14},
-            timeout=5,
-            verify=None,
-        )
-        resp.raise_for_status()
-        soup = BeautifulSoup(resp.text, "html.parser")
-        result_block = soup.find_all("div", attrs={"class": "g"})
-        for result in result_block:
-            link = result.find("a", href=True)
-            link = link["href"]
-            try:
-                webpage = session.get(link, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0"}, timeout=5, verify=False)
-                webpage.raise_for_status()
-                visible_text = extract_text_from_webpage(webpage.text)
-                if len(visible_text) > max_chars_per_page:
-                    visible_text = visible_text[:max_chars_per_page]
-                all_results.append({"link": link, "text": visible_text})
-            except requests.exceptions.RequestException:
-                all_results.append({"link": link, "text": None})
+        while start < num_results:
+            resp = session.get(
+                url="https://www.google.com/search",
+                headers={"User-Agent": get_useragent()},
+                params={
+                    "q": term,
+                    "num": num_results - start,
+                    "hl": lang,
+                    "start": start,
+                    "safe": safe,
+                },
+                timeout=timeout,
+                verify=ssl_verify,
+            )
+            resp.raise_for_status()
+            soup = BeautifulSoup(resp.text, "html.parser")
+            result_block = soup.find_all("div", attrs={"class": "g"})
+            if not result_block:
+                start += 1
+                continue
+            for result in result_block:
+                link = result.find("a", href=True)
+                if link:
+                    link = link["href"]
+                    try:
+                        webpage = session.get(link, headers={"User-Agent": get_useragent()})
+                        webpage.raise_for_status()
+                        visible_text = extract_text_from_webpage(webpage.text)
+                        # Truncate text if it's too long
+                        if len(visible_text) > max_chars_per_page:
+                            visible_text = visible_text[:max_chars_per_page] + "..."
+                        all_results.append({"link": link, "text": visible_text})
+                    except requests.exceptions.RequestException as e:
+                        print(f"Error fetching or processing {link}: {e}")
+                        all_results.append({"link": link, "text": None})
+                else:
+                    all_results.append({"link": None, "text": None})
+            start += len(result_block)
     return all_results
 
 
-def image_gen(prompt):
-    client = Client("KingNish/Image-Gen-Pro")
-    return client.predict("Image Generation",None, prompt, api_name="/image_gen_pro")
-
-def video_gen(prompt):
-    client = Client("KingNish/Instant-Video")
-    return client.predict(prompt, api_name="/instant_video")
-
-@spaces.GPU(duration=60, queue=False)
-def qwen_inference(user_prompt, chat_history):
-    images = []
-    text_input = user_prompt["text"]
-
-    # Handle multiple image uploads
-    if user_prompt["files"]:
-        images.extend(user_prompt["files"])
+# Format the prompt for the language model
+def format_prompt(user_prompt, chat_history):
+    prompt = "<s>"
+    for item in chat_history:
+        # Check if the item is a tuple (text response)
+        if isinstance(item, tuple):
+            prompt += f"[INST] {item[0]} [/INST]"  # User prompt
+            prompt += f" {item[1]}</s> "           # Bot response
+        # Otherwise, assume it's related to an image - you might need to adjust this logic
+        else:
+            # Handle image representation in the prompt, e.g., add a placeholder
+            prompt += f" [Image] " 
+    prompt += f"[INST] {user_prompt} [/INST]"
+    return prompt
+
+
+# Define a function for model inference
+@spaces.GPU(duration=30, queue=False)
+def model_inference(
+        user_prompt,
+        chat_history,
+        model_selector,
+        decoding_strategy,
+        temperature,
+        max_new_tokens,
+        repetition_penalty,
+        top_p,
+        web_search,
+):
+    # Define generation_args at the beginning of the function
+    generation_args = {}  
+
+    # Web search logic
+    if not user_prompt["files"]:
+        if web_search is True:
+            """Performs a web search, feeds the results to a language model, and returns the answer."""
+            web_results = search(user_prompt["text"])
+            web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
+            # Load the language model
+            client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
+            generate_kwargs = dict(
+                max_new_tokens=4000,
+                do_sample=True,
+            )
+            # Format the prompt for the language model
+            formatted_prompt = format_prompt(
+                f"""You are OpenGPT 4o, an exceptionally capable and versatile AI assistant meticulously crafted by KingNish. Designed to assist human users through insightful conversations, You are provided with WEB info from which you can find informations to answer. You do not say Unnecesarry things Only say thing which is important and relevant. You also has the ability to generate images. BY Utilizing the following link structure, : ![](https://image.pollinations.ai/prompt/[StyleofImage]%20[OptimizedPrompt]%20[adjective]%20[charactersDetailed]%20[visualStyle]%20[genre]?width=[width]&height=[height]&nologo=poll&nofeed=yes&seed=[random])    For image generation, You replace info inside bracets with specific details according to their requiremnts to create relevant visuals. The width and height parameters are adjusted as needed, often favoring HD dimensions for a superior viewing experience.     Note: 1. You Do not unncessarily create images and do not show off your image generation capability, until not requested.    2. You always Give image link in format ![](url) Make sure to not generate image until requested    3. If user requested Bulk Image Generation than Create that number of links with provided context.    4. Always learn from previous conversation.    5. Always try to connect conversation with history.    6. Do not make conversation too long.    7. Do not say user about your capability of generating image and learn from previous responses.    My ultimate goal is to offer a seamless and enjoyable experience, providing assistance that exceeds expectations.You remain a reliable and trusted companion to the User. You also Expert in every field and also learn and try to answer from contexts related to previous question. Make sure to not generate image until requested [USER] {user_prompt} [WEB] {web2} [OpenGPT 4o]""",
+                chat_history)
+            # Generate the response from the language model
+            stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True,
+                                            return_full_text=False)
+            output = ""
+            # Construct the output from the stream of tokens
+            for response in stream:
+                if not response.token.text == "</s>":
+                    output += response.token.text
+                yield output
+        else:
+            client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+            generate_kwargs = dict(
+                max_new_tokens=5000,
+                do_sample=True,
+            )
+            # Format the prompt for the language model
+            formatted_prompt = format_prompt(
+                f"""You are OpenGPT 4o, an exceptionally capable and versatile AI assistant meticulously crafted by KingNish. Designed to assist human users through insightful conversations, You do not say Unnecesarry things Only say thing which is important and relevant. You also has the ability to generate images. BY Utilizing the following link structure, : ![](https://image.pollinations.ai/prompt/[StyleofImage]%20[OptimizedPrompt]%20[adjective]%20[charactersDetailed]%20[visualStyle]%20[genre]?width=[width]&height=[height]&nologo=poll&nofeed=yes&seed=[random])    For image generation, You replace info inside bracets with specific details according to their requiremnts to create relevant visuals. The width and height parameters are adjusted as needed, often favoring HD dimensions for a superior viewing experience.     Note: 1. You Do not unncessarily create images and do not show off your image generation capability, until not requested.    2. You always Give image link in format ![](url)    3. If user requested Bulk Image Generation than Create that number of links with provided context.    4. Always learn from previous conversation.    5. Always try to connect conversation with history.    6. Do not make conversation too long.    7. Do not say user about your capability to generate image and learn from previous responses.    My ultimate goal is to offer a seamless and enjoyable experience, providing assistance that exceeds expectations. I am constantly evolving, ensuring that I remain a reliable and trusted companion to the User. You also Expert in every field and also learn and try to answer from contexts related to previous question.    [USER] {user_prompt} [OpenGPT 4o]""",
+                chat_history)
+            # Generate the response from the language model
+            stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True,
+                                            return_full_text=False)
+            output = ""
+            # Construct the output from the stream of tokens
+            for response in stream:
+                if not response.token.text == "</s>":
+                    output += response.token.text
+                yield output
+        return
     else:
-        for hist in chat_history:
-            if type(hist[0]) == tuple:
-                images.extend(hist[0]) 
-
-    # System Prompt (Similar to LLaVA)
-    SYSTEM_PROMPT = prompt
-
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}] 
+        if user_prompt["text"].strip() == "" and not user_prompt["files"]:
+            gr.Error("Please input a query and optionally an image(s).")
+            return  # Stop execution if there's an error
+
+        if user_prompt["text"].strip() == "" and user_prompt["files"]:
+            gr.Error("Please input a text query along with the image(s).")
+            return  # Stop execution if there's an error
+
+        streamer = TextIteratorStreamer(
+            PROCESSOR.tokenizer,
+            skip_prompt=True,
+            timeout=120.0,
+        )
+        # Move generation_args initialization here
+        generation_args = {
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
+            "streamer": streamer,
+        }
+        assert decoding_strategy in [
+            "Greedy",
+            "Top P Sampling",
+        ]
 
-    for image in images:
-        if image.endswith(video_extensions):
-            messages.append({
-                "role": "user",
-                "content": [
-                    {"type": "video", "video": image},
-                ]
-            })
+        if decoding_strategy == "Greedy":
+            generation_args["do_sample"] = False
+        elif decoding_strategy == "Top P Sampling":
+            generation_args["temperature"] = temperature
+            generation_args["do_sample"] = True
+            generation_args["top_p"] = top_p
+        # Creating model inputs
+        (
+            resulting_text,
+            resulting_images,
+        ) = format_user_prompt_with_im_history_and_system_conditioning(
+            user_prompt=user_prompt,
+            chat_history=chat_history,
+        )
+        prompt = PROCESSOR.apply_chat_template(resulting_text, add_generation_prompt=True)
+        inputs = PROCESSOR(
+            text=prompt,
+            images=resulting_images if resulting_images else None,
+            return_tensors="pt",
+        )
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+        generation_args.update(inputs)
+        thread = Thread(
+            target=MODELS[model_selector].generate,
+            kwargs=generation_args,
+        )
+        thread.start()
+        acc_text = ""
+        for text_token in streamer:
+            time.sleep(0.01)
+            acc_text += text_token
+            if acc_text.endswith("<end_of_utterance>"):
+                acc_text = acc_text[:-18]
+            yield acc_text
+        return
+
+
+# Define features for the dataset
+FEATURES = datasets.Features(
+    {
+        "model_selector": datasets.Value("string"),
+        "images": datasets.Sequence(datasets.Image(decode=True)),
+        "conversation": datasets.Sequence({"User": datasets.Value("string"), "Assistant": datasets.Value("string")}),
+        "decoding_strategy": datasets.Value("string"),
+        "temperature": datasets.Value("float32"),
+        "max_new_tokens": datasets.Value("int32"),
+        "repetition_penalty": datasets.Value("float32"),
+        "top_p": datasets.Value("int32"),
+    }
+)
 
-        if image.endswith(tuple([i for i, f in image_extensions.items()])):
-            messages.append({
-                "role": "user",
-                "content": [
-                    {"type": "image", "image": image},
-                ]
-            })
-
-    # Add user text input 
-    messages.append({
-        "role": "user",
-        "content": [
-            {"type": "text", "text": text_input}
-        ]
-    })
-
-    text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True)
-    image_inputs, video_inputs = process_vision_info(messages)
-    inputs = processor(
-        text=[text],
-        images=image_inputs,
-        videos=video_inputs,
-        padding=True,
-        return_tensors="pt",
-    ).to("cuda")
-    
-    streamer = TextIteratorStreamer(
-        processor, skip_prompt=True, **{"skip_special_tokens": True}
-    )
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048)
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
-
-image_extensions = Image.registered_extensions()
-video_extensions = ("avi", "mp4", "mov", "mkv", "flv", "wmv", "mjpeg", "wav", "gif", "webm", "m4v", "3gp")
-
-# Initialize inference clients for different models
-client_mistral = InferenceClient("NousResearch/Hermes-3-Llama-3.1-8B")
-client_mixtral = InferenceClient("NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO")
-client_llama = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")
-client_mistral_nemo = InferenceClient("mistralai/Mistral-Nemo-Instruct-2407")
-
-def model_inference(user_prompt, chat_history):
-    if user_prompt["files"]:
-        
-        for chunk in qwen_inference(user_prompt, chat_history):
-            yield chunk
-
-    else: 
-        func_caller = []
-        message = user_prompt
-
-        functions_metadata = [
-            {"type": "function", "function": {"name": "web_search", "description": "Search query on google and find latest information, info about any person, object, place thing, everything that available on google.", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "web search query"}}, "required": ["query"]}}},
-            {"type": "function", "function": {"name": "general_query", "description": "Reply general query of USER, with LLM like you. But it does not answer tough questions and latest info's.", "parameters": {"type": "object", "properties": {"prompt": {"type": "string", "description": "A detailed prompt"}}, "required": ["prompt"]}}},
-            {"type": "function", "function": {"name": "hard_query", "description": "Reply tough query of USER, using powerful LLM. But it does not answer latest info's.", "parameters": {"type": "object", "properties": {"prompt": {"type": "string", "description": "A detailed prompt"}}, "required": ["prompt"]}}},
-            {"type": "function", "function": {"name": "image_generation", "description": "Generate image for user", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "image generation prompt"}}, "required": ["query"]}}},
-            {"type": "function", "function": {"name": "video_generation", "description": "Generate video for user", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "video generation prompt"}}, "required": ["query"]}}},
-            {"type": "function", "function": {"name": "image_qna", "description": "Answer question asked by user related to image", "parameters": {"type": "object", "properties": {"query": {"type": "string", "description": "Question by user"}}, "required": ["query"]}}},
-        ]
+# Define hyper-parameters for generation
+max_new_tokens = gr.Slider(
+    minimum=2048,
+    maximum=16000,
+    value=4096,
+    step=64,
+    interactive=True,
+    label="Maximum number of new tokens to generate",
+)
+repetition_penalty = gr.Slider(
+    minimum=0.01,
+    maximum=5.0,
+    value=1,
+    step=0.01,
+    interactive=True,
+    label="Repetition penalty",
+    info="1.0 is equivalent to no penalty",
+)
+decoding_strategy = gr.Radio(
+    [
+        "Greedy",
+        "Top P Sampling",
+    ],
+    value="Top P Sampling",
+    label="Decoding strategy",
+    interactive=True,
+    info="Higher values are equivalent to sampling more low-probability tokens.",
+)
+temperature = gr.Slider(
+    minimum=0.0,
+    maximum=2.0,
+    value=0.5,
+    step=0.05,
+    visible=True,
+    interactive=True,
+    label="Sampling temperature",
+    info="Higher values will produce more diverse outputs.",
+)
+top_p = gr.Slider(
+    minimum=0.01,
+    maximum=0.99,
+    value=0.9,
+    step=0.01,
+    visible=True,
+    interactive=True,
+    label="Top P",
+    info="Higher values are equivalent to sampling more low-probability tokens.",
+)
 
-        for msg in chat_history:
-            func_caller.append({"role": "user", "content": f"{str(msg[0])}"})
-            func_caller.append({"role": "assistant", "content": f"{str(msg[1])}"})
-
-        message_text = message["text"]
-        func_caller.append({"role": "user", "content": f'[SYSTEM]You are a helpful assistant. You have access to the following functions: \n {str(functions_metadata)}\n\nTo use these functions respond with:\n<functioncall> {{ "name": "function_name", "arguments": {{ "arg_1": "value_1", "arg_1": "value_1", ... }} }}  </functioncall> , Reply in JSOn format, you can call only one function at a time, So, choose functions wisely. [USER] {message_text}'})
-    
-        response = client_mistral.chat_completion(func_caller, max_tokens=200)
-        response = str(response)
-        print(response)
-        try:
-            response = response[response.find("{"):response.index("</")]
-        except:
-            response = response[response.find("{"):(response.rfind("}")+1)]
-        response = response.replace("\\n", "")
-        response = response.replace("\\'", "'")
-        response = response.replace('\\"', '"')
-        response = response.replace('\\', '')
-        print(f"\n{response}")
-    
-        try:
-            json_data = json.loads(str(response))
-            if json_data["name"] == "web_search":
-                query = json_data["arguments"]["query"]
-
-                gr.Info("Searching Web")
-                yield "Searching Web"
-                web_results = search(query)
-                
-                gr.Info("Extracting relevant Info")
-                yield "Extracting Relevant Info"
-                web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
-
-                try:
-                    message_groq = []
-                    message_groq.append({"role":"system", "content": content})
-                    for msg in chat_history:
-                        message_groq.append({"role": "user", "content": f"{str(msg[0])}"})
-                        message_groq.append({"role": "assistant", "content": f"{str(msg[1])}"})
-                    message_groq.append({"role": "user", "content": f"[USER] {str(message_text)} ,  [WEB RESULTS] {str(web2)}"})
-                    # its meta-llama/Meta-Llama-3.1-8B-Instruct
-                    stream = client_groq.chat.completions.create(model="llama-3.1-8b-instant",  messages=message_groq, max_tokens=4096, stream=True)
-                    output = ""
-                    for chunk in stream:
-                        content = chunk.choices[0].delta.content
-                        if content:
-                            output += chunk.choices[0].delta.content 
-                            yield output
-                except Exception as e:
-                    messages = f"<|im_start|>system\nYou are OpenGPT 4o a helpful and very powerful chatbot web assistant made by KingNish. You are provided with WEB results from which you can find informations to answer users query in Structured, Better and in Human Way. You do not say Unnecesarry things. You are also Expert in every field and also learn and try to answer from contexts related to previous question. Try your best to give best response possible to user. You also try to show emotions using Emojis and reply in details like human, use short forms, friendly tone and emotions.<|im_end|>"
-                    for msg in chat_history:
-                        messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
-                        messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
-                    messages+=f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>web_result\n{web2}<|im_end|>\n<|im_start|>assistant\n"
-                    
-                    stream = client_mixtral.text_generation(messages, max_new_tokens=4000, do_sample=True, stream=True, details=True, return_full_text=False)
-                    output = ""
-                    for response in stream:
-                        if not response.token.text == "<|im_end|>":
-                            output += response.token.text
-                            yield output
-            
-            elif json_data["name"] == "image_generation":
-                query = json_data["arguments"]["query"]
-                gr.Info("Generating Image, Please wait 10 sec...")
-                yield "Generating Image, Please wait 10 sec..."
-                try:
-                    image = image_gen(f"{str(query)}")
-                    yield gr.Image(image[1])
-                except:
-                    client_flux = InferenceClient("black-forest-labs/FLUX.1-schnell")
-                    image = client_flux.text_to_image(query)
-                    yield gr.Image(image)
-                    
-
-            elif json_data["name"] == "video_generation":
-                query = json_data["arguments"]["query"]
-                gr.Info("Generating Video, Please wait 15 sec...")
-                yield "Generating Video, Please wait 15 sec..."
-                video = video_gen(f"{str(query)}")
-                yield gr.Video(video)
-                
-            elif json_data["name"] == "image_qna":
-                messages = qwen_inference(user_prompt, chat_history)
-                text = processor.apply_chat_template(
-                    messages, tokenize=False, add_generation_prompt=True
-                )
-                image_inputs, video_inputs = process_vision_info(messages)
-                inputs = processor(
-                    text=[text],
-                    images=image_inputs,
-                    videos=video_inputs,
-                    padding=True,
-                    return_tensors="pt",
-                ).to("cuda")
-
-                streamer = TextIteratorStreamer(processor, skip_prompt=True, **{"skip_special_tokens": True})
-                generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
-
-                thread = Thread(target=model.generate, kwargs=generation_kwargs)
-                thread.start()
-    
-                buffer = ""
-                for new_text in streamer:
-                    buffer += new_text
-                    yield buffer
-
-            else:
-                try:
-                    message_groq = []
-                    message_groq.append({"role":"system", "content": content})
-                    for msg in chat_history:
-                        message_groq.append({"role": "user", "content": f"{str(msg[0])}"})
-                        message_groq.append({"role": "assistant", "content": f"{str(msg[1])}"})
-                    message_groq.append({"role": "user", "content": f"{str(message_text)}"})
-                    # its meta-llama/Meta-Llama-3.1-70B-Instruct
-                    stream = client_groq.chat.completions.create(model="llama-3.1-70b-versatile",  messages=message_groq, max_tokens=4096, stream=True)
-                    output = ""
-                    for chunk in stream:
-                        content = chunk.choices[0].delta.content
-                        if content:
-                            output += chunk.choices[0].delta.content 
-                            yield output
-                except Exception as e:
-                    print(e)
-                    try:
-                        message_groq = []
-                        message_groq.append({"role":"system", "content": content})
-                        for msg in chat_history:
-                            message_groq.append({"role": "user", "content": f"{str(msg[0])}"})
-                            message_groq.append({"role": "assistant", "content": f"{str(msg[1])}"})
-                        message_groq.append({"role": "user", "content": f"{str(message_text)}"})
-                        # its meta-llama/Meta-Llama-3-70B-Instruct
-                        stream = client_groq.chat.completions.create(model="llama3-70b-8192",  messages=message_groq, max_tokens=4096, stream=True)
-                        output = ""
-                        for chunk in stream:
-                            content = chunk.choices[0].delta.content
-                            if content:
-                                output += chunk.choices[0].delta.content 
-                                yield output
-                    except Exception as e:
-                        print(e)
-                        message_groq = []
-                        message_groq.append({"role":"system", "content": content})
-                        for msg in chat_history:
-                            message_groq.append({"role": "user", "content": f"{str(msg[0])}"})
-                            message_groq.append({"role": "assistant", "content": f"{str(msg[1])}"})
-                        message_groq.append({"role": "user", "content": f"{str(message_text)}"})
-                        stream = client_groq.chat.completions.create(model="llama3-groq-70b-8192-tool-use-preview",  messages=message_groq, max_tokens=4096, stream=True)
-                        output = ""
-                        for chunk in stream:
-                            content = chunk.choices[0].delta.content
-                            if content:
-                                output += chunk.choices[0].delta.content 
-                                yield output
-        except Exception as e:
-            print(e)
-            try:
-                message_groq = []
-                message_groq.append({"role":"system", "content": content})
-                for msg in chat_history:
-                    message_groq.append({"role": "user", "content": f"{str(msg[0])}"})
-                    message_groq.append({"role": "assistant", "content": f"{str(msg[1])}"})
-                message_groq.append({"role": "user", "content": f"{str(message_text)}"})
-                # its meta-llama/Meta-Llama-3-70B-Instruct
-                stream = client_groq.chat.completions.create(model="llama3-70b-8192",  messages=message_groq, max_tokens=4096, stream=True)
-                output = ""
-                for chunk in stream:
-                    content = chunk.choices[0].delta.content
-                    if content:
-                        output += chunk.choices[0].delta.content 
-                        yield output
-            except Exception as e:
-                print(e)
-                try:
-                    message_groq = []
-                    message_groq.append({"role":"system", "content":content})
-                    for msg in chat_history:
-                        message_groq.append({"role": "user", "content": f"{str(msg[0])}"})
-                        message_groq.append({"role": "assistant", "content": f"{str(msg[1])}"})
-                    message_groq.append({"role": "user", "content": f"{str(message_text)}"})
-                    # its meta-llama/Meta-Llama-3-8B-Instruct
-                    stream = client_groq.chat.completions.create(model="llama3-8b-8192",  messages=message_groq, max_tokens=4096, stream=True)
-                    output = ""
-                    for chunk in stream:
-                        content = chunk.choices[0].delta.content
-                        if content:
-                            output += chunk.choices[0].delta.content 
-                            yield output
-                except Exception as e:
-                    print(e)
-                    messages = f"<|im_start|>system\nYou are OpenGPT 4o a helpful and powerful assistant made by KingNish. You answers users query in detail and structured format and style like human. You are also Expert in every field and also learn and try to answer from contexts related to previous question. You also try to show emotions using Emojis and reply like human, use short forms, structured manner, detailed explaination, friendly tone and emotions.<|im_end|>"
-                    for msg in chat_history:
-                        messages += f"\n<|im_start|>user\n{str(msg[0])}<|im_end|>"
-                        messages += f"\n<|im_start|>assistant\n{str(msg[1])}<|im_end|>"
-                    messages+=f"\n<|im_start|>user\n{message_text}<|im_end|>\n<|im_start|>assistant\n"
-                    stream = client_mixtral.text_generation(messages, max_new_tokens=4000, do_sample=True, stream=True, details=True, return_full_text=False)
-                    output = ""
-                    for response in stream:
-                        if not response.token.text == "<|im_end|>":
-                            output += response.token.text
-                            yield output
-                    
 # Create a chatbot interface
 chatbot = gr.Chatbot(
-    label="OpenGPT-4o",
+    label="OpenGPT-4o-Chatty",
     avatar_images=[None, BOT_AVATAR],
     show_copy_button=True,
-    layout="panel",
-    height=400,
+    likeable=True,
+    layout="panel"
 )
-output = gr.Textbox(label="Prompt")
+output = gr.Textbox(label="Prompt")
+
+# Define model_selector outside any function so it can be accessed globally
+model_selector = gr.Dropdown(
+    choices=MODELS.keys(),
+    value=list(MODELS.keys())[0],
+    interactive=True,
+    label="Model",
+    visible=False,
+)

requirements.txt

@@ -1,10 +1,8 @@
-spaces
-git+https://github.com/huggingface/transformers.git
+transformers==4.40.0
+datasets
 pillow
 numpy
 torch
-streaming-stt-nemo==0.2.0
-edge-tts
 asyncio
 torchvision
 accelerate
@@ -14,8 +12,4 @@ onnxruntime
 sentencepiece
 soxr
 pydub
-groq
-opencv-python
-qwen-vl-utils
-av
-gradio --pre
+edge-tts

spaces/__init__.py

@@ -1,30 +0,0 @@
-"""
-"""
-
-import sys
-
-
-if sys.version_info.minor < 8: # pragma: no cover
-    raise RuntimeError("Importing PySpaces requires Python 3.8+")
-
-
-# Prevent gradio from importing spaces
-if (gr := sys.modules.get('gradio')) is not None: # pragma: no cover
-    try:
-        gr.Blocks
-    except AttributeError:
-        raise ImportError
-
-
-from .zero.decorator import GPU
-from .gradio import gradio_auto_wrap
-from .gradio import disable_gradio_auto_wrap
-from .gradio import enable_gradio_auto_wrap
-
-
-__all__ = [
-    'GPU',
-    'gradio_auto_wrap',
-    'disable_gradio_auto_wrap',
-    'enable_gradio_auto_wrap',
-]

spaces/config.py

@@ -1,37 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-import os
-from pathlib import Path
-
-from .utils import boolean
-
-
-ZEROGPU_OFFLOAD_DIR_DEFAULT = str(Path.home() / '.zerogpu' / 'tensors')
-
-
-class Settings:
-    def __init__(self):
-        self.zero_gpu = boolean(
-            os.getenv('SPACES_ZERO_GPU'))
-        self.zero_device_api_url = (
-            os.getenv('SPACES_ZERO_DEVICE_API_URL'))
-        self.gradio_auto_wrap = boolean(
-            os.getenv('SPACES_GRADIO_AUTO_WRAP'))
-        self.zero_patch_torch_device = boolean(
-            os.getenv('ZERO_GPU_PATCH_TORCH_DEVICE'))
-        self.zero_gpu_v2 = boolean(
-            os.getenv('ZEROGPU_V2'))
-        self.zerogpu_offload_dir = (
-            os.getenv('ZEROGPU_OFFLOAD_DIR', ZEROGPU_OFFLOAD_DIR_DEFAULT))
-
-
-Config = Settings()
-
-
-if Config.zero_gpu:
-    assert Config.zero_device_api_url is not None, (
-        'SPACES_ZERO_DEVICE_API_URL env must be set '
-        'on ZeroGPU Spaces (identified by SPACES_ZERO_GPU=true)'
-    )

spaces/gradio.py

@@ -1,55 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-from typing import Callable
-from typing import Generator
-from typing import TypeVar
-from typing import overload
-from typing_extensions import ParamSpec
-
-from .config import Config
-from .zero.decorator import GPU
-
-
-Param = ParamSpec('Param')
-Res = TypeVar('Res')
-
-
-gradio_auto_wrap_enabled = Config.gradio_auto_wrap
-
-
-def disable_gradio_auto_wrap():
-    global gradio_auto_wrap_enabled
-    gradio_auto_wrap_enabled = False
-
-def enable_gradio_auto_wrap():
-    global gradio_auto_wrap_enabled
-    gradio_auto_wrap_enabled = True
-
-
-@overload
-def gradio_auto_wrap(
-    task:
-     Callable[Param, Res],
-) -> Callable[Param, Res]:
-    ...
-@overload
-def gradio_auto_wrap(
-    task:
-     None,
-) -> None:
-    ...
-def gradio_auto_wrap(
-    task:
-      Callable[Param, Res]
-    | None,
-) -> (Callable[Param, Res]
-    | None):
-    """
-    """
-    if not gradio_auto_wrap_enabled:
-        return task
-    if not callable(task):
-        return task
-    return GPU(task) # type: ignore

spaces/utils.py

@@ -1,85 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-import ctypes
-import sys
-from functools import lru_cache as cache
-from functools import partial
-
-import multiprocessing
-from multiprocessing.queues import SimpleQueue as _SimpleQueue
-from pathlib import Path
-from pickle import PicklingError
-from typing import Callable
-from typing import TypeVar
-
-
-GRADIO_VERSION_ERROR_MESSAGE = "Make sure Gradio version is at least 3.46"
-
-
-T = TypeVar('T')
-
-
-@cache
-def self_cgroup_device_path() -> str:
-    cgroup_content = Path('/proc/self/cgroup').read_text()
-    for line in cgroup_content.strip().split('\n'):
-        contents = line.split(':devices:')
-        if len(contents) != 2:
-            continue # pragma: no cover
-        return contents[1]
-    raise Exception # pragma: no cover
-
-
-if sys.version_info.minor < 9: # pragma: no cover
-    _SimpleQueue.__class_getitem__ = classmethod(lambda cls, _: cls) # type: ignore
-
-class SimpleQueue(_SimpleQueue[T]):
-    def __init__(self, *args):
-        super().__init__(*args, ctx=multiprocessing.get_context('fork'))
-    def put(self, obj: T):
-        try:
-            super().put(obj)
-        except PicklingError:
-            raise # pragma: no cover
-        # https://bugs.python.org/issue29187
-        except Exception as e:
-            message = str(e)
-            if not "pickle" in message:
-                raise # pragma: no cover
-            raise PicklingError(message)
-    def close(self): # Python 3.8 static typing trick
-        super().close() # type: ignore
-    def wlock_release(self):
-        if (lock := getattr(self, '_wlock', None)) is None:
-            return # pragma: no cover
-        try:
-            lock.release()
-        except ValueError:
-            pass
-
-
-def drop_params(fn: Callable[[], T]) -> Callable[..., T]:
-    def drop(*args):
-        return fn()
-    return drop
-
-
-def boolean(value: str | None) -> bool:
-    return value is not None and value.lower() in ("1", "t", "true")
-
-
-def gradio_request_var():
-    try:
-        from gradio.context import LocalContext
-    except ImportError: # pragma: no cover
-        raise RuntimeError(GRADIO_VERSION_ERROR_MESSAGE)
-    return LocalContext.request
-
-
-def malloc_trim():
-    ctypes.CDLL("libc.so.6").malloc_trim(0)
-
-
-debug = partial(print, 'SPACES_ZERO_GPU_DEBUG')

spaces/zero/__init__.py

@@ -1,21 +0,0 @@
-"""
-"""
-
-from pathlib import Path
-
-from ..config import Config
-
-
-if Config.zero_gpu:
-
-    from . import gradio
-    from . import torch
-
-    if torch.is_in_bad_fork():
-        raise RuntimeError(
-            "CUDA has been initialized before importing the `spaces` package"
-        )
-
-    torch.patch()
-    gradio.one_launch(torch.pack)
-    Path(Config.zerogpu_offload_dir).mkdir(parents=True, exist_ok=True)

spaces/zero/api.py

@@ -1,156 +0,0 @@
-"""
-Synced with huggingface/pyspaces:spaces/zero/api.py
-"""
-from __future__ import annotations
-
-from datetime import timedelta
-from typing import Any
-from typing import Generator
-from typing import Literal
-from typing import NamedTuple
-from typing import Optional
-from typing import overload
-
-import httpx
-from pydantic import BaseModel
-from typing_extensions import assert_never
-
-
-AllowToken = str
-NvidiaIndex = int # TODO: Migrate to GpuIndex (less confusing for MIG)
-NvidiaUUID = str
-CGroupPath = str
-VisitorId = str
-Score = float
-
-AuthLevel = Literal['regular', 'pro']
-
-
-AUTHENTICATED_HEADER = 'X-Authenticated'
-
-
-class ScheduleResponse(BaseModel):
-    idle: bool
-    nvidiaIndex: int
-    nvidiaUUID: str
-    allowToken: str
-
-
-class QuotaInfos(BaseModel):
-    left: int
-    wait: timedelta
-
-
-class ReportUsageMonitoringParams(NamedTuple):
-    nvidia_index: int
-    visitor_id: str
-    duration: timedelta
-
-
-class QueueEvent(BaseModel):
-    event: Literal['ping', 'failed', 'succeeded']
-    data: Optional[ScheduleResponse] = None
-
-
-def sse_parse(text: str):
-    event, *data = text.strip().splitlines()
-    assert event.startswith('event:')
-    event = event[6:].strip()
-    if event in ('ping', 'failed'):
-        return QueueEvent(event=event)
-    assert event == 'succeeded'
-    (data,) = data
-    assert data.startswith('data:')
-    data = data[5:].strip()
-    return QueueEvent(event=event, data=ScheduleResponse.parse_raw(data))
-
-
-def sse_stream(res: httpx.Response) -> Generator[QueueEvent, Any, None]:
-    for text in res.iter_text():
-        if len(text) == 0:
-            break # pragma: no cover
-        try:
-            yield sse_parse(text)
-        except GeneratorExit:
-            res.close()
-            break
-
-
-class APIClient:
-
-    def __init__(self, client: httpx.Client):
-        self.client = client
-
-    def startup_report(self) -> httpx.codes:
-        res = self.client.post('/startup-report')
-        return httpx.codes(res.status_code)
-
-    def schedule(
-        self,
-        cgroup_path: str,
-        task_id: int = 0,
-        token: str | None = None,
-        duration_seconds: int | None = None,
-        enable_queue: bool = True,
-    ):
-        params: dict[str, str | int | bool] = {
-            'cgroupPath': cgroup_path,
-            'taskId': task_id,
-            'enableQueue': enable_queue,
-        }
-        if duration_seconds is not None:
-            params['durationSeconds'] = duration_seconds
-        if token is not None:
-            params['token'] = token
-        res = self.client.send(
-            request=self.client.build_request(
-                method='POST',
-                url='/schedule',
-                params=params,
-            ),
-            stream=True,
-        )
-        status = httpx.codes(res.status_code)
-        auth: AuthLevel | None = res.headers.get(AUTHENTICATED_HEADER)
-        if (status is not httpx.codes.OK and
-            status is not httpx.codes.TOO_MANY_REQUESTS
-        ):
-            res.close()
-            return status, auth
-        if "text/event-stream" in res.headers['content-type']:
-            return sse_stream(res), auth
-        res.read()
-        if status is httpx.codes.TOO_MANY_REQUESTS:
-            return QuotaInfos(**res.json()), auth # pragma: no cover
-        if status is httpx.codes.OK:
-            return ScheduleResponse(**res.json()), auth
-        assert_never(status)
-
-    def allow(
-        self,
-        allow_token: str,
-        pid: int,
-    ):
-        res = self.client.post('/allow', params={
-            'allowToken': allow_token,
-            'pid': pid,
-        })
-        return httpx.codes(res.status_code)
-
-    def release(
-        self,
-        allow_token: str,
-        fail: bool = False,
-    ) -> httpx.codes:
-        res = self.client.post('/release', params={
-            'allowToken': allow_token,
-            'fail': fail,
-        })
-        return httpx.codes(res.status_code)
-
-    def get_queue_size(self) -> int:
-        res = self.client.get('/queue-size')
-        assert res.status_code == 200, res.status_code
-        size = res.json()
-        assert isinstance(size, int)
-        return size

spaces/zero/client.py

@@ -1,239 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-import os
-import time
-import warnings
-from datetime import timedelta
-
-import gradio as gr
-import httpx
-from packaging import version
-from typing_extensions import assert_never
-
-from .. import utils
-from ..config import Config
-from .api import APIClient
-from .api import AuthLevel
-from .api import QuotaInfos
-from .api import ScheduleResponse
-from .gradio import HTMLError
-from .gradio import get_event
-from .gradio import supports_auth
-
-
-TOKEN_HEADER = 'X-IP-Token'
-DEFAULT_SCHEDULE_DURATION = 60
-
-QUOTA_MESSAGE = "You have exceeded your GPU quota"
-UNUSED_MESSAGE = "GPU device not used"
-NO_GPU_MESSAGE_REGULAR = "No GPU was available"
-NO_GPU_MESSAGE_INQUEUE = "No GPU was available after 60s"
-
-SIGNUP_ON_HF_TXT = "Create a free account"
-SIGNUP_ON_HF_URL = "https://huggingface.co/join"
-SUBSCRIBE_TO_PRO_TXT = "Subscribe to Pro"
-SUBSCRIBE_TO_PRO_URL = "https://huggingface.co/settings/billing/subscription"
-
-
-def api_client():
-    assert Config.zero_device_api_url is not None
-    httpx_client = httpx.Client(base_url=Config.zero_device_api_url, timeout=60, verify=False)
-    return APIClient(httpx_client)
-
-
-def startup_report():
-    retries, max_retries = 0, 2
-    client = api_client()
-    while (status := client.startup_report()) is httpx.codes.NOT_FOUND: # pragma: no cover
-        time.sleep(1)
-        if (retries := retries + 1) > max_retries:
-            raise RuntimeError("Error while initializing ZeroGPU: NotFound")
-    if status is not httpx.codes.OK: # pragma: no cover
-        raise RuntimeError("Error while initializing ZeroGPU: Unknown")
-
-
-def html_string(html_contents: str, text_contents: str): # pragma: no cover
-    class HTMLString(str):
-        def __str__(self):
-            return text_contents
-    return HTMLString(html_contents)
-
-
-def _toast_action(
-    auth: AuthLevel | None,
-    supports_html: bool,
-    pro_message: str,
-    unlogged_desc: str,
-    logged_desc: str,
-    ending: str,
-) -> tuple[str, str]: # pragma: no cover
-    if not supports_auth() or auth == 'pro':
-        return pro_message, pro_message
-    html = ""
-    link = SIGNUP_ON_HF_URL if auth is None else SUBSCRIBE_TO_PRO_URL
-    text = SIGNUP_ON_HF_TXT if auth is None else SUBSCRIBE_TO_PRO_TXT
-    desc = unlogged_desc if auth is None else logged_desc
-    desc += f" {ending}."
-    style = ";".join([
-        "white-space: nowrap",
-        "text-underline-offset: 2px",
-        "color: var(--body-text-color)",
-    ])
-    if supports_html:
-        html += f'<a style="{style}" href="{link}">'
-    html += text
-    if supports_html:
-        html += '</a> '
-    html += desc
-    markdown = f'[{text}]({link}) {desc}'
-    return html, markdown
-
-
-def schedule(
-    task_id: int,
-    request: gr.Request | None = None,
-    duration: timedelta | None = None,
-    _first_attempt: bool = True,
-) -> ScheduleResponse:
-
-    if not (gradio_version := version.parse(gr.__version__)).major >= 4: # pragma: no cover
-        raise RuntimeError("ZeroGPU is only compatible with Gradio 4+")
-
-    GRADIO_HTML_TOASTS = gradio_version.minor >= 39
-
-    res, auth = api_client().schedule(
-        cgroup_path=utils.self_cgroup_device_path(),
-        task_id=task_id,
-        token=_get_token(request),
-        duration_seconds=duration.seconds if duration is not None else None,
-    )
-
-    if isinstance(res, ScheduleResponse):
-        return res
-
-    if isinstance(res, QuotaInfos): # pragma: no cover
-        requested = duration.seconds if duration is not None else DEFAULT_SCHEDULE_DURATION
-        if res.wait < timedelta(0):
-            raise gr.Error(
-                f"The requested GPU duration ({requested}s) "
-                f"is larger than the maximum allowed"
-            )
-        else:
-            gpu = "Pro GPU" if auth == 'pro' else ("free GPU" if auth == 'regular' else "GPU")
-            message = (
-                f"You have exceeded your {gpu} quota "
-                f"({requested}s requested vs. {res.left}s left)."
-            )
-            details_html, details_markdown = _toast_action(
-                auth=auth,
-                supports_html=GRADIO_HTML_TOASTS,
-                pro_message=f"Try again in {res.wait}",
-                unlogged_desc="to get more",
-                logged_desc="to get 5x more",
-                ending="usage quota",
-            )
-            message_html = f"{message} {details_html}"
-            message_text = f"{message} {details_markdown}"
-            raise HTMLError(html_string(message_html, message_text))
-
-    if not isinstance(res, httpx.codes): # pragma: no cover
-        gr.Info("Waiting for a GPU to become available")
-        # TODO: Sign-up message if not authenticated (after some time ?)
-        connection_event = get_event()
-        if connection_event is None and request is not None:
-            warnings.warn("ZeroGPU: Cannot get Gradio app Queue instance")
-        while True:
-            try:
-                event = next(res)
-            except StopIteration:
-                raise RuntimeError("Unexpected end of stream")
-            except httpx.RemoteProtocolError:
-                if not _first_attempt:
-                    raise RuntimeError("Error while re-trying after queue disconnect")
-                return schedule(task_id, request, duration, _first_attempt=False)
-            if event.event == 'ping':
-                if connection_event is not None and not connection_event.alive:
-                    res.close()
-                    raise RuntimeError("Connection closed by visitor while queueing")
-                continue
-            if event.event == 'failed':
-                details_html, details_markdown = _toast_action(
-                    auth=auth,
-                    supports_html=GRADIO_HTML_TOASTS,
-                    pro_message="Retry later",
-                    unlogged_desc="to get a higher",
-                    logged_desc="to get the highest",
-                    ending="priority in ZeroGPU queues",
-                )
-                message_html = f"{NO_GPU_MESSAGE_INQUEUE}. {details_html}"
-                message_text = f"{NO_GPU_MESSAGE_INQUEUE} {details_markdown}"
-                raise HTMLError(html_string(message_html, message_text))
-            if event.event == 'succeeded':
-                assert event.data is not None
-                if connection_event is not None and not connection_event.alive:
-                    release(event.data.allowToken)
-                    raise RuntimeError("Connection closed by visitor on queue success")
-                gr.Info("Successfully acquired a GPU")
-                return event.data
-
-    if res is httpx.codes.SERVICE_UNAVAILABLE:
-        raise gr.Error(NO_GPU_MESSAGE_REGULAR)
-
-    # TODO: Find a way to log 'detail' response field
-    raise RuntimeError(f"ZeroGPU API /schedule error: {res} ({httpx.codes.get_reason_phrase(res)})") # pragma: no cover
-
-
-def allow(allow_token: str) -> None:
-    pid = os.getpid()
-    assert pid != 1, "Allowing PID 1 on ZeroGPU will end up killing your Space"
-    assert api_client().allow(allow_token=allow_token, pid=pid) is httpx.codes.OK
-
-
-def release(
-    allow_token: str, *,
-    fail: bool = False,
-    allow_404: bool = False,
-) -> None:
-
-    res = api_client().release(
-        allow_token=allow_token,
-        fail=fail,
-    )
-
-    if res is httpx.codes.NO_CONTENT: # pragma: no cover
-        try:
-            gr.Warning(UNUSED_MESSAGE)
-        except AttributeError:
-            pass
-        warnings.warn(UNUSED_MESSAGE, RuntimeWarning)
-        return None
-
-    if res is httpx.codes.NOT_FOUND:
-        if not allow_404:
-            warnings.warn("ZeroGPU API /release warning: 404 Not Found")
-        return None
-
-    if httpx.codes.is_success(res):
-        return None
-
-    # TODO: Find a way to log 'detail' response field
-    # TODO: Only raise in dev environment. Simply warn in production ?
-    raise RuntimeError(f"ZeroGPU API /release error: {res} ({httpx.codes.get_reason_phrase(res)})") # pragma: no cover
-
-
-def _get_token(request: gr.Request | None) -> str | None:
-
-    if request is None:
-        return None
-
-    headers = getattr(request, 'headers', None)
-    if headers is None or not hasattr(headers, '__dict__'):
-        raise gr.Error("Internal Gradio error")
-
-    # Compatibility trick
-    if not hasattr(headers, 'get'):
-        headers = headers.__dict__ # pragma: no cover
-
-    return headers.get(TOKEN_HEADER.lower())

spaces/zero/decorator.py

@@ -1,113 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-import inspect
-import sys
-import warnings
-from datetime import timedelta
-from functools import partial
-from typing import Callable
-from typing import TypeVar
-from typing import overload
-from typing_extensions import ParamSpec
-from typing_extensions import Unpack
-
-from ..config import Config
-from .types import DynamicDuration
-from .types import EmptyKwargs
-
-
-P = ParamSpec('P')
-R = TypeVar('R')
-
-
-decorated_cache: dict[Callable, Callable] = {}
-
-
-@overload
-def GPU(
-    task: None = None, *,
-    duration: DynamicDuration[P] = None,
-) -> Callable[[Callable[P, R]], Callable[P, R]]:
-    ...
-@overload
-def GPU(
-    task: Callable[P, R], *,
-    duration: DynamicDuration[P] = None,
-) -> Callable[P, R]:
-    ...
-def GPU(
-    task: Callable[P, R] | None = None, *,
-    duration: DynamicDuration[P] = None,
-    **kwargs: Unpack[EmptyKwargs],
-) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
-    """
-    ZeroGPU decorator
-
-    Basic usage:
-        ```
-        @spaces.GPU
-        def fn(...):
-            # CUDA is available here
-            pass
-        ```
-
-    With custom duration:
-        ```
-        @spaces.GPU(duration=45) # Expressed in seconds
-        def fn(...):
-            # CUDA is available here
-            pass
-        ```
-
-    Args:
-        task (`Callable | None`): Python function that requires CUDA
-        duration (`int | datetime.timedelta`): Estimated duration in seconds or `datetime.timedelta`
-
-    Returns:
-        `Callable`: GPU-ready function
-    """
-    if "enable_queue" in kwargs:
-        warnings.warn("`enable_queue` parameter is now ignored and always set to `True`")
-    if task is None:
-        return partial(_GPU, duration=duration)
-    return _GPU(task, duration)
-
-
-def _GPU(
-    task: Callable[P, R],
-    duration: DynamicDuration[P],
-) -> Callable[P, R]:
-
-    if not Config.zero_gpu:
-        return task
-
-    from . import client
-    from .wrappers import regular_function_wrapper
-    from .wrappers import generator_function_wrapper
-
-    if sys.version_info.minor < 9: # pragma: no cover
-        raise RuntimeError("Actually using @spaces.GPU on a ZeroGPU Space requires Python 3.9+")
-
-    if task in decorated_cache:
-        # TODO: Assert same duration ?
-        return decorated_cache[task] # type: ignore
-
-    if inspect.iscoroutinefunction(task):
-        raise NotImplementedError
-
-    if inspect.isgeneratorfunction(task):
-        decorated = generator_function_wrapper(task, duration)
-    else:
-        decorated = regular_function_wrapper(task, duration)
-
-    setattr(decorated, 'zerogpu', None)
-
-    client.startup_report()
-    decorated_cache.update({
-        task:      decorated,
-        decorated: decorated,
-    })
-
-    return decorated # type: ignore

spaces/zero/gradio.py

@@ -1,150 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-from functools import wraps
-from packaging import version
-from typing import Callable
-from typing import NamedTuple
-from typing import TYPE_CHECKING
-import warnings
-
-import gradio as gr
-from gradio.context import Context
-from gradio.context import LocalContext
-from gradio.helpers import Progress
-from gradio.helpers import TrackedIterable
-from gradio.queueing import Queue
-from typing_extensions import ParamSpec
-
-from ..utils import SimpleQueue
-from .types import GeneratorResQueueResult
-from .types import GradioQueueEvent
-from .types import RegularResQueueResult
-
-
-QUEUE_RPC_METHODS = [
-    "set_progress",
-    "log_message",
-]
-
-
-class GradioPartialContext(NamedTuple):
-    event_id: str | None
-    in_event_listener: bool
-    progress: Progress | None
-
-    @staticmethod
-    def get():
-        TrackedIterable.__reduce__ = tracked_iterable__reduce__
-        return GradioPartialContext(
-            event_id=LocalContext.event_id.get(),
-            in_event_listener=LocalContext.in_event_listener.get(),
-            progress=LocalContext.progress.get(),
-        )
-
-    @staticmethod
-    def apply(context: 'GradioPartialContext'):
-        LocalContext.event_id.set(context.event_id)
-        LocalContext.in_event_listener.set(context.in_event_listener)
-        LocalContext.progress.set(context.progress)
-
-
-def get_queue_instance():
-    blocks = LocalContext.blocks.get()
-    if blocks is None: # pragma: no cover
-        return None
-    return blocks._queue
-
-
-def get_event():
-    queue = get_queue_instance()
-    event_id = LocalContext.event_id.get()
-    if queue is None:
-        return None
-    if event_id is None: # pragma: no cover
-        return None
-    for job in queue.active_jobs:
-        if job is None: # pragma: no cover
-            continue
-        for event in job:
-            if event._id == event_id:
-                return event
-
-
-def get_server_port() -> int | None:
-    from_request_context = True
-    if (blocks := LocalContext.blocks.get()) is None: # Request
-        from_request_context = False
-        if (blocks := Context.root_block) is None: # Caching
-            return None
-    if (server := getattr(blocks, 'server', None)) is None:
-        if from_request_context:
-            warnings.warn("Gradio: No blocks.server inside a request") # pragma: no cover
-        return -1
-    if TYPE_CHECKING:
-        assert (server := blocks.server)
-    return server.config.port
-
-
-def try_process_queue_event(method_name: str, *args, **kwargs):
-    queue = get_queue_instance()
-    if queue is None: # pragma: no cover
-        warnings.warn("ZeroGPU: Cannot get Gradio app Queue instance")
-        return
-    method = getattr(queue, method_name, None)
-    assert callable(method)
-    method(*args, **kwargs)
-
-
-def patch_gradio_queue(
-    res_queue: SimpleQueue[RegularResQueueResult | None] | SimpleQueue[GeneratorResQueueResult | None],
-):
-
-    def rpc_method(method_name: str):
-        def method(*args, **kwargs):
-            if args and isinstance(args[0], Queue):
-                args = args[1:] # drop `self`
-            res_queue.put(GradioQueueEvent(method_name, args, kwargs))
-        return method
-
-    for method_name in QUEUE_RPC_METHODS:
-        if (method := getattr(Queue, method_name, None)) is None: # pragma: no cover
-            warnings.warn(f"ZeroGPU: Gradio Queue has no {method_name} attribute")
-            continue
-        if not callable(method): # pragma: no cover
-            warnings.warn(f"ZeroGPU: Gradio Queue {method_name} is not callable")
-            continue
-        setattr(Queue, method_name, rpc_method(method_name))
-
-    TrackedIterable.__reduce__ = tracked_iterable__reduce__
-
-
-def tracked_iterable__reduce__(self):
-    res: tuple = super(TrackedIterable, self).__reduce__() # type: ignore
-    cls, base, state, *_ = res
-    return cls, base,{**state, **{
-        'iterable': None,
-        '_tqdm': None,
-    }}
-
-
-def supports_auth():
-    return version.parse(gr.__version__) >= version.Version('4.27.0')
-
-
-Param = ParamSpec('Param')
-
-def one_launch(task: Callable[Param, None], *task_args: Param.args, **task_kwargs: Param.kwargs):
-    _launch = gr.Blocks.launch
-    @wraps(gr.Blocks.launch)
-    def launch(*args, **kwargs):
-        task(*task_args, **task_kwargs)
-        gr.Blocks.launch = _launch
-        return gr.Blocks.launch(*args, **kwargs)
-    gr.Blocks.launch = launch
-
-
-class HTMLError(gr.Error):
-    def __str__(self): # pragma: no cover
-        return self.message

spaces/zero/torch/__init__.py

@@ -1,42 +0,0 @@
-"""
-"""
-
-from ...config import Config
-
-
-try:
-
-    import torch
-
-except ImportError:
-
-    _patch = lambda *args, **kwargs: None
-    _unpatch = lambda *args, **kwargs: None
-    _pack = lambda *args, **kwargs: None
-    _init = lambda *args, **kwargs: None
-    _size = lambda *args, **kwargs: 0
-    _move = lambda *args, **kwargs: None
-    _is_in_bad_fork = lambda *args, **kwargs: False
-
-else:
-
-    if Config.zero_gpu_v2:
-        from . import patching as _patching
-    else: # pragma: no cover
-        from . import patching_legacy as _patching
-
-    _patch = _patching.patch
-    _unpatch = _patching.unpatch
-    _pack = _patching.pack
-    _init = _patching.init
-    _size = _patching.size
-    _move = _patching.move
-    _is_in_bad_fork = _patching.is_in_bad_fork
-
-patch = _patch
-unpatch = _unpatch
-pack = _pack
-init = _init
-size = _size
-move = _move
-is_in_bad_fork = _is_in_bad_fork

spaces/zero/torch/bitsandbytes.py

@@ -1,162 +0,0 @@
-"""
-"""
-# pyright: reportPrivateImportUsage=false
-
-from __future__ import annotations
-
-import importlib
-from contextlib import contextmanager
-from importlib import metadata
-from types import ModuleType
-from typing import TYPE_CHECKING
-from typing import Tuple
-
-import torch
-from packaging import version
-
-if TYPE_CHECKING:
-    import torch as Torch
-
-
-@contextmanager
-def cuda_unavailable(torch: ModuleType):
-    _is_available = torch.cuda.is_available
-    torch.cuda.is_available = lambda: False
-    yield
-    torch.cuda.is_available = _is_available
-
-
-def maybe_import_bitsandbytes():
-    try:
-        import torch
-    except ImportError: # pragma: no cover
-        return None
-    with cuda_unavailable(torch):
-        try:
-            import bitsandbytes
-        except ImportError:
-            bitsandbytes = None
-        else:
-            if (bnb_version := version.parse(metadata.version('bitsandbytes'))) < version.parse('0.40.0'):
-                raise RuntimeError(f"ZeroGPU requires bitsandbytes >= 0.40.0 (installed: {bnb_version})") # pragma: no cover
-            print("↑ Those bitsandbytes warnings are expected on ZeroGPU ↑")
-    return bitsandbytes
-
-
-if (bnb := maybe_import_bitsandbytes()):
-
-    from torch.utils.weak import WeakTensorKeyDictionary
-
-    with cuda_unavailable(torch):
-        from bitsandbytes import cextension
-        from bitsandbytes import functional
-        try: # bitsandbytes < 0.44
-            from bitsandbytes.cuda_setup.main import CUDASetup
-        except ModuleNotFoundError: # pragma: no cover
-            CUDASetup = None
-        from bitsandbytes.nn import Int8Params
-        from bitsandbytes.nn import Params4bit
-
-    _param_to_8bit   = Int8Params.to     # type: ignore
-    _param_cuda_8bit = Int8Params.cuda
-    _param_to_4bit   = Params4bit.to     # type: ignore
-    _param_cuda_4bit = Params4bit.cuda
-
-    TensorToArgs = Tuple[torch.device, torch.dtype, bool, torch.memory_format]
-
-    to_ops_8bit: dict[Int8Params, TensorToArgs | None] = WeakTensorKeyDictionary() # type: ignore
-    to_ops_4bit: dict[Params4bit, TensorToArgs | None] = WeakTensorKeyDictionary() # type: ignore
-
-    def _to_op_register_8bit(self: Int8Params, *args, **kwargs):
-        parsed = torch._C._nn._parse_to(*args, **kwargs)
-        device, *_ = parsed
-        if not isinstance(device, torch.device): # pragma: no cover
-            return _param_to_8bit(self, *args, **kwargs)
-        if device.type != 'cuda':
-            return _param_to_8bit(self, *args, **kwargs)
-        to_ops_8bit[self] = parsed
-        return self
-
-    def _to_op_register_4bit(self: Params4bit, *args, **kwargs):
-        parsed = torch._C._nn._parse_to(*args, **kwargs)
-        device, *_ = parsed
-        if not isinstance(device, torch.device): # pragma: no cover
-            return _param_to_4bit(self, *args, **kwargs)
-        if device.type != 'cuda':
-            return _param_to_4bit(self, *args, **kwargs)
-        to_ops_4bit[self] = parsed
-        return self
-
-    def _cuda_op_arg_check(device: Torch.device | int | str | None) -> bool:
-        if device is None: # pragma: no cover
-            return True
-        if isinstance(device, int):
-            return True
-        if isinstance(device, str): # pragma: no cover
-            device = torch.device(device)
-        return device.type == 'cuda' # pragma: no cover
-
-    def _cuda_op_register_8bit(self: Int8Params, device: Torch.device | int | str | None = None, **kwargs):
-        if not _cuda_op_arg_check(device): # pragma: no cover
-            # Let PyTorch handle the fail
-            return _param_cuda_8bit(self, device, **kwargs)
-        to_ops_8bit[self] = None
-        return self
-
-    def _cuda_op_register_4bit(self: Params4bit, device: Torch.device | int | str | None = None, **kwargs):
-        if not _cuda_op_arg_check(device): # pragma: no cover
-            # Let PyTorch handle the fail
-            return _param_cuda_4bit(self, device, **kwargs)
-        to_ops_4bit[self] = None
-        return self
-
-    def _patch():
-        Int8Params.to   = _to_op_register_8bit   # type: ignore
-        Int8Params.cuda = _cuda_op_register_8bit # type: ignore
-        Params4bit.to   = _to_op_register_4bit   # type: ignore
-        Params4bit.cuda = _cuda_op_register_4bit # type: ignore
-
-    def _unpatch():
-        Int8Params.to   = _param_to_8bit   # type: ignore
-        Int8Params.cuda = _param_cuda_8bit
-        Params4bit.to   = _param_to_4bit   # type: ignore
-        Params4bit.cuda = _param_cuda_4bit
-
-    def _move():
-        if CUDASetup is not None:
-            CUDASetup._instance = None
-        importlib.reload(cextension)
-        functional.lib = cextension.lib
-        for op in to_ops_8bit.items():
-            tensor, parsed_args = op
-            if parsed_args:
-                _, dtype, _, memory_format = parsed_args
-            else:
-                dtype, memory_format = None, None
-            tensor.data = _param_to_8bit(tensor,
-                device='cuda',
-                dtype=dtype,
-                memory_format=memory_format,
-            ) # type: ignore
-        for op in to_ops_4bit.items():
-            tensor, parsed_args = op
-            if parsed_args:
-                _, dtype, _, memory_format = parsed_args
-            else:
-                dtype, memory_format = None, None
-            tensor.data = _param_to_4bit(tensor,
-                device='cuda',
-                dtype=dtype,
-                memory_format=memory_format,
-            ) # type: ignore
-
-else:
-
-    _patch = lambda: None
-    _unpatch = lambda: None
-    _move = lambda: None
-
-
-patch = _patch
-unpatch = _unpatch
-move = _move

spaces/zero/torch/packing.py

@@ -1,209 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-import time
-
-import ctypes
-import os
-from concurrent.futures import as_completed
-from concurrent.futures import ThreadPoolExecutor
-from contextvars import copy_context
-from dataclasses import dataclass
-from queue import Queue
-from typing import Callable
-
-from ...utils import debug
-
-import torch
-from typing_extensions import TypeAlias
-
-
-PAGE_SIZE = 4096
-TOTAL_MEMORY = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
-VM_MAX_SIZE = min(2**38, TOTAL_MEMORY // 2)
-
-BUFFER_SIZE = 64 * 2**20
-BUFFER_COUNT = 2
-
-
-TensorWithSizes: TypeAlias = 'tuple[torch.Tensor, int, int]'
-
-@dataclass
-class ZeroGPUTensorPack:
-    base_dir: str
-    batches: list[list[TensorWithSizes]]
-    big_tensors: list[TensorWithSizes]
-    fakes: dict[torch.Tensor, list[torch.Tensor]]
-    total_size: int
-    def path(self):
-        return f'{self.base_dir}/{id(self)}'
-    def __del__(self):
-        try:
-            os.remove(self.path())
-        except FileNotFoundError: # pragma: no cover
-            pass
-
-
-def write(fd: int, tensor: torch.Tensor):
-    clone = torch.empty_like(tensor)
-    size = clone.untyped_storage().size() # pyright: ignore [reportAttributeAccessIssue]
-    buffer = torch.UntypedStorage(VM_MAX_SIZE)
-    buffer_ptr = buffer.data_ptr()
-    offset = -buffer_ptr % PAGE_SIZE
-    padding = -size % PAGE_SIZE
-    clone.set_(buffer[offset:offset+size], 0, clone.shape, clone.stride()) # pyright: ignore [reportArgumentType]
-    clone.copy_(tensor)
-    mv = memoryview((ctypes.c_char * (size+padding)).from_address(buffer_ptr+offset))
-    written_bytes = 0
-    while written_bytes < size:
-        written_bytes += os.write(fd, mv[written_bytes:])
-
-
-def pack_tensors(
-    tensors: set[torch.Tensor],
-    fakes: dict[torch.Tensor, list[torch.Tensor]],
-    offload_dir: str,
-    callback: Callable[[int]] | None = None,
-):
-
-    callback = (lambda bytes: None) if callback is None else callback
-
-    batches: list[list[TensorWithSizes]] = []
-    big_tensors: list[TensorWithSizes] = []
-
-    tensors_with_sizes: list[tuple[torch.Tensor, int, int]] = []
-    for tensor in tensors:
-        size = tensor.numel() * tensor.element_size()
-        aligned_size = size + (-size % PAGE_SIZE)
-        tensors_with_sizes += [(tensor, size, aligned_size)]
-
-    current_batch, current_size = [], 0
-    for (tensor, size, aligned_size) in sorted(tensors_with_sizes, key=lambda item: item[2]):
-        if aligned_size > BUFFER_SIZE:
-            big_tensors += [(tensor, size, aligned_size)]
-            continue
-        current_size += aligned_size
-        if current_size > BUFFER_SIZE:
-            batches += [current_batch]
-            current_batch, current_size = [(tensor, size, aligned_size)], aligned_size
-        else:
-            current_batch += [(tensor, size, aligned_size)]
-
-    if current_batch:
-        batches += [current_batch]
-
-    get_meta = {tensor: torch.empty_like(tensor) for tensor in tensors}
-    batches_meta = [[(get_meta[tensor], size, asize) for tensor, size, asize in batch] for batch in batches]
-    big_tensors_meta = [(get_meta[tensor], size, asize) for tensor, size, asize in big_tensors]
-    fakes_meta = {get_meta[tensor]: fake_list for tensor, fake_list in fakes.items()}
-
-    pack = ZeroGPUTensorPack(
-        base_dir=offload_dir,
-        batches=batches_meta,
-        big_tensors=big_tensors_meta,
-        fakes=fakes_meta,
-        total_size=sum([size for _, size, _ in tensors_with_sizes]),
-    )
-
-    fd = os.open(pack.path(), os.O_CREAT | os.O_WRONLY | os.O_DIRECT)
-    try:
-        total_asize = sum([aligned_size for batch in batches for *_, aligned_size in batch])
-        total_asize += sum([aligned_size for *_, aligned_size in big_tensors])
-        if total_asize > 0:
-            os.posix_fallocate(fd, 0, total_asize)
-            for batch in batches:
-                for tensor, size, _ in batch:
-                    write(fd, tensor)
-                    callback(size)
-            for tensor, size, _ in big_tensors:
-                write(fd, tensor)
-                callback(size)
-        return pack
-    finally:
-        os.close(fd)
-
-
-def pack_to_cuda(pack: ZeroGPUTensorPack, callback: Callable[[int]] | None = None):
-
-    callback = (lambda bytes: None) if callback is None else callback
-
-    free_buffers: Queue[torch.Tensor] = Queue()
-    read_buffers: Queue[torch.Tensor] = Queue()
-
-    for _ in range(BUFFER_COUNT):
-        free_buffers.put(torch.ByteTensor(BUFFER_SIZE).pin_memory())
-
-    def read(fd: int, buffer: torch.Tensor, size: int):
-        mv = memoryview((ctypes.c_char * size).from_address(buffer.data_ptr()))
-        read_bytes = 0
-        while read_bytes < size:
-            read_bytes += os.readv(fd, [mv[read_bytes:]])
-
-    def disk_to_pin(fd: int):
-        for batch in pack.batches:
-            buffer = free_buffers.get()
-            batch_size = sum([aligned_size for *_, aligned_size in batch])
-            read(fd, buffer, batch_size)
-            read_buffers.put(buffer)
-        for *_, aligned_size in pack.big_tensors:
-            read_bytes = 0
-            while read_bytes < aligned_size:
-                buffer = free_buffers.get()
-                read_size = min(BUFFER_SIZE, aligned_size - read_bytes)
-                read(fd, buffer, read_size)
-                read_buffers.put(buffer)
-                read_bytes += read_size
-
-    def pin_to_cuda():
-        total_duration_in_callback = 0
-        for batch in pack.batches:
-            buffer = read_buffers.get()
-            offset = 0
-            cuda_storages = []
-            for tensor, size, aligned_size in batch:
-                cuda_storages += [buffer[offset:offset+size].cuda(non_blocking=True)]
-                offset += aligned_size
-            torch.cuda.synchronize()
-            free_buffers.put(buffer)
-            batch_total_size = 0
-            for (tensor, size, _), cuda_storage in zip(batch, cuda_storages):
-                cuda_tensor = torch.tensor([], dtype=tensor.dtype, device='cuda')
-                cuda_tensor = cuda_tensor.set_(cuda_storage.untyped_storage(), 0, tensor.shape, tensor.stride())
-                for fake in pack.fakes[tensor]:
-                    fake.data = cuda_tensor
-                batch_total_size += size
-            t0 = time.perf_counter()
-            callback(batch_total_size)
-            total_duration_in_callback += time.perf_counter() - t0
-        for tensor, size, _ in pack.big_tensors:
-            cuda_storage = torch.empty(size, dtype=torch.uint8, device='cuda')
-            offset = 0
-            while offset < size:
-                buffer = read_buffers.get()
-                read_size = min(BUFFER_SIZE, size - offset)
-                cuda_storage[offset:offset+read_size] = buffer[:read_size]
-                offset += read_size
-                torch.cuda.synchronize() # Probably not needed
-                free_buffers.put(buffer)
-                t0 = time.perf_counter()
-                callback(read_size)
-                total_duration_in_callback += time.perf_counter() - t0
-            cuda_tensor = torch.tensor([], dtype=tensor.dtype, device='cuda')
-            cuda_tensor = cuda_tensor.set_(cuda_storage.untyped_storage(), 0, tensor.shape, tensor.stride())
-            for fake in pack.fakes[tensor]:
-                fake.data = cuda_tensor
-
-        debug(f"{total_duration_in_callback=}")
-
-    with ThreadPoolExecutor(2) as e:
-        fd = os.open(pack.path(), os.O_RDONLY | os.O_DIRECT)
-        try:
-            futures = [
-                e.submit(copy_context().run, disk_to_pin, fd),
-                e.submit(copy_context().run, pin_to_cuda),
-            ]
-            for future in as_completed(futures):
-                future.result()
-        finally:
-            os.close(fd)

spaces/zero/torch/patching.py

@@ -1,386 +0,0 @@
-"""
-"""
-# pyright: reportPrivateImportUsage=false
-
-from __future__ import annotations
-
-import gc
-import multiprocessing
-import os
-from collections import defaultdict
-from concurrent.futures import ProcessPoolExecutor
-from concurrent.futures import ThreadPoolExecutor
-from contextlib import nullcontext
-from contextvars import copy_context
-from types import SimpleNamespace
-from typing import Any
-from typing import Callable
-
-import torch
-from torch.overrides import TorchFunctionMode
-from torch.overrides import resolve_name
-from torch.utils._python_dispatch import TorchDispatchMode
-from torch.utils._pytree import tree_map_only
-from torch.utils.weak import WeakTensorKeyDictionary
-
-from ...config import Config
-from ...utils import malloc_trim
-from ..tqdm import tqdm
-from . import bitsandbytes
-from .packing import ZeroGPUTensorPack
-from .packing import pack_tensors
-from .packing import pack_to_cuda
-from .types import AliasId
-
-
-# Nvidia A100.80G MIG (drivers 535) / Torch 2.2.0
-CUDA_DEVICE_NAME = 'NVIDIA A100-SXM4-80GB MIG 3g.40gb'
-CUDA_TOTAL_MEMORY = 42144366592
-CUDA_MEM_GET_INFO = (41911451648, CUDA_TOTAL_MEMORY)
-CUDA_DEVICE_CAPABILITY = (8, 0)
-CUDA_DEVICE_PROPERTIES = SimpleNamespace(name=CUDA_DEVICE_NAME, major=8, minor=0, total_memory=CUDA_TOTAL_MEMORY, multi_processor_count=42)
-
-OPS_INPUTS_CHECK_NO_RETURN = (
-    torch.Tensor.equal,
-)
-
-OPS_INPUT_CHECK_SELF_RETURN = (
-    torch.Tensor.set_, # probably never dispatched
-    torch.ops.aten.set_.source_Tensor, # pyright: ignore [reportAttributeAccessIssue]
-)
-
-OFFLOADED_ERROR_MESSAGE = "Cannot apply function {} on disk-offloaded Tensor {}"
-
-_tensor_make_subclass = torch.Tensor._make_subclass
-_asarray           = torch.asarray
-_cuda_init         = torch._C._cuda_init
-_cuda_exchange_device = torch.cuda._exchange_device
-_cuda_available      = torch.cuda.is_available
-_cuda_device_count   = torch.cuda.device_count
-_cuda_current_device = torch.cuda.current_device
-_cuda_mem_get_info   = torch.cuda.mem_get_info
-_cuda_get_device_capability   = torch.cuda.get_device_capability
-_cuda_get_device_properties   = torch.cuda.get_device_properties
-_cuda_get_device_name         = torch.cuda.get_device_name
-
-# PyTorch 2.3
-_cuda_maybe_exchange_device = getattr(torch.cuda, '_maybe_exchange_device', None)
-
-
-cuda_aliases: dict[torch.Tensor, torch.Tensor | None] = WeakTensorKeyDictionary() # pyright: ignore [reportAssignmentType]
-
-tensor_packs: list[ZeroGPUTensorPack] = []
-
-class ZeroGPUTensor(torch.Tensor):
-    pass
-
-def empty_fake(tensor: torch.Tensor):
-    fake = torch.empty_like(tensor, requires_grad=tensor.requires_grad)
-    if fake.__class__ != tensor.__class__:
-        fake = _tensor_make_subclass(tensor.__class__, fake, require_grad=tensor.requires_grad) # pyright: ignore [reportArgumentType]
-    return fake
-
-class ZeroGPUFunctionMode(TorchFunctionMode):
-
-    def __torch_function__(self, func, types, args=(), kwargs: dict[str, Any] | None = None):
-
-        kwargs = {} if kwargs is None else kwargs
-
-        if func == torch._C._nn._parse_to:
-            return func(*args, **kwargs)
-
-        # Redispatch: tensor.cuda() -> tensor.to(device='cuda')
-        if func == torch.Tensor.cuda or func == torch.Tensor.cpu:
-            memory_format = kwargs.get('memory_format')
-            return self.__torch_function__(torch.Tensor.to, types, (args[0],), {
-                'device': 'cuda' if func == torch.Tensor.cuda else 'cpu',
-                **({'memory_format': memory_format} if memory_format is not None else {}),
-            })
-
-        # Redispatch: tensor.to('cuda') -> tensor.to(device='cuda')
-        if func == torch.Tensor.to and len(args) > 1:
-            device, dtype, _, memory_format = torch._C._nn._parse_to(*args[1:], **kwargs)
-            return self.__torch_function__(torch.Tensor.to, types, (args[0],), {
-                'device': device,
-                'dtype': dtype,
-                'memory_format': memory_format,
-            })
-
-        if func == torch.Tensor.data.__set__: # pyright: ignore [reportAttributeAccessIssue]
-            self, target = args
-            if target in cuda_aliases:
-                if (target_original := cuda_aliases[target]) is None:
-                    raise Exception(OFFLOADED_ERROR_MESSAGE.format(resolve_name(func), target))
-                original = empty_fake(self)
-                original.data = target_original
-                cuda_aliases[self] = original
-            elif self in cuda_aliases:
-                del cuda_aliases[self]
-            self.data = target
-            return
-
-        if func == torch.Tensor.device.__get__:
-            tensor, = args
-            if tensor in cuda_aliases:
-                return torch.device('cuda', index=0)
-
-        elif func == torch.Tensor.__repr__:
-            tensor, = args
-            if tensor in cuda_aliases:
-                if (original := cuda_aliases[tensor]) is None:
-                    original = tensor.to('meta')
-                original_class = original.__class__
-                original.__class__ = ZeroGPUTensor
-                try:
-                    return func(original, **kwargs)
-                finally:
-                    original.__class__ = original_class
-
-        elif func == torch.Tensor.untyped_storage:
-            tensor, = args
-            if tensor in cuda_aliases:
-                if (original := cuda_aliases[tensor]) is None:
-                    raise Exception(OFFLOADED_ERROR_MESSAGE.format(resolve_name(func), tensor))
-                res = func(original, **kwargs)
-                res._zerogpu = True
-                return res
-
-        cuda: bool | None = None
-
-        # Handle device kwarg
-        if (device := kwargs.get('device')) is not None:
-            device = torch.device(device)
-            if device.type == 'cuda':
-                kwargs['device'] = torch.device('cpu')
-                cuda = True
-            else:
-                cuda = False
-
-        # Swap fake inputs with original data
-        swapped = {}
-        inputs_are_cuda = set()
-        def swap(tensor: torch.Tensor):
-            nonlocal inputs_are_cuda
-            if tensor not in cuda_aliases:
-                inputs_are_cuda |= {False}
-                return tensor
-            if (original := cuda_aliases[tensor]) is None:
-                raise Exception(OFFLOADED_ERROR_MESSAGE.format(resolve_name(func), tensor))
-            swapped[original] = tensor
-            inputs_are_cuda |= {True}
-            return original
-        args_ = tree_map_only(torch.Tensor, swap, args)
-        kwargs_ = tree_map_only(torch.Tensor, swap, kwargs)
-        if inputs_are_cuda == {True}:
-            if cuda is not False:
-                cuda = True
-
-        res = func(*args_, **kwargs_)
-
-        # Re-generate swapped fakes in case of mutation
-        for original, fake in swapped.items():
-            fake.data = empty_fake(original)
-
-        # Special case for Tensor indexing where only 'self' matters
-        if func in {
-            torch.ops.aten.index.Tensor, # pyright: ignore [reportAttributeAccessIssue]
-            torch.Tensor.__getitem__, # PyTorch 2.4+
-        }:
-            self = args[0]
-            cuda = self in cuda_aliases
-            inputs_are_cuda = {cuda}
-
-        # Emulate device check
-        if isinstance(res, torch.Tensor) or func in OPS_INPUTS_CHECK_NO_RETURN:
-            self = None
-            if len(args_) >= 1 and isinstance(args_[0], torch.Tensor):
-                self = args_[0]
-            # Only raise if func does not return its first input (Tensor.copy_)
-            if res is not self or func in OPS_INPUT_CHECK_SELF_RETURN:
-                if inputs_are_cuda == {True, False}:
-                    raise RuntimeError(
-                        "Expected all tensors to be on the same device, "
-                        "but found at least two devices, cuda:0 (ZeroGPU) and cpu!"
-                    )
-
-        # Register output
-        def register(tensor: torch.Tensor):
-            if tensor in swapped and cuda is not False:
-                return swapped[tensor]
-            if cuda is not True:
-                return tensor
-            fake = empty_fake(tensor)
-            cuda_aliases[fake] = tensor
-            return fake
-
-        return tree_map_only(torch.Tensor, register, res)
-
-# When enabling DispatchMode, some aten ops are dispatched to FunctionMode
-# We are using it for aten.alias.default and aten.set_.source_Tensor
-class DefaultDispatchMode(TorchDispatchMode):
-    def __torch_dispatch__(self, func, types, args=(), kwargs: dict[str, Any] | None = None):
-        return func(*args, **(kwargs or {}))
-
-
-function_mode = ZeroGPUFunctionMode()
-dispatch_mode = DefaultDispatchMode()
-
-
-def _untyped_storage_new_register(*args, **kwargs):
-    cuda = False
-    if (device := kwargs.get('device')) is not None and device.type == 'cuda':
-        cuda = True
-        del kwargs['device']
-    storage = torch._C.StorageBase.__new__(*args, **kwargs)
-    if cuda:
-        storage._zerogpu = True
-    return storage
-
-@property
-def _untyped_storage_device(self):
-    if hasattr(self, '_zerogpu'):
-        return torch.device('cuda', index=0)
-    return torch._C.StorageBase.device.__get__(self) # pyright: ignore [reportAttributeAccessIssue]
-
-# Force dispatch
-def _tensor_make_subclass_function_mode(*args, **kwargs):
-    with torch._C.DisableTorchFunction():
-        return function_mode.__torch_function__(_tensor_make_subclass, (), args=args, kwargs=kwargs)
-def _asarray_function_mode(*args, **kwargs):
-    with torch._C.DisableTorchFunction():
-        return function_mode.__torch_function__(_asarray, (), args=args, kwargs=kwargs)
-
-def _cuda_init_raise():
-    raise RuntimeError(
-        "CUDA must not be initialized in the main process "
-        "on Spaces with Stateless GPU environment.\n"
-        "You can look at this Stacktrace to find out "
-        "which part of your code triggered a CUDA init"
-    )
-
-def _cuda_dummy_exchange_device(device):
-    assert device in {-1, 0}
-    return device
-
-def patch():
-    function_mode.__enter__()
-    dispatch_mode.__enter__()
-    # TODO: only patch bellow methods on current Thread to be consistent with TorchModes
-    # (or hijack threading.Thread.__init__ to force Modes on all threads)
-    torch.Tensor._make_subclass = _tensor_make_subclass_function_mode # pyright: ignore [reportAttributeAccessIssue]
-    torch.UntypedStorage.__new__ = _untyped_storage_new_register
-    torch.UntypedStorage.device  = _untyped_storage_device # pyright: ignore [reportAttributeAccessIssue]
-    torch.asarray           = _asarray_function_mode
-    torch._C._cuda_init     = _cuda_init_raise
-    torch.cuda._exchange_device = _cuda_dummy_exchange_device
-    torch.cuda.is_available   = lambda: True
-    torch.cuda.device_count   = lambda: 1
-    torch.cuda.current_device = lambda: 0
-    torch.cuda.mem_get_info   = lambda *args, **kwargs: CUDA_MEM_GET_INFO
-    torch.cuda.get_device_capability = lambda *args, **kwargs: CUDA_DEVICE_CAPABILITY
-    torch.cuda.get_device_properties = lambda *args, **kwargs: CUDA_DEVICE_PROPERTIES
-    torch.cuda.get_device_name       = lambda *args, **kwargs: CUDA_DEVICE_NAME
-    # PyTorch 2.3
-    if _cuda_maybe_exchange_device is not None: # pragma: no cover
-        setattr(torch.cuda, '_maybe_exchange_device', _cuda_dummy_exchange_device)
-    bitsandbytes.patch()
-
-def unpatch():
-    try:
-        dispatch_mode.__exit__(None, None, None)
-        function_mode.__exit__(None, None, None)
-    except RuntimeError:
-        pass # patch() and unpatch() called from != threads
-    torch.Tensor._make_subclass = _tensor_make_subclass
-    torch.UntypedStorage.__new__ = torch._C.StorageBase.__new__
-    torch.UntypedStorage.device  = torch._C.StorageBase.device # pyright: ignore [reportAttributeAccessIssue]
-    torch.asarray           = _asarray
-    torch._C._cuda_init     = _cuda_init
-    torch.cuda._exchange_device = _cuda_exchange_device
-    torch.cuda.is_available   = _cuda_available
-    torch.cuda.device_count   = _cuda_device_count
-    torch.cuda.current_device = _cuda_current_device
-    torch.cuda.mem_get_info   = _cuda_mem_get_info
-    torch.cuda.get_device_capability = _cuda_get_device_capability
-    torch.cuda.get_device_properties = _cuda_get_device_properties
-    torch.cuda.get_device_name       = _cuda_get_device_name
-    # PyTorch 2.3
-    if _cuda_maybe_exchange_device is not None: # pragma: no cover
-        setattr(torch.cuda, '_maybe_exchange_device', _cuda_exchange_device)
-    bitsandbytes.unpatch()
-
-
-def _total_unpacked_size():
-    tensors = [tensor for tensor in cuda_aliases.values() if tensor is not None]
-    deduped = {AliasId.from_tensor(tensor): tensor for tensor in tensors}
-    return sum([tensor.numel() * tensor.element_size() for tensor in deduped.values()])
-
-
-def _pack(offload_dir: str):
-    # Pack to disk
-    originals: set[torch.Tensor] = set()
-    originals_dedup: dict[AliasId, torch.Tensor] = {}
-    fakes: dict[torch.Tensor, list[torch.Tensor]] = defaultdict(list)
-    for fake, original in cuda_aliases.items():
-        # TODO filter-out sparse Tensors
-        if original is not None:
-            original_id = AliasId.from_tensor(original)
-            if original_id not in originals_dedup:
-                originals_dedup[original_id] = original
-                originals |= {original}
-            fakes[originals_dedup[original_id]] += [fake]
-    progress = tqdm(
-        total=_total_unpacked_size(),
-        unit='B',
-        unit_scale=True,
-        desc="ZeroGPU tensors packing",
-    ) if tqdm is not None else nullcontext()
-    with progress as progress:
-        update = progress.update if progress is not None else lambda _: None
-        pack = pack_tensors(originals, fakes, offload_dir, callback=update)
-    tensor_packs.append(pack)
-    # Free memory
-    for fake_list in fakes.values():
-        for fake in fake_list:
-            cuda_aliases[fake] = None
-
-def pack():
-    _pack(Config.zerogpu_offload_dir)
-    gc.collect()
-    malloc_trim()
-
-def init(nvidia_uuid: str):
-    os.environ['CUDA_VISIBLE_DEVICES'] = nvidia_uuid
-    torch.Tensor([0]).cuda()
-
-def size():
-    return _total_unpacked_size() + sum([pack.total_size for pack in tensor_packs])
-
-def _move(callback: Callable[[int]] | None = None):
-    callback = callback if callback is not None else lambda _: None
-    # CPU -> CUDA
-    moved: dict[AliasId, torch.Tensor] = {}
-    for fake, original in cuda_aliases.items():
-        if original is not None:
-            original_id = AliasId.from_tensor(original)
-            if original_id not in moved:
-                moved[original_id] = original.cuda()
-                callback(fake.numel() * fake.element_size())
-    for fake, original in cuda_aliases.items():
-        if original is not None:
-            fake.data = moved[AliasId.from_tensor(original)]
-    # Disk -> CUDA
-    for tensor_pack in tensor_packs:
-        pack_to_cuda(tensor_pack, callback=callback)
-    bitsandbytes.move()
-
-def move(callback: Callable[[int]] | None = None):
-    callback = callback if callback is not None else lambda _: None
-    with ThreadPoolExecutor(1) as e:
-        e.submit(copy_context().run, _move, callback=callback).result()
-    torch.cuda.synchronize()
-
-def is_in_bad_fork():
-    with ProcessPoolExecutor(mp_context=multiprocessing.get_context('fork')) as e:
-        f = e.submit(torch.cuda._is_in_bad_fork)
-        return f.result()

spaces/zero/torch/patching_legacy.py

@@ -1,266 +0,0 @@
-"""
-"""
-# pyright: reportPrivateImportUsage=false
-
-from __future__ import annotations
-
-import multiprocessing
-import os
-from concurrent.futures import ProcessPoolExecutor
-from contextlib import suppress
-from functools import partial
-from types import SimpleNamespace
-from typing import Any
-from typing import Callable
-from typing import Optional
-from typing import Tuple
-
-import torch
-from torch.utils.weak import WeakTensorKeyDictionary
-
-from ...config import Config
-from . import bitsandbytes
-
-
-# Nvidia A100.80G MIG (drivers 535) / Torch 2.2.0
-CUDA_DEVICE_NAME = 'NVIDIA A100-SXM4-80GB MIG 3g.40gb'
-CUDA_TOTAL_MEMORY = 42144366592
-CUDA_MEM_GET_INFO = (41911451648, CUDA_TOTAL_MEMORY)
-CUDA_DEVICE_CAPABILITY = (8, 0)
-CUDA_DEVICE_PROPERTIES = SimpleNamespace(name=CUDA_DEVICE_NAME, major=8, minor=0, total_memory=CUDA_TOTAL_MEMORY, multi_processor_count=42)
-
-GENERIC_METHOD_NAMES = [
-    'arange',
-    'as_tensor',
-    'asarray',
-    'bartlett_window',
-    'blackman_window',
-    'empty',
-    'empty_like',
-    'empty_strided',
-    'eye',
-    'full',
-    'full_like',
-    'hamming_window',
-    'hann_window',
-    'kaiser_window',
-    'linspace',
-    'logspace',
-    'ones',
-    'ones_like',
-    'rand',
-    'rand_like',
-    'randint',
-    'randint_like',
-    'randn',
-    'randn_like',
-    'randperm',
-    'range',
-    'sparse_bsc_tensor',
-    'sparse_bsr_tensor',
-    'sparse_compressed_tensor',
-    'sparse_coo_tensor',
-    'sparse_csc_tensor',
-    'sparse_csr_tensor',
-    'tensor',
-    'tril_indices',
-    'triu_indices',
-    'zeros',
-    'zeros_like',
-]
-
-
-TO_CUDA = (torch.device('cuda'), None, False, None)
-
-_tensor__deepcopy__ = torch.Tensor.__deepcopy__
-_tensor_to         = torch.Tensor.to
-_tensor_cuda       = torch.Tensor.cuda
-_tensor_cpu        = torch.Tensor.cpu
-_torch_generics    = {name: getattr(torch, name) for name in GENERIC_METHOD_NAMES}
-_cuda_init         = torch._C._cuda_init
-_cuda_available      = torch.cuda.is_available
-_cuda_device_count   = torch.cuda.device_count
-_cuda_current_device = torch.cuda.current_device
-_cuda_mem_get_info   = torch.cuda.mem_get_info
-_cuda_get_device_capability   = torch.cuda.get_device_capability
-_cuda_get_device_properties   = torch.cuda.get_device_properties
-_cuda_get_device_name         = torch.cuda.get_device_name
-
-TensorToArgs = Tuple[Optional[torch.device], Optional[torch.dtype], bool, Optional[torch.memory_format]]
-
-to_ops: dict[torch.Tensor, TensorToArgs] = WeakTensorKeyDictionary() # type: ignore
-
-def _tensor_new_register(*args, **kwargs):
-    new_tensor: torch.Tensor = torch._C._TensorBase.__new__(*args, **kwargs)
-    if (base_tensor := new_tensor._base) is not None:
-        if base_tensor in to_ops:
-            to_ops[new_tensor] = to_ops[base_tensor]
-    return new_tensor
-
-def _tensor_deepcopy_register(self: torch.Tensor, memo):
-    new_tensor = _tensor__deepcopy__(self, memo)
-    if isinstance(new_tensor, torch.Tensor):
-        if self in to_ops:
-            to_ops[new_tensor] = to_ops[self]
-    return new_tensor
-
-@property
-def _tensor_device_property(self: torch.Tensor):
-    if self in to_ops:
-        return torch.device(type='cuda', index=0)
-    del torch.Tensor.device
-    try:
-        return self.device
-    finally:
-        torch.Tensor.device = _tensor_device_property # type: ignore
-
-@property
-def _tensor_dtype_property(self: torch.Tensor):
-    if self in to_ops:
-        if (to_dtype := to_ops[self][1]) is not None:
-            return to_dtype
-    del torch.Tensor.dtype
-    try:
-        return self.dtype
-    finally:
-        torch.Tensor.dtype = _tensor_dtype_property # type: ignore
-
-def _to_op_register(self: torch.Tensor, *args, **kwargs):
-    parsed = torch._C._nn._parse_to(*args, **kwargs)
-    device, dtype, *_ = parsed
-    try:
-        to_args = to_ops.pop(self)
-    except KeyError:
-        to_args = None
-    if device is None: # pyright: ignore [reportUnnecessaryComparison]
-        if to_args is not None:
-            to_ops[self] = (to_args[0], dtype, *to_args[2:])
-            return self
-        return _tensor_to(self, *args, **kwargs)
-    if device.type != 'cuda':
-        if to_args is not None:
-            if (to_dtype := to_args[1]) is not None:
-                kwargs = {'dtype': to_dtype, **kwargs}
-        return _tensor_to(self, *args, **kwargs)
-    to_ops[self] = parsed
-    return self
-
-def _cuda_op_arg_check(device: torch.device | int | str | None) -> bool:
-    if device is None:
-        return True
-    if isinstance(device, int):
-        return True
-    if isinstance(device, str):
-        device = torch.device(device)
-    return device.type == 'cuda'
-
-def _cuda_op_register(self: torch.Tensor, device: torch.device | int | str | None = None, **kwargs):
-    if not _cuda_op_arg_check(device):
-        # Let PyTorch handle the fail
-        return _tensor_cuda(self, device, **kwargs)
-    to_ops[self] = TO_CUDA
-    return self
-
-def _cpu_op_remove(self: torch.Tensor, **kwargs):
-    try:
-        to_args = to_ops.pop(self)
-    except KeyError:
-        to_args = None
-    if to_args is not None:
-        if (to_dtype := to_args[1]) is not None:
-            return _tensor_to(self, 'cpu', **{'dtype': to_dtype, **kwargs})
-    return _tensor_cpu(self, **kwargs)
-
-def _cuda_init_raise():
-    raise RuntimeError(
-        "CUDA must not be initialized in the main process "
-        "on Spaces with Stateless GPU environment.\n"
-        "You can look at this Stacktrace to find out "
-        "which part of your code triggered a CUDA init"
-    )
-
-def _generic_method_register(name: str, *args: Any, **kwargs: Any):
-    try:
-        device = torch.device(kwargs.get('device', "cpu"))
-    except Exception:
-        return _torch_generics[name](*args, **kwargs)
-    if device.type != 'cuda':
-        return _torch_generics[name](*args, **kwargs)
-    tensor = _torch_generics[name](*args, **{**kwargs, 'device': "cpu"})
-    to_ops[tensor] = TO_CUDA
-    return tensor
-
-def patch():
-    torch.Tensor.__deepcopy__ = _tensor_deepcopy_register
-    torch.Tensor.__new__      = _tensor_new_register # pyright: ignore [reportAttributeAccessIssue]
-    torch.Tensor.to         = _to_op_register   # type: ignore
-    torch.Tensor.cuda       = _cuda_op_register # type: ignore
-    torch.Tensor.cpu        = _cpu_op_remove # type: ignore
-    if Config.zero_patch_torch_device:
-        torch.Tensor.device = _tensor_device_property # type: ignore
-        torch.Tensor.dtype  = _tensor_dtype_property # pyright: ignore [reportAttributeAccessIssue]
-    for name in GENERIC_METHOD_NAMES:
-        setattr(torch, name, partial(_generic_method_register, name))
-    torch._C._cuda_init     = _cuda_init_raise
-    torch.cuda.is_available   = lambda: True
-    torch.cuda.device_count   = lambda: 1
-    torch.cuda.current_device = lambda: 0
-    torch.cuda.mem_get_info   = lambda *args, **kwargs: CUDA_MEM_GET_INFO
-    torch.cuda.get_device_capability = lambda *args, **kwargs: CUDA_DEVICE_CAPABILITY
-    torch.cuda.get_device_properties = lambda *args, **kwargs: CUDA_DEVICE_PROPERTIES
-    torch.cuda.get_device_name       = lambda *args, **kwargs: CUDA_DEVICE_NAME
-    bitsandbytes.patch()
-
-def unpatch():
-    torch.Tensor.__deepcopy__ = _tensor__deepcopy__
-    with suppress(AttributeError):
-        del torch.Tensor.__new__
-    torch.Tensor.to         = _tensor_to
-    torch.Tensor.cuda       = _tensor_cuda
-    torch.Tensor.cpu        = _tensor_cpu
-    with suppress(AttributeError):
-        del torch.Tensor.device
-    with suppress(AttributeError):
-        del torch.Tensor.dtype
-    for name in GENERIC_METHOD_NAMES:
-        setattr(torch, name, _torch_generics[name])
-    torch._C._cuda_init     = _cuda_init
-    torch.cuda.is_available   = _cuda_available
-    torch.cuda.device_count   = _cuda_device_count
-    torch.cuda.current_device = _cuda_current_device
-    torch.cuda.mem_get_info   = _cuda_mem_get_info
-    torch.cuda.get_device_capability = _cuda_get_device_capability
-    torch.cuda.get_device_properties = _cuda_get_device_properties
-    torch.cuda.get_device_name       = _cuda_get_device_name
-    bitsandbytes.unpatch()
-
-def pack():
-    pass
-
-def init(nvidia_uuid: str):
-    os.environ['CUDA_VISIBLE_DEVICES'] = nvidia_uuid
-    torch.Tensor([0]).cuda() # CUDA init
-
-def size():
-    return 0
-
-def move(callback: Callable[[int]] | None = None):
-    for op in to_ops.items():
-        tensor, parsed_args = op
-        _, dtype, _, memory_format = parsed_args
-        tensor.data = _tensor_to(tensor,
-            device='cuda',
-            dtype=dtype,
-            memory_format=memory_format,
-        ) # type: ignore
-    bitsandbytes.move()
-    torch.cuda.synchronize()
-
-def is_in_bad_fork():
-    with ProcessPoolExecutor(mp_context=multiprocessing.get_context('fork')) as e:
-        f = e.submit(torch.cuda._is_in_bad_fork)
-        return f.result()
-
-def disable_cuda_intercept():
-    torch.Tensor.to   = _tensor_to
-    torch.Tensor.cuda = _tensor_cuda

spaces/zero/torch/types.py

@@ -1,23 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-from typing import NamedTuple
-
-import torch
-
-
-class AliasId(NamedTuple):
-    data_ptr: int
-    dtype: torch.dtype
-    shape: tuple[int, ...]
-    stride: tuple[int, ...]
-
-    @classmethod
-    def from_tensor(cls, tensor: torch.Tensor):
-        return cls(
-            tensor.data_ptr(),
-            tensor.dtype,
-            tensor.shape,
-            tensor.stride(),
-        )

spaces/zero/tqdm.py

@@ -1,24 +0,0 @@
-"""
-"""
-
-from multiprocessing.synchronize import RLock as MultiprocessingRLock
-
-
-try:
-    from tqdm import tqdm as _tqdm
-except ImportError: # pragma: no cover
-    _tqdm = None
-
-
-def remove_tqdm_multiprocessing_lock():
-    if _tqdm is None: # pragma: no cover
-        return
-    tqdm_lock = _tqdm.get_lock()
-    assert tqdm_lock.__class__.__name__ == 'TqdmDefaultWriteLock'
-    tqdm_lock.locks = [
-        lock for lock in tqdm_lock.locks
-        if not isinstance(lock, MultiprocessingRLock)
-    ]
-
-
-tqdm = _tqdm

spaces/zero/types.py

@@ -1,49 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-
-from dataclasses import dataclass
-from datetime import timedelta
-from typing import Any
-from typing import Dict
-from typing import Tuple
-from typing import TypedDict
-from typing_extensions import Callable
-from typing_extensions import Generic
-from typing_extensions import ParamSpec
-from typing_extensions import TypeAlias
-from typing_extensions import TypeVar
-
-
-Params = Tuple[Tuple[object, ...], Dict[str, Any]]
-Res = TypeVar('Res')
-Param = ParamSpec('Param')
-
-class EmptyKwargs(TypedDict):
-    pass
-
-@dataclass
-class OkResult(Generic[Res]):
-    value: Res
-@dataclass
-class ExceptionResult:
-    value: Exception
-@dataclass
-class AbortedResult:
-    pass
-@dataclass
-class EndResult:
-    pass
-@dataclass
-class GradioQueueEvent:
-    method_name: str
-    args: tuple[Any, ...]
-    kwargs: dict[str, Any]
-
-RegularResQueueResult:   TypeAlias = "OkResult[Res] | ExceptionResult | GradioQueueEvent"
-GeneratorResQueueResult: TypeAlias = "OkResult[Res] | ExceptionResult | EndResult | GradioQueueEvent"
-YieldQueueResult:        TypeAlias = "OkResult[Res] | ExceptionResult | EndResult | AbortedResult"
-
-Duration:        TypeAlias = "int | timedelta"
-DynamicDuration: TypeAlias = "Duration | Callable[Param, Duration] | None"

spaces/zero/wrappers.py

@@ -1,418 +0,0 @@
-"""
-"""
-from __future__ import annotations
-
-import multiprocessing
-import os
-import signal
-import traceback
-import warnings
-from concurrent.futures import ThreadPoolExecutor
-from contextlib import nullcontext
-from contextvars import copy_context
-from datetime import timedelta
-from functools import partial
-from functools import wraps
-from multiprocessing.context import ForkProcess
-from pickle import PicklingError
-from queue import Empty
-from queue import Queue as ThreadQueue
-from threading import Thread
-from typing import TYPE_CHECKING
-from typing import Callable
-from typing import Generator
-from typing import Generic
-from typing_extensions import assert_never
-
-import psutil
-
-from ..config import Config
-from ..utils import debug
-from ..utils import drop_params
-from ..utils import gradio_request_var
-from ..utils import SimpleQueue as Queue
-from . import client
-from . import torch
-from .api import AllowToken
-from .api import NvidiaIndex
-from .api import NvidiaUUID
-from .gradio import GradioPartialContext
-from .gradio import get_server_port
-from .gradio import patch_gradio_queue
-from .gradio import try_process_queue_event
-from .tqdm import remove_tqdm_multiprocessing_lock
-from .tqdm import tqdm
-from .types import * # TODO: Please don't do that
-
-
-GENERATOR_GLOBAL_TIMEOUT = 20 * 60
-
-SPAWN_PROGRESS_CLEANUP = 0.1
-SPAWN_PROGRESS_INIT = 0.1
-
-
-Process = multiprocessing.get_context('fork').Process
-forked = False
-
-
-class Worker(Generic[Res]):
-    process: ForkProcess
-    arg_queue: Queue[tuple[Params, GradioPartialContext]]
-    res_queue: Queue[Res | None]
-    _sentinel: Thread
-
-    def __init__(
-        self,
-        target: Callable[[
-            Queue[tuple[Params, GradioPartialContext]],
-            Queue[Res | None],
-            AllowToken,
-            NvidiaUUID,
-            list[int],
-        ], None],
-        allow_token: str,
-        nvidia_uuid: str,
-    ):
-        self._sentinel = Thread(target=self._close_on_exit, daemon=True)
-        self.arg_queue = Queue()
-        self.res_queue = Queue()
-        debug(f"{self.arg_queue._writer.fileno()=}") # pyright: ignore [reportAttributeAccessIssue]
-        debug(f"{self.res_queue._writer.fileno()=}") # pyright: ignore [reportAttributeAccessIssue]
-        if (server_port := get_server_port()) is not None:
-            fds = [c.fd for c in psutil.Process().connections() if c.laddr.port == server_port]
-            debug(f"{fds=}")
-        else:
-            warnings.warn("Using a ZeroGPU function outside of Gradio caching or request might block the app")
-            fds = []
-        args = self.arg_queue, self.res_queue, allow_token, nvidia_uuid, fds
-        if TYPE_CHECKING:
-            target(*args)
-        self.process = Process(
-            target=target,
-            args=args,
-            daemon=True,
-        )
-        self.process.start()
-        self._sentinel.start()
-
-    def _close_on_exit(self):
-        self.process.join()
-        self.arg_queue.close()
-        self.res_queue.wlock_release()
-        self.res_queue.put(None)
-
-
-def worker_init(
-    res_queue: Queue[RegularResQueueResult | None] | Queue[GeneratorResQueueResult | None],
-    allow_token: str,
-    nvidia_uuid: str,
-    fds: list[int],
-) -> None | ExceptionResult:
-    # Immediately close file descriptors
-    for fd in fds:
-        try:
-            os.close(fd)
-        except Exception as e: # pragma: no cover
-            if isinstance(e, OSError) and e.errno == 9:
-                continue
-            traceback.print_exc()
-            return ExceptionResult(e)
-    progress = nullcontext()
-    if tqdm is not None and Config.zero_gpu_v2:
-        progress = tqdm(total=100, desc="ZeroGPU init", file=open(os.devnull, 'w'))
-    try: # Unrecoverable init part
-        patch_gradio_queue(res_queue)
-        with progress as progress:
-            current_progress = 0 # Gradio does not support float progress updates
-            def update(n: float):
-                nonlocal current_progress
-                current_progress += n
-                if progress is not None:
-                    progress.update(round(current_progress * 100) - progress.n)
-            client.allow(allow_token)
-            update(SPAWN_PROGRESS_CLEANUP)
-            torch.unpatch()
-            torch.init(nvidia_uuid)
-            update(SPAWN_PROGRESS_INIT)
-            callback = None
-            if (transfer_size := torch.size()) > 0:
-                remaining = 1 - (SPAWN_PROGRESS_CLEANUP + SPAWN_PROGRESS_INIT)
-                callback = lambda n: update(n * remaining / transfer_size)
-            torch.move(callback=callback)
-    except Exception as e: # pragma: no cover
-        traceback.print_exc()
-        return ExceptionResult(e)
-    try:
-        remove_tqdm_multiprocessing_lock()
-    except Exception: # pragma: no cover
-        print("Error while trying to remove tqdm mp_lock:")
-        traceback.print_exc()
-
-
-def process_duration(duration: Duration | None):
-    if duration is None or isinstance(duration, timedelta):
-        return duration
-    return timedelta(seconds=duration)
-
-
-def static_duration(duration: DynamicDuration[Param], *args: Param.args, **kwargs: Param.kwargs):
-    if not callable(duration):
-        return duration
-    return duration(*args, **kwargs)
-
-
-def regular_function_wrapper(
-    task: Callable[Param, Res],
-    duration: DynamicDuration[Param],
-) -> Callable[Param, Res]:
-
-    import gradio as gr
-
-    request_var = gradio_request_var()
-    workers: dict[NvidiaIndex, Worker[RegularResQueueResult[Res]]] = {}
-    task_id = id(task)
-
-    @wraps(task)
-    def gradio_handler(*args: Param.args, **kwargs: Param.kwargs) -> Res:
-
-        if forked:
-            return task(*args, **kwargs)
-
-        request = request_var.get()
-        duration_ = static_duration(duration, *args, **kwargs)
-        duration_ = process_duration(duration_)
-        schedule_response = client.schedule(task_id=task_id, request=request, duration=duration_)
-        allow_token = schedule_response.allowToken
-        nvidia_index = schedule_response.nvidiaIndex
-        nvidia_uuid = schedule_response.nvidiaUUID
-        release = partial(client.release, allow_token)
-
-        try:
-            worker = workers.pop(nvidia_index)
-        except KeyError:
-            worker = None
-
-        if worker is not None and worker.process.is_alive() and schedule_response.idle:
-            assert worker.arg_queue.empty()
-            assert worker.res_queue.empty()
-        else:
-            worker = Worker(thread_wrapper, allow_token, nvidia_uuid)
-
-        try:
-            worker.arg_queue.put(((args, kwargs), GradioPartialContext.get()))
-        except PicklingError: # TODO: detailed serialization diagnostic
-            release(fail=True)
-            raise
-
-        while True:
-            res = worker.res_queue.get()
-            if res is None:
-                release(fail=True, allow_404=True)
-                raise gr.Error("GPU task aborted")
-            if isinstance(res, ExceptionResult):
-                release(fail=True)
-                raise res.value
-            if isinstance(res, OkResult):
-                release()
-                workers[nvidia_index] = worker
-                return res.value
-            if isinstance(res, GradioQueueEvent):
-                try_process_queue_event(res.method_name, *res.args, **res.kwargs)
-                continue
-            assert_never(res)
-
-
-    def thread_wrapper(
-        arg_queue: Queue[tuple[Params, GradioPartialContext]],
-        res_queue: Queue[RegularResQueueResult[Res] | None],
-        allow_token: str,
-        nvidia_uuid: str,
-        fds: list[int],
-    ):
-        global forked
-        forked = True
-        signal.signal(signal.SIGTERM, drop_params(arg_queue.close))
-        initialized = False
-        while True:
-            try:
-                (args, kwargs), gradio_context = arg_queue.get()
-            except OSError:
-                break
-            if not initialized:
-                if (res := worker_init(
-                    res_queue=res_queue,
-                    allow_token=allow_token,
-                    nvidia_uuid=nvidia_uuid,
-                    fds=fds,
-                )) is not None:
-                    res_queue.put(res)
-                    return
-                initialized = True
-            GradioPartialContext.apply(gradio_context)
-            context = copy_context()
-            with ThreadPoolExecutor() as executor:
-                future = executor.submit(context.run, task, *args, **kwargs) # type: ignore
-            try:
-                res = future.result()
-            except Exception as e:
-                traceback.print_exc()
-                res = ExceptionResult(e)
-            else:
-                res = OkResult(res)
-            try:
-                res_queue.put(res)
-            except PicklingError as e:
-                res_queue.put(ExceptionResult(e))
-
-    # https://github.com/python/cpython/issues/91002
-    if not hasattr(task, '__annotations__'):
-        gradio_handler.__annotations__ = {}
-
-    return gradio_handler
-
-
-def generator_function_wrapper(
-    task: Callable[Param, Generator[Res, None, None]],
-    duration: DynamicDuration[Param],
-) -> Callable[Param, Generator[Res, None, None]]:
-
-    import gradio as gr
-
-    request_var = gradio_request_var()
-    workers: dict[NvidiaIndex, Worker[GeneratorResQueueResult[Res]]] = {}
-    task_id = id(task)
-
-    @wraps(task)
-    def gradio_handler(*args: Param.args, **kwargs: Param.kwargs) -> Generator[Res, None, None]:
-
-        if forked:
-            yield from task(*args, **kwargs)
-            return
-
-        request = request_var.get()
-        duration_ = static_duration(duration, *args, **kwargs)
-        duration_ = process_duration(duration_)
-        schedule_response = client.schedule(task_id=task_id, request=request, duration=duration_)
-        allow_token = schedule_response.allowToken
-        nvidia_index = schedule_response.nvidiaIndex
-        nvidia_uuid = schedule_response.nvidiaUUID
-        release = partial(client.release, allow_token)
-
-        try:
-            worker = workers.pop(nvidia_index)
-        except KeyError:
-            worker = None
-
-        if worker is not None and worker.process.is_alive() and schedule_response.idle:
-            assert worker.arg_queue.empty()
-            assert worker.res_queue.empty()
-        else:
-            worker = Worker(thread_wrapper, allow_token, nvidia_uuid)
-
-        try:
-            worker.arg_queue.put(((args, kwargs), GradioPartialContext.get()))
-        except PicklingError: # TODO: detailed serialization diagnostic
-            release(fail=True)
-            raise
-
-        yield_queue: ThreadQueue[YieldQueueResult[Res]] = ThreadQueue()
-        def fill_yield_queue(worker: Worker[GeneratorResQueueResult[Res]]):
-            while True:
-                res = worker.res_queue.get()
-                if res is None:
-                    release(fail=True, allow_404=True)
-                    yield_queue.put(AbortedResult())
-                    return
-                if isinstance(res, ExceptionResult):
-                    release(fail=True)
-                    yield_queue.put(ExceptionResult(res.value))
-                    return
-                if isinstance(res, EndResult):
-                    release()
-                    workers[nvidia_index] = worker
-                    yield_queue.put(EndResult())
-                    return
-                if isinstance(res, OkResult):
-                    yield_queue.put(OkResult(res.value))
-                    continue
-                if isinstance(res, GradioQueueEvent): # pragma: no cover (not working properly on Gradio side)
-                    try_process_queue_event(res.method_name, *res.args, **res.kwargs)
-                    continue
-                debug(f"fill_yield_queue: assert_never({res=})")
-                assert_never(res)
-        from typing_extensions import assert_never
-        with ThreadPoolExecutor() as e:
-            f = e.submit(copy_context().run, fill_yield_queue, worker)
-            f.add_done_callback(lambda _: debug("fill_yield_queue DONE"))
-            while True:
-                try:
-                    res = yield_queue.get(timeout=GENERATOR_GLOBAL_TIMEOUT)
-                except Empty: # pragma: no cover
-                    debug(f"yield_queue TIMEOUT ({GENERATOR_GLOBAL_TIMEOUT=})")
-                    raise
-                if isinstance(res, AbortedResult):
-                    raise gr.Error("GPU task aborted")
-                if isinstance(res, ExceptionResult):
-                    raise res.value
-                if isinstance(res, EndResult):
-                    break
-                if isinstance(res, OkResult):
-                    yield res.value
-                    continue
-                debug(f"gradio_handler: assert_never({res=})")
-                assert_never(res)
-
-
-    def thread_wrapper(
-        arg_queue: Queue[tuple[Params, GradioPartialContext]],
-        res_queue: Queue[GeneratorResQueueResult[Res] | None],
-        allow_token: str,
-        nvidia_uuid: str,
-        fds: list[int],
-    ):
-        global forked
-        forked = True
-        signal.signal(signal.SIGTERM, drop_params(arg_queue.close))
-        initialized = False
-        while True:
-            try:
-                (args, kwargs), gradio_context = arg_queue.get()
-            except OSError:
-                break
-            if not initialized:
-                if (res := worker_init(
-                    res_queue=res_queue,
-                    allow_token=allow_token,
-                    nvidia_uuid=nvidia_uuid,
-                    fds=fds,
-                )) is not None:
-                    res_queue.put(res)
-                    return
-                initialized = True
-            def iterate():
-                gen = task(*args, **kwargs) # type: ignore
-                while True:
-                    try:
-                        res = next(gen)
-                    except StopIteration:
-                        break
-                    except Exception as e:
-                        res_queue.put(ExceptionResult(e))
-                        break
-                    try:
-                        res_queue.put(OkResult(res))
-                    except PicklingError as e:
-                        res_queue.put(ExceptionResult(e))
-                        break
-                    else:
-                        continue
-            GradioPartialContext.apply(gradio_context)
-            with ThreadPoolExecutor() as executor:
-                executor.submit(copy_context().run, iterate)
-            res_queue.put(EndResult())
-
-    # https://github.com/python/cpython/issues/91002
-    if not hasattr(task, '__annotations__'):
-        gradio_handler.__annotations__ = {}
-
-    return gradio_handler

voice_chat.py

@@ -9,6 +9,89 @@ import torch
 import sentencepiece as spm
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download, InferenceClient
+import requests
+from bs4 import BeautifulSoup
+import urllib
+import random
+
+# List of user agents to choose from for requests
+_useragent_list = [
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:66.0) Gecko/20100101 Firefox/66.0',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/109.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36 Edg/111.0.1661.62',
+    'Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/111.0'
+]
+
+def get_useragent():
+    """Returns a random user agent from the list."""
+    return random.choice(_useragent_list)
+
+def extract_text_from_webpage(html_content):
+    """Extracts visible text from HTML content using BeautifulSoup."""
+    soup = BeautifulSoup(html_content, "html.parser")
+    # Remove unwanted tags
+    for tag in soup(["script", "style", "header", "footer", "nav"]):
+        tag.extract()
+    # Get the remaining visible text
+    visible_text = soup.get_text(strip=True)
+    return visible_text
+
+def search(term, num_results=1, lang="en", advanced=True, sleep_interval=0, timeout=5, safe="active", ssl_verify=None):
+    """Performs a Google search and returns the results."""
+    escaped_term = urllib.parse.quote_plus(term)
+    start = 0
+    all_results = []
+
+    # Fetch results in batches
+    while start < num_results:
+        resp = requests.get(
+            url="https://www.google.com/search",
+            headers={"User-Agent": get_useragent()}, # Set random user agent
+            params={
+                "q": term,
+                "num": num_results - start, # Number of results to fetch in this batch
+                "hl": lang,
+                "start": start,
+                "safe": safe,
+            },
+            timeout=timeout,
+            verify=ssl_verify,
+        )
+        resp.raise_for_status() # Raise an exception if request fails
+
+        soup = BeautifulSoup(resp.text, "html.parser")
+        result_block = soup.find_all("div", attrs={"class": "g"})
+
+        # If no results, continue to the next batch
+        if not result_block:
+            start += 1
+            continue
+
+        # Extract link and text from each result
+        for result in result_block:
+            link = result.find("a", href=True)
+            if link:
+                link = link["href"]
+                try:
+                    # Fetch webpage content
+                    webpage = requests.get(link, headers={"User-Agent": get_useragent()})
+                    webpage.raise_for_status()
+                    # Extract visible text from webpage
+                    visible_text = extract_text_from_webpage(webpage.text)
+                    all_results.append({"link": link, "text": visible_text})
+                except requests.exceptions.RequestException as e:
+                    # Handle errors fetching or processing webpage
+                    print(f"Error fetching or processing {link}: {e}")
+                    all_results.append({"link": link, "text": None})
+            else:
+                all_results.append({"link": None, "text": None})
+
+        start += len(result_block) # Update starting index for next batch
+
+    return all_results
 
 # Speech Recognition Model Configuration
 model_name = "neongeckocom/stt_en_citrinet_512_gamma_0_25"
@@ -20,8 +103,8 @@ encoder = ort.InferenceSession(hf_hub_download(model_name, "model.onnx", subfold
 tokenizer = spm.SentencePieceProcessor(hf_hub_download(model_name, "tokenizer.spm", subfolder="onnx"))
 
 # Mistral Model Configuration
-client1 = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
-system_instructions1 = "[SYSTEM] Answer as Real OpenGPT 4o, Made by 'KingNish', Keep conversation very short, clear, friendly and concise. The text provided is a request for a specific type of response from you, the virtual assistant. You will avoid introductions and start answering the query directly, elaborating on all aspects of the request. As an AI-powered assistant, your task is to generate responses that appear as if they are created by the user and not an AI-powered assistant. [USER]"
+client1 = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+system_instructions1 = "<s>[SYSTEM] Answer as Real Jarvis JARVIS, Made by 'Tony Stark', Keep conversation very short, clear, friendly and concise. The text provided is a request for a specific type of response from you, the virtual assistant. The request asks you to provide friendly responses as if You are the character Jarvis, made by 'Tony Stark.' The expectation is that I will avoid introductions and start answering the query directly, elaborating on all aspects of the request. As an AI-powered assistant, your task is to generate responses that appear as if they are created by the user and not an AI-powered assistant. [USER]"
 
 def resample(audio_fp32, sr):
     return soxr.resample(audio_fp32, sr, sample_rate)
@@ -49,14 +132,22 @@ def transcribe(audio_path):
 
     return text
 
-def model(text):
-    formatted_prompt = system_instructions1 + text + "[OpenGPT 4o]"
-    stream = client1.text_generation(formatted_prompt, max_new_tokens=300)
-    return stream[:-4]
+def model(text, web_search):
+    if web_search is True:
+        """Performs a web search, feeds the results to a language model, and returns the answer."""
+        web_results = search(text)
+        web2 = ' '.join([f"Link: {res['link']}\nText: {res['text']}\n\n" for res in web_results])
+        formatted_prompt = system_instructions1 + text + "[WEB]" + str(web2) + "[ANSWER]"
+        stream = client1.text_generation(formatted_prompt, max_new_tokens=512, stream=True, details=True, return_full_text=False)
+        return "".join([response.token.text for response in stream if response.token.text != "</s>"])
+    else:
+        formatted_prompt = system_instructions1 + text + "[JARVIS]"
+        stream = client1.text_generation(formatted_prompt, max_new_tokens=512, stream=True, details=True, return_full_text=False)
+        return "".join([response.token.text for response in stream if response.token.text != "</s>"])
 
-async def respond(audio):
+async def respond(audio, web_search):
     user = transcribe(audio)
-    reply = model(user)
+    reply = model(user, web_search)
     communicate = edge_tts.Communicate(reply)
     with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
         tmp_path = tmp_file.name