Tryin in Colab

Hi,
I'm trying to make it work in Google Colab, I got all the way to make the gradio app work but then I get Error.

Last part of the log is:

Traceback (most recent call last):
File "/usr/local/lib/python3.10/site-packages/gradio/queueing.py", line 536, in process_events
response = await route_utils.call_process_api(
File "/usr/local/lib/python3.10/site-packages/gradio/route_utils.py", line 322, in call_process_api
output = await app.get_blocks().process_api(
File "/usr/local/lib/python3.10/site-packages/gradio/blocks.py", line 1935, in process_api
result = await self.call_function(
File "/usr/local/lib/python3.10/site-packages/gradio/blocks.py", line 1520, in call_function
prediction = await anyio.to_thread.run_sync( # type: ignore
File "/usr/local/lib/python3.10/site-packages/anyio/to_thread.py", line 56, in run_sync
return await get_async_backend().run_sync_in_worker_thread(
File "/usr/local/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 2505, in run_sync_in_worker_thread
return await future
File "/usr/local/lib/python3.10/site-packages/anyio/_backends/_asyncio.py", line 1005, in run
result = context.run(func, *args)
File "/usr/local/lib/python3.10/site-packages/gradio/utils.py", line 826, in wrapper
response = f(*args, **kwargs)
File "/content/TRELLIS/app.py", line 100, in image_to_3d
outputs = pipeline.run(
File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/content/TRELLIS/trellis/pipelines/trellis_image_to_3d.py", line 281, in run
coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
File "/content/TRELLIS/trellis/pipelines/trellis_image_to_3d.py", line 183, in sample_sparse_structure
z_s = self.sparse_structure_sampler.sample(
File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/content/TRELLIS/trellis/pipelines/samplers/flow_euler.py", line 199, in sample
return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/content/TRELLIS/trellis/pipelines/samplers/flow_euler.py", line 112, in sample
out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
return func(*args, **kwargs)
File "/content/TRELLIS/trellis/pipelines/samplers/flow_euler.py", line 73, in sample_once
pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
File "/content/TRELLIS/trellis/pipelines/samplers/flow_euler.py", line 43, in _get_model_prediction
pred_v = self._inference_model(model, x_t, t, cond, **kwargs)
File "/content/TRELLIS/trellis/pipelines/samplers/guidance_interval_mixin.py", line 11, in _inference_model
pred = super()._inference_model(model, x_t, t, cond, **kwargs)
File "/content/TRELLIS/trellis/pipelines/samplers/flow_euler.py", line 40, in _inference_model
return model(x_t, t, cond, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/content/TRELLIS/trellis/models/sparse_structure_flow.py", line 192, in forward
h = block(h, t_emb, cond)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/content/TRELLIS/trellis/modules/transformer/modulated.py", line 156, in forward
return self._forward(x, mod, context)
File "/content/TRELLIS/trellis/modules/transformer/modulated.py", line 139, in _forward
h = self.self_attn(h)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
return self._call_impl(*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
return forward_call(*args, **kwargs)
File "/content/TRELLIS/trellis/modules/attention/modules.py", line 126, in forward
h = scaled_dot_product_attention(q, k, v)
File "/content/TRELLIS/trellis/modules/attention/full_attn.py", line 120, in scaled_dot_product_attention
out = flash_attn.flash_attn_func(q, k, v)
File "/usr/local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 1163, in flash_attn_func
return FlashAttnFunc.apply(
File "/usr/local/lib/python3.10/site-packages/torch/autograd/function.py", line 574, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
File "/usr/local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 810, in forward
out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward(
File "/usr/local/lib/python3.10/site-packages/torch/ops.py", line 1061, in call
return self._op(*args, **(kwargs or {}))
File "/usr/local/lib/python3.10/site-packages/torch/_library/autograd.py", line 98, in autograd_impl
result = Generated.apply(*args, Metadata(keyset, keyword_only_args)) # type: ignore[attr-defined]
File "/usr/local/lib/python3.10/site-packages/torch/autograd/function.py", line 574, in apply
return super().apply(*args, **kwargs) # type: ignore[misc]
File "/usr/local/lib/python3.10/site-packages/torch/_library/autograd.py", line 40, in forward
result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/ops.py", line 672, in redispatch
return self._handle.redispatch_boxed(keyset, *args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/torch/_library/custom_ops.py", line 236, in backend_impl
result = self._backend_fns[device_type](*args, **kwargs)
File "/usr/local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 91, in _flash_attn_forward
out, softmax_lse, S_dmask, rng_state = flash_attn_cuda.fwd(
RuntimeError: FlashAttention only supports Ampere GPUs or newer.

Is there a way, in setup.sh or in app.py, not to use FlashAttention (Colab is on T4) ?
Thanks

Reproduce the failure in Google Colab on a T4, then start at trellis/modules/attention/full_attn.py where scaled_dot_product_attention calls FlashAttention. Trace how app.py and setup.sh select the attention implementation, and verify that image-to-3D inference completes on a T4 without requiring Ampere-only FlashAttention.