Update app.py

app.py

@@ -136,8 +136,12 @@ class ModelLoader:
                 "torch_dtype": torch.float16 if self.device == "cuda" else torch.float32,
                 "device_map": "auto" if self.device == "cuda" else None,
                 "low_cpu_mem_usage": True,
-                "offload_folder": "offload",
-                "max_memory": {i: "20GiB" for i in range(torch.cuda.device_count())} if torch.cuda.device_count() > 1 else {}
+                "offload_folder": "offload"
+            }
+
+            # Add max_memory configuration if multiple GPUs available
+            if torch.cuda.device_count() > 1:
+                model_kwargs["max_memory"] = {i: "20GiB" for i in range(torch.cuda.device_count())}
             
             model = None
             for attempt in range(3):