Update app.py

app.py

@@ -1,81 +1,55 @@
 import os
 import cv2
 import tempfile
-import numpy as np
-import gradio as gr
+from modelscope.outputs import OutputKeys
 from modelscope.pipelines import pipeline
 from modelscope.utils.constant import Tasks
-from pathlib import Path
+import numpy as np
+import gradio as gr
 
-# Initialize model
-def load_model():
-    global img_colorization
-    img_colorization = pipeline(
-        Tasks.image_colorization,
-        model='iic/cv_ddcolor_image-colorization',
-        model_revision='v1.0.0'
-    )
+# Load model once at startup
+img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
 
 def inference(img):
-    if img is None:
-        raise gr.Error("Please upload an image first")
+    """Process input image and return colorized output"""
+    image = cv2.imread(str(img))
+    output = img_colorization(image[..., ::-1])
+    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
     
-    with tempfile.TemporaryDirectory() as temp_dir:
-        # Convert PIL image to numpy array if needed
-        if isinstance(img, np.ndarray):
-            image = img
-        else:
-            image = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
-        
-        # Process image
-        output = img_colorization(image[..., ::-1])
-        result = output['output_img'].astype(np.uint8)
-        
-        # Save result
-        out_path = os.path.join(temp_dir, 'colorized.png')
-        cv2.imwrite(out_path, result)
-        return Path(out_path), "✅ Colorization completed successfully!"
+    # Save result to temporary directory
+    temp_dir = tempfile.mkdtemp()
+    out_path = os.path.join(temp_dir, 'colorized.png')
+    cv2.imwrite(out_path, result)
+    return Path(out_path)
 
-# Create modern UI with Blocks
-with gr.Blocks(theme="soft", title="🎨 AI Color Restoration Studio") as demo:
-    gr.Markdown("""
-    # 🎨 AI Color Restoration Studio
-    Transform your black & white photos into vibrant colorized versions using state-of-the-art AI!
-    
-    Upload an image and watch as our deep learning model automatically adds natural colors.
-    """)
+# Modern UI design
+with gr.Blocks(theme="default", css=".container {max-width: 1000px; margin: auto;}") as demo:
+    # Header section
+    gr.Markdown("## 🎨 Image Colorization Studio\n*Transform your black-and-white images into vibrant color masterpieces*")
     
+    # Input/Output layout
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column():
             input_img = gr.Image(
-                label="Upload Monochrome Image",
-                type="pil",
-                height=400,
-                sources=["upload"],
-                interactive=True
-            )
-            submit_btn = gr.Button("✨ Colorize Image", variant="primary")
-            clear_btn = gr.ClearButton()
-            
-        with gr.Column(scale=1):
-            output_img = gr.Image(
-                label="Colorized Result",
-                type="pil",
-                height=400,
-                interactive=False
+                label="Grayscale Image",
+                type="filepath",
+                elem_id="input-image"
             )
+            submit_btn = gr.Button("🎨 Colorize Image", variant="primary")
+        
+        with gr.Column():
+            output_img = gr.Image(label="Colorized Result", elem_id="output-image")
             download_btn = gr.File(label="Download Result")
-            status = gr.Textbox(label="Status", interactive=False)
     
     # Examples section
     gr.Examples(
         examples=[
-            ["examples/1.jpg"],
-            ["examples/2.jpg"],
-            ["examples/3.jpg"]
+            ["examples/vintage.jpg"],
+            ["examples/portrait.png"],
+            ["examples/architecture.jpeg"]
         ],
-        inputs=[input_img],
-        outputs=[output_img, status],
+        inputs=input_img,
+        outputs=output_img,
         fn=inference,
         cache_examples=True
     )
@@ -84,11 +58,13 @@ with gr.Blocks(theme="soft", title="🎨 AI Color Restoration Studio") as demo:
     submit_btn.click(
         fn=inference,
         inputs=[input_img],
-        outputs=[output_img, status]
+        outputs=[output_img]
+    )
+    output_img.change(
+        fn=lambda img: gr.File.update(value=img) if img else None,
+        inputs=[output_img],
+        outputs=[download_btn]
     )
-    
-    clear_btn.add([input_img, output_img, status])
 
-if __name__ == "__main__":
-    load_model()
-    demo.launch(debug=True)
+# Launch configuration
+demo.launch(enable_queue=True)