Update app.py

app.py

@@ -6,11 +6,12 @@ import os
 from PIL import Image
 import base64
 import time
-from accelerate import Accelerator
 import logging
 
 # Disable GPU detection
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
+os.environ["CUDA_DEVICE_ORDER"] = ""
+os.environ["TORCH_CUDA_ARCH_LIST"] = ""
 torch.set_default_device("cpu")
 
 app = Flask(__name__, static_folder='static')
@@ -20,8 +21,8 @@ CORS(app)
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Initialize Accelerator for CPU
-accelerator = Accelerator(device_placement=False)
+# Log device in use
+logger.info(f"Device in use: {torch.device('cpu')}")
 
 # Model cache
 model_cache = {}
@@ -49,10 +50,8 @@ def load_model(model_id):
                 low_cpu_mem_usage=True
             )
             pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-            pipe = accelerator.prepare(pipe)
             pipe.enable_attention_slicing()
-            pipe.enable_sequential_cpu_offload()
-            pipe.to("cpu")
+            pipe.to(torch.device("cpu"))
             model_cache[model_id] = pipe
             logger.info(f"Model {model_id} loaded successfully")
         except Exception as e:
@@ -90,7 +89,7 @@ def generate():
 
         width, height = ratio_to_dims.get(ratio, (256, 256))
         pipe = load_model(model_id)
-        pipe.to("cpu")
+        pipe.to(torch.device("cpu"))
 
         images = []
         num_inference_steps = 20 if model_id == 'ssd-1b' else 30