Update app.py

app.py

@@ -1,46 +1,72 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
+from torchvision import transforms
 from PIL import Image
+from model import Generator  # Assuming you are using Hammad712's model structure
 
-# Load the model
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "rsortino/ColorizeNet",  # YOUR MODEL
-    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,
-        strength=0.8,
-        guidance_scale=7.5,
-        num_inference_steps=30
-    )
-    return result.images[0]
+# Load model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = Generator().to(device)
+model.load_state_dict(torch.load('generator.pth', map_location=device))
+model.eval()
+
+# Define preprocessing and postprocessing
+preprocess = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.Grayscale(num_output_channels=1),
+    transforms.ToTensor()
+])
+
+postprocess = transforms.ToPILImage()
+
+def colorize_image(input_image):
+    input_tensor = preprocess(input_image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        output_tensor = model(input_tensor)
+    output_image = postprocess(output_tensor.squeeze(0).cpu().clamp(0, 1))
+    return output_image
+
+def reset():
+    return None, None
 
 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.")
+    gr.Markdown("# 🎨 Image Colorization App")
 
     with gr.Row():
         with gr.Column():
-            input_img = gr.Image(label="Grayscale Input", type="pil")
-            submit_btn = gr.Button("Colorize")
-
+            input_image = gr.Image(label="Upload your grayscale image", type="pil")
+            clear_button = gr.Button("🔄 Reset / Clear")
+            download_button = gr.File(label="Download Colorized Image")
         with gr.Column():
-            output_img = gr.Image(label="Colorized Output", type="pil")
-            download_btn = gr.Button("Download")
+            output_image = gr.Image(label="Colorized Image")
+
+    colorize_btn = gr.Button("✨ Colorize Image")
 
-    submit_btn.click(fn=colorize, inputs=input_img, outputs=output_img)
-    download_btn.click(fn=lambda img: img, inputs=output_img, outputs=gr.File())
+    colorize_btn.click(
+        colorize_image,
+        inputs=input_image,
+        outputs=output_image
+    )
+
+    clear_button.click(
+        reset,
+        inputs=[],
+        outputs=[input_image, output_image]
+    )
+
+    # Allow download after processing
+    def prepare_download(image):
+        if image:
+            path = "colorized_output.png"
+            image.save(path)
+            return path
+        else:
+            return None
+
+    output_image.change(
+        prepare_download,
+        inputs=output_image,
+        outputs=download_button
+    )
 
 demo.launch()