Update app.py

app.py

@@ -1,39 +1,138 @@
-
-import os, copy
-os.environ["RWKV_V7_ON"] = '1'
-os.environ["RWKV_JIT_ON"] = '1'
-os.environ["RWKV_CUDA_ON"] = '1' # if '1' then use CUDA kernel for seq mode (much faster)
-
-from rwkv.model import RWKV
-
-import gc, re
-import gradio as gr
-import base64
-from io import BytesIO
+import spaces
+import os, copy, gc, re, sys
+import traceback
 import torch
 import torch.nn.functional as F
 from datetime import datetime
+import gradio as gr
 from huggingface_hub import hf_hub_download
-from pynvml import *
-nvmlInit()
-gpu_h = nvmlDeviceGetHandleByIndex(0)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-ctx_limit = 4000
-gen_limit = 32000
+# Check for CUDA availability
+cuda_available = torch.cuda.is_available()
+device = torch.device("cuda" if cuda_available else "cpu")
+print(f"Using device: {device}")
+
+# Function to set up CUDA environment if available
+def setup_cuda_environment():
+    if not cuda_available:
+        print("CUDA not available, falling back to CPU")
+        os.environ["RWKV_V7_ON"] = '1'
+        os.environ["RWKV_JIT_ON"] = '1'
+        os.environ["RWKV_CUDA_ON"] = '0'
+        return False
+
+    print("CUDA is available, setting up environment")
+    
+    # Try to detect CUDA location automatically
+    possible_cuda_paths = [
+        "/usr/local/cuda",
+        "/opt/cuda",
+        "/usr/lib/cuda",
+        "/usr/cuda",
+        "/usr/local/nvidia/cuda",
+        "/usr/lib/nvidia-cuda-toolkit",
+        "/usr/lib/x86_64-linux-gnu/cuda"
+    ]
+    
+    cuda_found = False
+    for path in possible_cuda_paths:
+        if os.path.exists(path):
+            os.environ["CUDA_HOME"] = path
+            print(f"Found CUDA at: {path}")
+            cuda_found = True
+            break
+    
+    if not cuda_found:
+        # If we can't find the CUDA path but CUDA is available,
+        # try looking for common libraries
+        try:
+            import ctypes
+            cuda_runtime = ctypes.cdll.LoadLibrary("libcudart.so")
+            print("Found CUDA runtime library, proceeding without explicit CUDA_HOME")
+            cuda_found = True
+        except:
+            print("Could not locate CUDA runtime library")
+    
+    # Set RWKV environment variables
+    if cuda_found:
+        os.environ["RWKV_V7_ON"] = '1'
+        os.environ["RWKV_JIT_ON"] = '1'
+        os.environ["RWKV_CUDA_ON"] = '1'
+    else:
+        print("CUDA is available but environment couldn't be set up correctly, falling back to CPU")
+        os.environ["RWKV_V7_ON"] = '1'
+        os.environ["RWKV_JIT_ON"] = '1'
+        os.environ["RWKV_CUDA_ON"] = '0'
+        return False
+    
+    return cuda_found
 
-########################## text rwkv ################################################################
-from rwkv.utils import PIPELINE, PIPELINE_ARGS
+# Initialize NVML for GPU monitoring if available
+has_nvml = False
+if cuda_available:
+    try:
+        from pynvml import nvmlInit, nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo
+        nvmlInit()
+        gpu_h = nvmlDeviceGetHandleByIndex(0)
+        has_nvml = True
+        print("NVML initialized for GPU monitoring")
+    except:
+        print("NVML not available, GPU monitoring disabled")
 
+# Set up CUDA environment
+use_cuda = setup_cuda_environment()
+
+# Model parameters
+ctx_limit = 4000
+gen_limit = 32000
 title_v6 = "rwkv7-g1-0.1b-20250307-ctx4096"
-model_path_v6 = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title_v6}.pth")
-model_v6 = RWKV(model=model_path_v6.replace('.pth',''), strategy='cuda fp16')
-pipeline_v6 = PIPELINE(model_v6, "rwkv_vocab_v20230424")
 
-args = model_v6.args
+# Load RWKV with fallback mechanisms
+try:
+    # First try importing normally
+    from rwkv.model import RWKV
+    from rwkv.utils import PIPELINE, PIPELINE_ARGS
+    print("RWKV imported successfully")
+except Exception as e:
+    print(f"Error importing RWKV: {e}")
+    print("Attempting fallback import method...")
+    
+    # Fallback method - reinstall the package
+    try:
+        import subprocess
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "rwkv"])
+        from rwkv.model import RWKV
+        from rwkv.utils import PIPELINE, PIPELINE_ARGS
+        print("RWKV imported after reinstall")
+    except Exception as e:
+        print(f"Failed to import RWKV after reinstall: {e}")
+        raise
 
