Update app.py

app.py

@@ -1,30 +1,14 @@
 import gradio as gr
+from diffusers import StableDiffusionPipeline
 import torch
-from torchvision import transforms
-from PIL import Image
-from model import Generator  # Assuming you are using Hammad712's model structure
+from transformers import pipeline
 
-# 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()
+# Use an available colorization pipeline from Hugging Face
+colorizer = pipeline("image-to-image", model="HuggingFaceM4/DeOldify", device_map="auto")
 
 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
+    output = colorizer(input_image)
+    return output[0]['image']
 
 def reset():
     return None, None
@@ -54,7 +38,6 @@ with gr.Blocks() as demo:
         outputs=[input_image, output_image]
     )
 
-    # Allow download after processing
     def prepare_download(image):
         if image:
             path = "colorized_output.png"