kavinel commited on
Commit
134ba43
1 Parent(s): 0bffdb0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -10,9 +10,11 @@ API_URL = "https://api-inference.huggingface.co/models/ZB-Tech/Text-to-Image"
10
  headers = {"Authorization": "Bearer HF_TOKEN"}
11
 
12
  model_id = "instruction-tuning-sd/cartoonizer"
 
 
13
  pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
14
  model_id, torch_dtype=torch.float16, use_auth_token=True
15
- ).to("cuda")
16
 
17
  def query(payload):
18
  response = requests.post(API_URL, headers=headers, json=payload)
 
10
  headers = {"Authorization": "Bearer HF_TOKEN"}
11
 
12
  model_id = "instruction-tuning-sd/cartoonizer"
13
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
14
+
15
  pipeline = StableDiffusionInstructPix2PixPipeline.from_pretrained(
16
  model_id, torch_dtype=torch.float16, use_auth_token=True
17
+ ).to(device)
18
 
19
  def query(payload):
20
  response = requests.post(API_URL, headers=headers, json=payload)