Update app.py

app.py

@@ -1,15 +1,23 @@
 import gradio as gr
-from diffusers import StableDiffusionUpscalePipeline
-import torch
 from PIL import Image
+import torch
+from diffusers import StableDiffusionUpscalePipeline
 
-# Load model
+# Load UltraSharpV2 model
 model_id = "Kim2091/UltraSharpV2"
-pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
 
-def enhance_image(input_image):
-    # input_image is a PIL.Image
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = StableDiffusionUpscalePipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float16 if device=="cuda" else torch.float32
+)
+pipe = pipe.to(device)
+
+def enhance_image(input_image: Image.Image) -> Image.Image:
+    """
+    Enhance and upscale input image using UltraSharpV2.
+    """
     result = pipe(input_image).images[0]
     return result