Update app.py

app.py

@@ -4,13 +4,13 @@ from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, U
 from datasets import load_dataset
 from torch.utils.data import DataLoader
 from torchvision import transforms
-
+import spaces
 css = "* {font-family: sans-serif;} .gradio-container {max-width:950px!important; padding:10px!important;} .box {background:#161b22!important; border:1px solid #30363d!important; border-radius:8px!important; padding:12px!important;} .btn {background:linear-gradient(135deg,#2563eb,#1d4ed8)!important; color:white!important; font-weight:600!important;} input, select, textarea {background:#0d1117!important; color:#fff!important; border:1px solid #30363d!important;}"
 pipe, logs = None, []
 
 def log(msg): global logs; logs.append(f"[{time.strftime('%H:%M:%S')}] {msg}")
 def get_logs(): return "\n".join(logs)
-
+@spaces.GPU
 def train_from_scratch(token, repo, epochs, lr, batch_size):
     global pipe, logs
     logs = []
@@ -65,7 +65,7 @@ def train_from_scratch(token, repo, epochs, lr, batch_size):
 def start_train(t, r, e, l, b):
     threading.Thread(target=train_from_scratch, args=(t, r, e, l, b)).start()
     return "Initialisation de l'architecture vierge..."
-
+@spaces.GPU 
 def generate(prompt, neg, steps, cfg, w, h, img=None, strg=0.75):
     global pipe
     if pipe is None: return None, "Veuillez d'abord entraîner le modèle ou initialiser l'état."