Update app.py

app.py

@@ -59,11 +59,13 @@ class BuildTheFuture:
         """Load YOLOv11 model for structural detection"""
         try:
             from ultralytics import YOLO
-            # You can replace this with your custom trained model
-            self.yolo_model = YOLO('yolov11n.pt')
+            # Prefer local trained model if available, fallback to small public model
+            model_path = 'best.pt' if os.path.exists('best.pt') else 'yolov11n.pt'
+            self.yolo_model = YOLO(model_path)
+            logger.info(f"YOLO model loaded: {model_path}")
             return True
         except Exception as e:
-            print(f"Error loading YOLO model: {e}")
+            logger.error(f"Error loading YOLO model: {e}")
             return False
     
     def detect_structures(self, image):
@@ -90,7 +92,7 @@ class BuildTheFuture:
             return image, f"Detection error: {str(e)}"
     
     def complete_construction(self, image, style="realistic"):
-        """Complete the construction using Gemini 2.5 Flash Image"""
+        """Complete the construction using Gemini image generation"""
         if self.model is None:
             return image, "Gemini model not available. Please check your API key."
         
@@ -107,18 +109,29 @@ class BuildTheFuture:
             # Add additional context for better results
             enhanced_prompt = f"{prompt} The image should maintain the same perspective and lighting as the original. Focus on completing the unfinished parts while maintaining architectural coherence."
             
-            # Generate the completed image
-            response = self.model.generate_content([enhanced_prompt, image])
-            
+            # Prepare image bytes for Gemini
+            buffered = io.BytesIO()
+            image.save(buffered, format='PNG')
+            image_bytes = buffered.getvalue()
+
+            # Use a stable Gemini image-capable model
+            gemini_image_model = genai.GenerativeModel('gemini-1.5-flash')
+            response = gemini_image_model.generate_content([
+                enhanced_prompt,
+                {
+                    'mime_type': 'image/png',
+                    'data': base64.b64encode(image_bytes).decode('utf-8')
+                }
+            ])
+
             # Extract the generated image
-            if response.candidates and response.candidates[0].content.parts:
+            if getattr(response, 'candidates', None):
                 for part in response.candidates[0].content.parts:
-                    if hasattr(part, 'inline_data') and part.inline_data:
-                        # Convert base64 to image
+                    if hasattr(part, 'inline_data') and part.inline_data and getattr(part.inline_data, 'data', None):
                         image_data = base64.b64decode(part.inline_data.data)
-                        completed_image = Image.open(io.BytesIO(image_data))
+                        completed_image = Image.open(io.BytesIO(image_data)).convert('RGB')
                         return completed_image, f"Construction completed successfully with {style} style"
-            
+
             return image, "No image generated - please try again"
             
         except Exception as e: