Update app.py

app.py

@@ -1 +1,140 @@
-print("hello")
+import os
+import torch
+from PIL import Image
+import gradio as gr
+from transformers import DetrImageProcessor, DetrForObjectDetection
+from diffusers import StableDiffusionPipeline
+from huggingface_hub import login
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+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.")
+
+# Login to Hugging Face using the token
+login(token=HF_TOKEN)
+
+# Load DETR model for object detection
+def load_detr_model():
+    try:
+        model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50')
+        processor = DetrImageProcessor.from_pretrained('facebook/detr-resnet-50')
+        return model, processor, None
+    except Exception as e:
+        return None, None, f"Error loading DETR model: {str(e)}"
+
+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")
+            outputs = detr_model(**inputs)
+            target_sizes = torch.tensor([image.size[::-1]])
+            results = detr_processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]
+
+            detected_objects = [
+                {"label": detr_model.config.id2label[label.item()],
+                 "box": box.tolist()}
+                for label, box in zip(results['labels'], results['boxes'])
+            ]
+            return detected_objects, None
+        except Exception as e:
+            return None, f"Error in detect_objects: {str(e)}"
+    else:
+        return None, "DETR models not loaded. Skipping object detection."
+
+# Load Stable Diffusion model for image generation
+def load_stable_diffusion_model():
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
+        return pipeline, None
+    except Exception as e:
+        return None, f"Error loading Stable Diffusion model: {str(e)}"
+
+sd_pipeline, sd_error = load_stable_diffusion_model()
+
+def adjust_dimensions(width, height):
+    # Adjust width and height to be divisible by 8
+    adjusted_width = (width // 8) * 8
+    adjusted_height = (height // 8) * 8
+    return adjusted_width, adjusted_height
+
+def generate_image(prompt, width, height):
+    if sd_pipeline is not None:
+        try:
+            adjusted_width, adjusted_height = adjust_dimensions(width, height)
+            image = sd_pipeline(prompt, width=adjusted_width, height=adjusted_height).images[0]
+            # Resize back to original dimensions if needed
+            image = image.resize((width, height), Image.LANCZOS)
+            return image, None
+        except Exception as e:
+            return None, f"Error in generate_image: {str(e)}"
+    else:
+        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)
+        if detect_error:
+            return None, detect_error
+
+        # Create a prompt based on detected objects
+        prompt = "modern redesign of an interior room with "
+        if detected_objects:
+            prompt += ", ".join([obj['label'] for obj in detected_objects])
+        else:
+            prompt += "empty room"
+
+        # Generate a redesigned image based on the prompt
+        width, height = image.size
+        generated_image, gen_image_error = generate_image(prompt, width, height)
+        if gen_image_error:
+            return None, gen_image_error
+
+        return generated_image, None
+    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")],
+    title="Interior Redesign",
+    css=custom_css
+)
+
+try:
+    iface.launch()
+except Exception as e:
+    print(f"Error occurred while launching the interface: {str(e)}")