Update app.py

app.py

@@ -2,8 +2,15 @@ from transformers import Blip2Processor, Blip2ForConditionalGeneration, AutoMode
 import torch
 from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
+import torch
 import os
 
+# Check CUDA availability
+if not torch.cuda.is_available():
+    os.environ["BITSANDBYTES_NOWELCOME"] = "1"
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    os.environ["LIBRARY_PATH"] = "/usr/local/cuda/lib64/stubs:$LIBRARY_PATH"
+
 # Initialize environment
 os.makedirs("./offload", exist_ok=True)
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -28,7 +35,7 @@ phi_model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
     device_map="auto",
     torch_dtype=torch.float16,
-    load_in_4bit=True,
+    load_in_4bit=torch.cuda.is_available(),  # Only use 4bit if CUDA available
     token=HF_TOKEN
 ).eval()
 phi_tokenizer = AutoTokenizer.from_pretrained(
@@ -123,4 +130,4 @@ with gr.Blocks(title="AI Meme Generator") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)