Update app.py

app.py

@@ -1,21 +1,26 @@
 import gradio as gr
-from diffusion_single_file import load_model, enhance_image
+import torch
+from diffusers import StableDiffusionPipeline
 
-# Load the ZenCtrl model using the correct library
-model = load_model("fotographerai/zenctrl_tools")
+# Load the image enhancement model
+model_id = "fotographerai/zenctrl_tools"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def enhance_image_fn(input_image):
-    # Enhance the image using the model
-    enhanced_image = enhance_image(model, input_image)
-    return enhanced_image
+pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
+pipe.to(device)
+
+def enhance_image(image):
+    # Process the image using the AI model
+    result = pipe(image).images[0]  # Get the first enhanced image
+    return result
 
 # Set up the Gradio interface
 interface = gr.Interface(
-    fn=enhance_image_fn,
+    fn=enhance_image,
     inputs=gr.Image(type="pil", label="Upload Image to Enhance"),
     outputs=gr.Image(type="pil", label="Enhanced Image"),
-    title="Image Enhancement with ZenCtrl",
-    description="Upload an image to enhance it using FotographerAI's ZenCtrl tools. Enhance your photos with AI!",
+    title="AI Image Enhancement",
+    description="Enhance your images using FotographerAI's ZenCtrl tools with AI!",
     live=True
 )