Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -1,20 +1,7 @@
-# Copyright 2024 the Llamole team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import os 
 import yaml
 import json
+import spaces
 
 import torch
 import torch.nn as nn
@@ -234,6 +221,7 @@ class GraphDiT(nn.Module):
         return noisy_data
 
     @torch.no_grad()
+    @spaces.GPU(duration=60)
     def generate(
         self,
         properties,