Update app.py

app.py

@@ -10,9 +10,11 @@ API_URL = "https://api-inference.huggingface.co/models/ZB-Tech/Text-to-Image"
 headers = {"Authorization": "Bearer HF_TOKEN"}
 
 model_id = "instruction-tuning-sd/cartoonizer"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
     model_id, torch_dtype=torch.float16, use_auth_token=True
-  ).to("cuda")
+  ).to(device)
 
 def query(payload):
     response = requests.post(API_URL, headers=headers, json=payload)