Update main.py
Browse files
main.py
CHANGED
@@ -12,7 +12,7 @@ app = FastAPI()
|
|
12 |
|
13 |
|
14 |
|
15 |
-
model =
|
16 |
|
17 |
def generate_cartoon(image_bytes: bytes) -> bytes:
|
18 |
image = Image.open(BytesIO(image_bytes))
|
|
|
12 |
|
13 |
|
14 |
|
15 |
+
model = cartoon1(device='cuda' if torch.cuda.is_available() else 'cpu')
|
16 |
|
17 |
def generate_cartoon(image_bytes: bytes) -> bytes:
|
18 |
image = Image.open(BytesIO(image_bytes))
|