Create app.py

app.py

@@ -0,0 +1,61 @@
+import gradio as gr
+import torch
+from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel
+from PIL import Image
+
+# Load ControlNet model
+controlnet = ControlNetModel.from_pretrained(
+    "rsortino/ColorizeNet",
+    torch_dtype=torch.float16,
+    use_safetensors=True
+)
+
+# Load the pipeline
+pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1",
+    controlnet=controlnet,
+    torch_dtype=torch.float16
+)
+
+# Move to CUDA if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = pipe.to(device)
+
+# Disable safety checker
+pipe.safety_checker = lambda images, **kwargs: (images, False)
+
+def colorize(image: Image.Image) -> Image.Image:
+    image = image.convert("RGB").resize((512, 512))
+    result = pipe(
+        prompt="A realistic colorized version of this image.",
+        image=image,
+        control_image=image,
+        strength=1.0,
+        guidance_scale=9.0,
+        num_inference_steps=30
+    )
+    return result.images[0]
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🎨 ColorizeNet - Grayscale to Color Image")
+    gr.Markdown("Upload a grayscale image. The model will generate a realistic colorized version.")
+
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(label="Grayscale Input", type="pil")
+            submit_btn = gr.Button("Colorize")
+
+        with gr.Column():
+            output_img = gr.Image(label="Colorized Output", type="pil")
+            download_btn = gr.Button("Download")
+
+    def handle_colorize(img):
+        return colorize(img)
+
+    def download_image(img):
+        return img
+
+    submit_btn.click(fn=handle_colorize, inputs=input_img, outputs=output_img)
+    download_btn.click(fn=download_image, inputs=output_img, outputs=gr.File())
+
+demo.launch()