-penalty_decay = 0.996
+# Download and initialize the model
+try:
+    print(f"Downloading model {title_v6}...")
+    model_path_v6 = hf_hub_download(repo_id="BlinkDL/rwkv7-g1", filename=f"{title_v6}.pth")
+    print(f"Model downloaded to {model_path_v6}")
+    
+    # Select strategy based on available hardware
+    strategy = 'cuda fp16' if use_cuda else 'cpu fp32'
+    print(f"Using strategy: {strategy}")
+    
+    # Initialize model with appropriate strategy
+    model_v6 = RWKV(model=model_path_v6.replace('.pth',''), strategy=strategy)
+    pipeline_v6 = PIPELINE(model_v6, "rwkv_vocab_v20230424")
+    args = model_v6.args
+    print("Model initialized successfully")
+    
+    model_loaded = True
+except Exception as e:
+    print(f"Error loading model: {e}")
+    traceback.print_exc()
+    model_loaded = False
 
+# Text generation parameters
+penalty_decay = 0.996
+@spaces.GPU
 def generate_prompt(instruction, input=""):
     instruction = instruction.strip().replace('\r\n','\n').replace('\n\n','\n')
     input = input.strip().replace('\r\n','\n').replace('\n\n','\n')
@@ -55,57 +154,74 @@ def evaluate(
     presencePenalty = 0.1,
     countPenalty = 0.1,
 ):
-    args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_p),
-                     alpha_frequency = countPenalty,
-                     alpha_presence = presencePenalty,
-                     token_ban = [], # ban the generation of some tokens
-                     token_stop = [0]) # stop generation whenever you see any token here
-    ctx = ctx.strip()
-    all_tokens = []
-    out_last = 0
-    out_str = ''
-    occurrence = {}
-    state = None
-    for i in range(int(token_count)):
-
-        input_ids = pipeline_v6.encode(ctx)[-ctx_limit:] if i == 0 else [token]
-        out, state = model_v6.forward(input_ids, state)
-        for n in occurrence:
-            out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
-
-        token = pipeline_v6.sample_logits(out, temperature=args.temperature, top_p=args.top_p)
-        if token in args.token_stop:
-            break
-        all_tokens += [token]
-        for xxx in occurrence:
-            occurrence[xxx] *= penalty_decay
+    if not model_loaded:
+        yield "Error: Model failed to load. Please check logs for details."
+        return
+    
+    try:
+        args = PIPELINE_ARGS(temperature = max(0.2, float(temperature)), top_p = float(top_p),
+                         alpha_frequency = countPenalty,
+                         alpha_presence = presencePenalty,
+                         token_ban = [], # ban the generation of some tokens
+                         token_stop = [0]) # stop generation whenever you see any token here
+        ctx = ctx.strip()
+        all_tokens = []
+        out_last = 0
+        out_str = ''
+        occurrence = {}
+        state = None
+        for i in range(int(token_count)):
+    
+            input_ids = pipeline_v6.encode(ctx)[-ctx_limit:] if i == 0 else [token]
+            out, state = model_v6.forward(input_ids, state)
+            for n in occurrence:
+                out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
+    
+            token = pipeline_v6.sample_logits(out, temperature=args.temperature, top_p=args.top_p)
+            if token in args.token_stop:
+                break
+            all_tokens += [token]
+            for xxx in occurrence:
+                occurrence[xxx] *= penalty_decay
+                
+            ttt = pipeline_v6.decode([token])
+            www = 1
+            if ttt in ' \t0123456789':
+                www = 0
+            if token not in occurrence:
+                occurrence[token] = www
+            else:
+                occurrence[token] += www
+                
+            tmp = pipeline_v6.decode(all_tokens[out_last:])
+            if '\ufffd' not in tmp:
+                out_str += tmp
+                yield out_str.strip()
+                out_last = i + 1
+    
+        # Log GPU info if available
+        if use_cuda and has_nvml:
+            try:
+                gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
+                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                print(f'{timestamp} - vram total: {gpu_info.total/1024**2:.2f}MB, used: {gpu_info.used/1024**2:.2f}MB, free: {gpu_info.free/1024**2:.2f}MB')
+            except:
+                print("Error getting GPU info")
+        
+        # Clean up to free memory
+        del out
+        del state
+        gc.collect()
+        if use_cuda:
+            torch.cuda.empty_cache()
             
-        ttt = pipeline_v6.decode([token])
-        www = 1
-        if ttt in ' \t0123456789':
-            www = 0
-        #elif ttt in '\r\n,.;?!"\':+-*/=#@$%^&_`~|<>\\()[]{}，。；“”：？！（）【】':
-        #    www = 0.5
-        if token not in occurrence:
-            occurrence[token] = www
-        else:
-            occurrence[token] += www
-            
-        tmp = pipeline_v6.decode(all_tokens[out_last:])
-        if '\ufffd' not in tmp:
-            out_str += tmp
-            yield out_str.strip()
-            out_last = i + 1
-
-    gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
-    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-    print(f'{timestamp} - vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}')
-    del out
-    del state
-    gc.collect()
-    torch.cuda.empty_cache()
-    yield out_str.strip()
+        yield out_str.strip()
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        traceback.print_exc()
+        yield f"Error during generation: {str(e)}"
 
