"CUDA must not be initialized in main process" Error during inference

#1
by zhiqiulin - opened

I see other spaces initializing the model and moving it to device("cuda") outside of the @spaces.GPU decorated function which I assume to be the main process. However, when I try to do that I receive the following error. I've tried to initialize the model within @spaces.GPU as well but I get the same error.

File "/usr/local/lib/python3.10/site-packages/spaces/zero/wrappers.py", line 171, in gradio_handler
    res = worker.res_queue.get()
  File "/usr/local/lib/python3.10/multiprocessing/queues.py", line 367, in get
    return _ForkingPickler.loads(res)
  File "/usr/local/lib/python3.10/site-packages/torch/multiprocessing/reductions.py", line 147, in rebuild_cuda_tensor
    torch.cuda._lazy_init()
  File "/usr/local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 302, in _lazy_init
    torch._C._cuda_init()
  File "/usr/local/lib/python3.10/site-packages/spaces/zero/torch.py", line 181, in _cuda_init_raise
    raise RuntimeError(
RuntimeError: CUDA must not be initialized in the main process on Spaces with Stateless GPU environment.
You can look at this Stacktrace to find out which part of your code triggered a CUDA init

My app.py looks like

import gradio as gr
import spaces
import torch
torch.jit.script = lambda f: f  # Avoid script error in lambda

from t2v_metrics import VQAScore, list_all_vqascore_models

# Global model variable, but do not initialize or move to CUDA here
model_pipe = VQAScore(model="clip-flant5-xl", device="cuda")  # our recommended scoring model

@spaces.GPU(duration = 20)
def generate(model_name, image, text):
    
    result = model_pipe(images=[image], texts=[text])  # Perform the model inference
    
    return result  # Return the result

demo = gr.Interface(
    fn=generate,  # function to call
    inputs=[gr.Dropdown(["clip-flant5-xl", "clip-flant5-xxl"], label="Model Name"), gr.Image(type="filepath"), gr.Textbox(label="Prompt")],  # define the types of inputs
    outputs="number",  # define the type of output
    title="VQAScore",  # title of the app
    description="This model evaluates the similarity between an image and a text prompt."
)

demo.queue()
demo.launch()
zhiqiulin changed discussion status to closed

Sign up or log in to comment