Update app.py

app.py

@@ -1,8 +1,7 @@
 import os
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
 import numpy as np
-import traceback
 import gradio as gr
 from transformers import DetrImageProcessor, DetrForObjectDetection
 from diffusers import StableDiffusionPipeline
@@ -14,8 +13,7 @@ from dotenv import load_dotenv
 load_dotenv()
 
 # Retrieve Hugging Face token from environment variable
-HF_TOKEN = os.getenv("HF_TOKEN")
-
+HF_TOKEN = os.getenv('HF_TOKEN')
 if HF_TOKEN is None:
     raise ValueError("Hugging Face token not found in environment variables.")
 
@@ -40,7 +38,13 @@ def detect_objects(image):
             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]
-            return results, None
+
+            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:
@@ -57,27 +61,42 @@ def load_stable_diffusion_model():
 
 sd_pipeline, sd_error = load_stable_diffusion_model()
 
-def generate_image(prompt):
+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:
-            image = sd_pipeline(prompt).images[0]
+            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."
 
-# Gradio Interface
 def process_image(image):
     try:
-        # Detect objects
-        object_results, detect_error = detect_objects(image)
+        # Detect objects in the provided image
+        detected_objects, detect_error = detect_objects(image)
         if detect_error:
             return None, detect_error
 
-        # Generate a modern redesign of the image based on the detected objects
-        prompt = "modern redesign of an interior room"
-        generated_image, gen_image_error = generate_image(prompt)
+        # 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
 
@@ -100,4 +119,3 @@ try:
     iface.launch()
 except Exception as e:
     print(f"Error occurred while launching the interface: {str(e)}")
-    traceback.print_exc()