added multi gpu config

app.py

@@ -23,6 +23,36 @@ from PIL import Image
 from queue import Queue
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
+from dataclasses import dataclass
+from typing import Optional, List
+
+@dataclass
+class MultiGPUConfig:
+    count: int = 2  # Number of GPUs to request
+    memory: int = 16  # Memory per GPU in GB
+    duration: int = 3600  # Duration in seconds
+
+class SpacesMultiGPU:
+    def __init__(self, config: Optional[MultiGPUConfig] = None):
+        self.config = config or MultiGPUConfig()
+    
+    def __call__(self, func):
+        # Apply multiple GPU decorators
+        decorated_func = func
+        for gpu_idx in range(self.config.count):
+            decorated_func = spaces.GPU(
+                device=gpu_idx,  # Specify which GPU to request
+                memory=self.config.memory,
+                duration=self.config.duration
+            )(decorated_func)
+        return decorated_func
+
+# Example usage in your generation code
+gpu_config = MultiGPUConfig(
+    count=2,  # Request 2 GPUs
+    duration=400  # 1 hour duration
+)
+
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -142,7 +172,7 @@ def generate_image_on_gpu(args):
         release_gpu(gpu_id)
         raise e
         
-@spaces.GPU(duration=400)
+@SpacesMultiGPU(gpu_config)
 def generate_all(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress()):
     outputs = [None] * (len(MODEL_CONFIGS) * 2)