main

app.py

@@ -1,3 +1,4 @@
+import spaces
 import sys
 import os
 import torch
@@ -35,9 +36,10 @@ model_path = model_id
 pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float16)
 pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 #pipe.enable_model_cpu_offload()
+pipe.to('cuda')
 
+@spaces.GPU
 def core(bundle):
-
     generator = torch.Generator().manual_seed(int(bundle["seed"]))
     result = pipe(
         bundle=bundle,