+# Example prompts
 examples = [
     ["User: simulate SpaceX mars landing using python\n\nAssistant: <think", gen_limit, 1, 0.3, 0.5, 0.5],
     [generate_prompt("Please give the pros and cons of hodl versus active trading."), gen_limit, 1, 0.3, 0.5, 0.5],
@@ -119,12 +235,16 @@ examples = [
     ["En una pequeña aldea escondida entre las montañas de Andalucía, donde las calles aún conservaban el eco de antiguas leyendas, vivía un joven llamado Alejandro.", gen_limit, 1, 0.3, 0.5, 0.5],
     ["Dans le cœur battant de Paris, sous le ciel teinté d'un crépuscule d'or et de pourpre, se tenait une petite librairie oubliée par le temps.", gen_limit, 1, 0.3, 0.5, 0.5],
     ["في تطور مذهل وغير مسبوق، أعلنت السلطات المحلية في العاصمة عن اكتشاف أثري قد يغير مجرى التاريخ كما نعرفه.", gen_limit, 1, 0.3, 0.5, 0.5],
-    ['''“当然可以，大宇宙不会因为这五公斤就不坍缩了。”关一帆说，他还有一个没说出来的想法：也许大宇宙真的会因为相差一个原子的质量而由封闭转为开放。大自然的精巧有时超出想象，比如生命的诞生，就需要各项宇宙参数在几亿亿分之一精度上的精确配合。但程心仍然可以留下她的生态球，因为在那无数文明创造的无数小宇宙中，肯定有相当一部分不响应回归运动的号召，所以，大宇宙最终被夺走的质量至少有几亿吨，甚至可能是几亿亿亿吨。\n但愿大宇宙能够忽略这个误差。\n程心和关一帆进入了飞船，智子最后也进来了。她早就不再穿那身华丽的和服了，她现在身着迷彩服，再次成为一名轻捷精悍的战士，她的身上佩带着许多武器和生存装备，最引人注目的是那把插在背后的武士刀。\n“放心，我在，你们就在！”智子对两位人类朋友说。\n聚变发动机启动了，推进器发出幽幽的蓝光，''', gen_limit, 1, 0.3, 0.5, 0.5],
+    ['''"当然可以，大宇宙不会因为这五公斤就不坍缩了。"关一帆说，他还有一个没说出来的想法：也许大宇宙真的会因为相差一个原子的质量而由封闭转为开放。大自然的精巧有时超出想象，比如生命的诞生，就需要各项宇宙参数在几亿亿分之一精度上的精确配合。但程心仍然可以留下她的生态球，因为在那无数文明创造的无数小宇宙中，肯定有相当一部分不响应回归运动的号召，所以，大宇宙最终被夺走的质量至少有几亿吨，甚至可能是几亿亿亿吨。\n但愿大宇宙能够忽略这个误差。\n程心和关一帆进入了飞船，智子最后也进来了。她早就不再穿那身华丽的和服了，她现在身着迷彩服，再次成为一名轻捷精悍的战士，她的身上佩带着许多武器和生存装备，最引人注目的是那把插在背后的武士刀。\n"放心，我在，你们就在！"智子对两位人类朋友说。\n聚变发动机启动了，推进器发出幽幽的蓝光，''', gen_limit, 1, 0.3, 0.5, 0.5],
 ]
 
 ##################################################################################################################
+# Create Gradio UI
 with gr.Blocks(title=title_v6) as demo:
-    gr.HTML(f"<div style=\"text-align: center;\">\n<h1>{title_v6}</h1>\n</div>")
+    model_status = "✅ Model loaded successfully" if model_loaded else "❌ Model failed to load"
+    device_status = f"Using {'CUDA' if use_cuda else 'CPU'}"
+    
+    gr.HTML(f"<div style=\"text-align: center;\">\n<h1>{title_v6}</h1>\n<p>{model_status} - {device_status}</p>\n</div>")
 
     with gr.Tab("=== Base Model (Raw Generation) ==="):
         gr.Markdown(f'This is [RWKV7 G1](https://huggingface.co/BlinkDL/rwkv7-g1) 0.1B (!!!) L12-D768 reasoning base LM - an attention-free pure RNN [RWKV-LM](https://github.com/BlinkDL/RWKV-LM). Supports 100+ world languages and code. Check [400+ Github RWKV projects](https://github.com/search?o=desc&p=1&q=rwkv&s=updated&type=Repositories). *** Can try examples (bottom of page) *** (can edit them). Demo limited to ctxlen {ctx_limit}.')
@@ -146,5 +266,7 @@ with gr.Blocks(title=title_v6) as demo:
         clear.click(lambda: None, [], [output])
         data.click(lambda x: x, [data], [prompt, token_count, temperature, top_p, presence_penalty, count_penalty])
 
+# Launch the app
+print("Starting Gradio app...")
 demo.queue(concurrency_count=1, max_size=10)
 demo.launch(share=False)