Spaces:
Sleeping
Sleeping
| from tools.generate import generate_image | |
| from tools.edit import edit_image | |
| from tools.analyze import analyze_image | |
| from tools.upload import save_uploads | |
| class ImageAgent: | |
| def __init__(self, router, memory, planner): | |
| self.router = router | |
| self.memory = memory | |
| self.planner = planner | |
| # -------------------------------------------------- | |
| # Upload Files | |
| # -------------------------------------------------- | |
| def upload(self, files): | |
| uploaded = save_uploads(files) | |
| self.memory.uploaded_files = uploaded | |
| self.memory.image_context = None | |
| self.memory.analysis = None | |
| for file in uploaded: | |
| if file.get("type", "").startswith("image/"): | |
| self.memory.last_used_image = file["path"] | |
| break | |
| return uploaded | |
| # -------------------------------------------------- | |
| # Analyze Image | |
| # -------------------------------------------------- | |
| def auto_analyze(self, image_path=None): | |
| image = image_path or self.memory.last_used_image | |
| if not image: | |
| return { | |
| "description": "No uploaded image found." | |
| } | |
| analysis = analyze_image(image) | |
| self.memory.analysis = analysis | |
| self.memory.image_context = analysis | |
| return analysis | |
| # -------------------------------------------------- | |
| # Convert object list safely | |
| # -------------------------------------------------- | |
| def format_objects(self, objects): | |
| if not objects: | |
| return "" | |
| names = [] | |
| for obj in objects: | |
| if isinstance(obj, str): | |
| names.append(obj) | |
| elif isinstance(obj, dict): | |
| if "name" in obj: | |
| names.append(str(obj["name"])) | |
| elif "label" in obj: | |
| names.append(str(obj["label"])) | |
| else: | |
| names.append(str(obj)) | |
| else: | |
| names.append(str(obj)) | |
| return ", ".join(names) | |
| # -------------------------------------------------- | |
| # Build Prompt | |
| # -------------------------------------------------- | |
| def build_context(self, analysis): | |
| if not analysis: | |
| return "" | |
| objects = self.format_objects( | |
| analysis.get("objects", []) | |
| ) | |
| colors = analysis.get("colors", []) | |
| if isinstance(colors, list): | |
| colors = ", ".join(colors) | |
| suggestions = analysis.get("suggestions", []) | |
| if isinstance(suggestions, list): | |
| suggestions = ", ".join(suggestions) | |
| context = f""" | |
| Image Description: | |
| {analysis.get("description","")} | |
| Objects: | |
| {objects} | |
| People: | |
| {analysis.get("people","")} | |
| Style: | |
| {analysis.get("style","")} | |
| Lighting: | |
| {analysis.get("lighting","")} | |
| Colors: | |
| {colors} | |
| Background: | |
| {analysis.get("background","")} | |
| Camera Angle: | |
| {analysis.get("camera_angle","")} | |
| Quality: | |
| {analysis.get("quality","")} | |
| Editing Prompt: | |
| {analysis.get("editing_prompt","")} | |
| Suggestions: | |
| {suggestions} | |
| """ | |
| return context | |
| # -------------------------------------------------- | |
| # Main Agent | |
| # -------------------------------------------------- | |
| def run(self, user_input): | |
| task = self.router.route(user_input) | |
| steps = self.planner.plan(task) | |
| image = self.memory.last_used_image | |
| if image and self.memory.image_context is None: | |
| self.auto_analyze(image) | |
| analysis = self.memory.image_context | |
| context = self.build_context(analysis) | |
| prompt = f""" | |
| {context} | |
| User Request: | |
| {user_input} | |
| """ | |
| result = None | |
| for step in steps: | |
| if step == "generate": | |
| result = generate_image(prompt) | |
| elif step == "edit": | |
| result = edit_image( | |
| image_path=image, | |
| prompt=prompt | |
| ) | |
| # If edit_image returns a file path | |
| if isinstance(result, str): | |
| self.memory.last_used_image = result | |
| # If edit_image returns a PIL.Image, | |
| # app.py will display it directly. | |
| elif step == "analyze": | |
| result = self.auto_analyze(image) | |
| elif step == "chat": | |
| result = analysis | |
| return result |