update

app.py

@@ -11,6 +11,7 @@ import torch.nn as nn
 import math
 import logging
 import sys
+import devicetorch
 from qwen2_vl.modeling_qwen2_vl import Qwen2VLSimplifiedModel
 from huggingface_hub import snapshot_download
 
@@ -24,7 +25,7 @@ logger = logging.getLogger(__name__)
 
 MODEL_ID = "Djrango/Qwen2vl-Flux"
 MODEL_CACHE_DIR = "model_cache"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = devictorch.get(torch)
 DTYPE = torch.bfloat16
 
 # Aspect ratio options
@@ -171,7 +172,7 @@ def process_image(image):
         # Move models back to CPU
         qwen2vl.cpu()
         connector.cpu()
-        torch.cuda.empty_cache()
+        devicetorch.empty_cache(torch)
         
         return result
             
@@ -429,4 +430,4 @@ if __name__ == "__main__":
         server_name="0.0.0.0",  # Listen on all network interfaces
         server_port=7860,       # Use a specific port
         share=False,             # Disable public URL sharing
-    )
+    )

requirements.txt

@@ -1,6 +1,6 @@
 # Core dependencies
-torch==2.4.0
-torchvision==0.19.0
+#torch==2.4.0
+#torchvision==0.19.0
 transformers==4.45.0
 diffusers==0.30.0
 accelerate==0.33.0
@@ -12,4 +12,4 @@ numpy>=1.24.0
 protobuf==4.23.4
 sentencepiece==0.2.0
 gradio==5.6.0
-huggingface-hub
+huggingface-hub