Update app.py

app.py

@@ -1,12 +1,10 @@
 import os
 import torch
-from PIL import Image, ImageDraw
-import numpy as np
+from PIL import Image
 import gradio as gr
 from transformers import DetrImageProcessor, DetrForObjectDetection
 from diffusers import StableDiffusionPipeline
 from huggingface_hub import login
-import torchvision.transforms as T
 from dotenv import load_dotenv
 
 # Load environment variables from .env file
@@ -14,6 +12,7 @@ load_dotenv()
 
 # Retrieve Hugging Face token from environment variable
 HF_TOKEN = os.getenv('HF_TOKEN')
+
 if HF_TOKEN is None:
     raise ValueError("Hugging Face token not found in environment variables.")
 
@@ -32,6 +31,9 @@ def load_detr_model():
 detr_model, detr_processor, detr_error = load_detr_model()
 
 def detect_objects(image):
+    if image is None:
+        return None, "Invalid image: image is None."
+
     if detr_model is not None and detr_processor is not None:
         try:
             inputs = detr_processor(images=image, return_tensors="pt")
@@ -81,6 +83,9 @@ def generate_image(prompt, width, height):
         return None, "Stable Diffusion model not loaded. Skipping image generation."
 
 def process_image(image):
+    if image is None:
+        return None, "Invalid image: image is None."
+
     try:
         # Detect objects in the provided image
         detected_objects, detect_error = detect_objects(image)
@@ -104,15 +109,29 @@ def process_image(image):
     except Exception as e:
         return None, f"Error in process_image: {str(e)}"
 
+# Custom CSS for styling
+custom_css = """
+body {
+    background-color: black;
+}
+
+h1 {
+    background: linear-gradient(to right, blue, purple);
+    -webkit-background-clip: text;
+    color: transparent;
+    font-size: 3em;
+    text-align: center;
+    margin-bottom: 20px;
+}
+"""
+
+# Creating the Gradio interface with custom styling
 iface = gr.Interface(
     fn=process_image,
-    inputs=[
-        gr.Image(type="pil", label="Upload Room Image")
-    ],
-    outputs=[
-        gr.Image(type="pil", label="Redesigned Image"),
-        gr.Textbox(label="Error Message")
-    ]
+    inputs=[gr.Image(type="pil", label="Upload Room Image")],
+    outputs=[gr.Image(type="pil", label="Redesigned Image"), gr.Textbox(label="Error Message")],
+    title="Interior Redesign",
+    css=custom_css
 )
 
 try: