Update app.py

app.py

@@ -1,9 +1,6 @@
-# Yolo (before mod-spaces)
-from ultralytics import YOLO
-
 # UI and Application Framework
 import gradio as gr
-import spaces # Thanks a lot!
+import spaces
 
 
 # Standard Libraries
@@ -19,18 +16,23 @@ from PIL import Image
 import torch
 from transformers import pipeline
 from diffusers import AutoPipelineForInpainting
+from ultralytics import YOLO
 
 # Text and Data Manipulation
 import difflib
 
 
+# Constants
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
+# Load
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 yoloModel = YOLO('yolov8x-seg.pt')
 sdxl = AutoPipelineForInpainting.from_pretrained(
     "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
-    torch_dtype=torch.float16
-)
-image_captioner = pipeline("image-to-text", model="Abdou/vit-swin-base-224-gpt2-image-captioning")
+    torch_dtype=torch.float32
+).to(DEVICE)
+image_captioner = pipeline("image-to-text", model="Abdou/vit-swin-base-224-gpt2-image-captioning", device=DEVICE)
 
 
 def image_to_base64(image: Image.Image):
@@ -54,7 +56,7 @@ def get_most_similar_string(target_string, string_array):
 
 # Yolo
 def getClasses(model, img1):
-    results = model([np.array(img1)])
+    results = model([np.array(img1)])  # Изменение для передачи изображения как массива NumPy
     out = []
     for r in results:
         im_array = r.plot()
@@ -106,7 +108,6 @@ def getSegments(yoloModel, img1):
     return allMask
 
 
-
 # Gradio UI
 @spaces.GPU
 def captionMaker(base64_img):
@@ -165,4 +166,4 @@ iface = gr.Interface(
     live=False
 )
 
-iface.launch()
+iface.launch()