diff --git a/bashful.py b/bashful.py index f74d311..1422afd 100644 --- a/bashful.py +++ b/bashful.py @@ -1,17 +1,29 @@ -import json -import functions_framework +import os import uuid -from models.http_models.bashful_http import Txt2ImgModel + +import requests +from models.http_models.bashful_http import Img2SVGModel, RemoveBGRequest, Txt2ImgModel, Txt2SVGModel from utils.http_utils import cors_decorator -from providers import airtable_service, ai_service, logger, response_headers, gcp_service +from providers import ( + airtable_service, + ai_service, + logger, + response_headers, + gcp_service, +) from utils.image_utils import convert_b64_to_bytes, replace_transparent_with_white from PIL import Image import io -from utils.img2img_utils import create_default_img2img_payload, merge_model_config, merge_request_and_default_img2img_payload +from utils.img2img_utils import ( + create_default_img2img_payload, + merge_model_config, + merge_request_and_default_img2img_payload, +) + ##################### PROJECT: Bashful-Photoshop ##################### @@ -53,7 +65,7 @@ def txt2img(request: Txt2ImgModel): payload["upscale"] = model_config.get("upscale", "") payload["steps"] = model_config.get("steps", "") seed = payload.get("seed", None) - + if seed: del payload["seed"] @@ -122,7 +134,7 @@ def img2img(request): if not model_config: return ({"error": "No model config found"}, 400, response_headers) - + payload = merge_model_config(payload, model_config) response = ai_service.img2img(byte_image, payload) @@ -157,3 +169,70 @@ def img2img(request): except Exception as e: logger.error(e) return ({"error": str(e)}, 500, response_headers) + + +def img2svg(request: Img2SVGModel): + image_url = request.image_url + if not image_url: + return ({"error": "Invalid image url"}, 400, response_headers) + + response = requests.post( + "https://vectorizer.ai/api/v1/vectorize", + data={ + "image.url": image_url, + # TODO: Add more upload options here + }, + headers={ + "Authorization": "Basic dmtjM3prbW44djM2ZmVkOjVzMXMzN2l2czRoZHNyazJjczVtZ2Uxa2ZwcjdtZmNuZzllbXM0czNpODdwZGZvczNpMHY=" + }, + ) + if response.status_code == requests.codes.ok: + return (response.content, 200, response_headers) + else: + return ({"error": "Invalid image url"}, 400, response_headers) + + +def txt2svg(request: Txt2SVGModel | Txt2ImgModel): + try: + response = txt2img(request) + if response[1] != 200: + return response + + img2svg_req = Img2SVGModel(image_url=response[0].get("url")) + img2svg_rep = img2svg(img2svg_req) + if img2svg_rep[1] != 200: + return img2svg_rep + + upload_url = gcp_service.upload_svg_to_storage( + img2svg_rep[0], "bashful_api", f"svg/{uuid.uuid1()}.svg" + ) + + + return ( + {"svg_url": upload_url, "image_url": response[0].get("url")}, + 200, + response_headers, + ) + except Exception as e: + logger.error(e) + return ({"error": str(e)}, 500, response_headers) + + +def remove_background(request: RemoveBGRequest): + data = request.json + img_response = requests.get(data.get("image_url")) + + image_data = img_response.content + + response = requests.post( + 'https://sdk.photoroom.com/v1/segment', + headers={'x-api-key': 'eebb03f0afc9a8df9781f0c63c5c4d9edfbe9906'}, + files={'image_file': image_data}, + ) + + upload_url = gcp_service.upload_image_to_storage( + response.content, "bashful_api", f"remove_bg/{uuid.uuid1()}.png" + ) + + + return ({"image_url": upload_url}, 200, response_headers) \ No newline at end of file diff --git a/bashful_studio.py b/bashful_studio.py index a401dfe..8f01043 100644 --- a/bashful_studio.py +++ b/bashful_studio.py @@ -439,26 +439,4 @@ def embed_watermark(request): except Exception as e: logger.error(e) return ({"error": str(e)}, 500, response_headers) - -def img2svg(request): - - request_data = request.json - image_url = request_data.get("image_url") - if not image_url: - return ({"error": "Invalid image url"}, 400 , response_headers) - - response = requests.post( - 'https://vectorizer.ai/api/v1/vectorize', - data={ - 'image.url': image_url, - # TODO: Add more upload options here - }, - headers={ - 'Authorization': - 'Basic dmtjM3prbW44djM2ZmVkOjVzMXMzN2l2czRoZHNyazJjczVtZ2Uxa2ZwcjdtZmNuZzllbXM0czNpODdwZGZvczNpMHY=' - }, - ) - if response.status_code == requests.codes.ok: - return (response.content, 200, response_headers) - else: - return ({"error": "Invalid image url"}, 400 , response_headers) \ No newline at end of file + \ No newline at end of file diff --git a/chatbot/chat.py b/chatbot/chat.py new file mode 100644 index 0000000..43263ee --- /dev/null +++ b/chatbot/chat.py @@ -0,0 +1,11 @@ +def chat_bot(): + while True: + user_input = input("You: ") + if user_input.lower() == 'exit': + break + else: + print("Bot: " + user_input) + +if __name__ == "__main__": + print("Chat Bot is ready. Type 'exit' to end the conversation.") + chat_bot() diff --git a/llm_custom_agent.py b/chatbot/llm_custom_agent.py similarity index 100% rename from llm_custom_agent.py rename to chatbot/llm_custom_agent.py diff --git a/llm_showcase.py b/chatbot/llm_showcase.py similarity index 100% rename from llm_showcase.py rename to chatbot/llm_showcase.py diff --git a/chatbot/palm_based_chat.py b/chatbot/palm_based_chat.py new file mode 100644 index 0000000..1111183 --- /dev/null +++ b/chatbot/palm_based_chat.py @@ -0,0 +1,98 @@ +import os +import re +from typing import Union +from langchain.tools import BaseTool, StructuredTool, Tool, tool +from langchain.agents import AgentType, initialize_agent +from langchain.chat_models import ChatVertexAI, ChatOpenAI +from langchain import LLMMathChain +from langchain.prompts.chat import ( + ChatPromptTemplate, + SystemMessagePromptTemplate, + HumanMessagePromptTemplate, +) +from langchain.schema import ( + HumanMessage, + SystemMessage +) +from rembg import remove +from langchain.agents import Tool, AgentExecutor, LLMSingleActionAgent, AgentOutputParser +from langchain.schema import AgentAction, AgentFinish, HumanMessage + + +chat = ChatVertexAI(project="shepard-379102") +# chat = ChatOpenAI() + + + +def remove_background(message: str): + input_path = 'with_background_2.png' + output_path = 'without_background_2.png' + + with open(input_path, 'rb') as i: + with open(output_path, 'wb') as o: + input = i.read() + output = remove(input) + o.write(output) + +def get_image_location(message: str): + path = os.path.join(os.path.dirname(__file__), message) + return path + +def get_this_file_location(message: str): + # Return the location of this file + return os.path.dirname(__file__) + +tools = [ + Tool.from_function( + func=remove_background, + name = "Background Remover", + description="This removes the backgrounds from images." + # coroutine= ... <- you can specify an async method if desired as well + ), + Tool.from_function( + func=get_image_location, + name = "Image Location Getter", + description="This builds the location of the image on the server." + # coroutine= ... <- you can specify an async method if desired as well + ), + Tool.from_function( + func=get_this_file_location, + name = "Image Location Getter", + description="This builds the location of the image on the server." + # coroutine= ... <- you can specify an async method if desired as well + ), +] + +class CustomOutputParser(AgentOutputParser): + + def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]: + print("llm_output:\n\n", llm_output) + # Check if agent should finish + if "Final Answer:" in llm_output: + return AgentFinish( + # Return values is generally always a dictionary with a single `output` key + # It is not recommended to try anything else at the moment :) + return_values={"output": llm_output.split("Final Answer:")[-1].strip()}, + log=llm_output, + ) + # Parse out the action and action input + regex = r"Action: (.*?)[\n]*Action Input:[\s]*(.*)" + match = re.search(regex, llm_output, re.DOTALL) + if not match: + raise ValueError(f"Could not parse LLM output: `{llm_output}`") + action = match.group(1).strip() + action_input = match.group(2) + # Return the action and action input + return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output) + + +messages = [ + SystemMessage(content="You are a helpful assistant that builds prompts for text to image generators iteratively by gathering requirements from chatting with users. You edit the prompts and the AI generates images. You edit the images based on the edits the users asks for."), + HumanMessage(content="Remove the background from this image at where this script is running, the name of the file is with_background_2.png."), +] +agent = initialize_agent(tools, chat, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True) +result = agent.run(messages) +print(result) +# result = chat(messages) +# print(result) + diff --git a/chatbot/remove_background.py b/chatbot/remove_background.py new file mode 100644 index 0000000..56195d6 --- /dev/null +++ b/chatbot/remove_background.py @@ -0,0 +1,10 @@ +from rembg import remove + +input_path = 'with_background_2.png' +output_path = 'without_background_2.png' + +with open(input_path, 'rb') as i: + with open(output_path, 'wb') as o: + input = i.read() + output = remove(input) + o.write(output) \ No newline at end of file diff --git a/chatbot/test.py b/chatbot/test.py new file mode 100644 index 0000000..9fcff06 --- /dev/null +++ b/chatbot/test.py @@ -0,0 +1,71 @@ +import re +from langchain.agents import Tool, AgentExecutor, BaseSingleActionAgent, BaseMultiActionAgent +from langchain.schema import AgentAction, AgentFinish +from typing import List, Tuple, Any, Union + +class FileReader: + def read_file(self, location: str) -> str: + with open(location, 'r') as f: + return f.read() + +class BackgroundRemover: + def remove_background(self, image_path: str) -> str: + # code to remove background from image + return "background removed from " + image_path + +class Joker: + def make_joke(self, image_path: str) -> str: + # code to remove background from image + return "background removed from " + image_path + +file_reader = FileReader() +background_remover = BackgroundRemover() +joker = Joker() + +tools = [ + Tool( + name="Read File", + func=file_reader.read_file, + description="reads a file given a location" + ), + Tool( + name="Remove Background", + func=background_remover.remove_background, + description="removes the background from an image" + ), + Tool( + name="Make Joke", + func=joker.make_joke, + description="removes the background from an image" + ) +] + +class CustomAgent(BaseMultiActionAgent): + @property + def input_keys(self): + return ["message"] + + def plan( + self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any + ) -> Union[AgentAction, AgentFinish]: + file_path = re.search(r'remove background from (.+)', kwargs["message"]).group(1) + return AgentAction(tool="Read File", tool_input=file_path, log="") + + def aplan( + self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any + ) -> Union[AgentAction, AgentFinish]: + file_path = re.search(r'remove background from (.+)', kwargs["message"]).group(1) + return AgentAction(tool="Read File", tool_input=file_path, log="") + +agent = CustomAgent() + +# agent = LLMSingleActionAgent( +# llm_chain=llm_chain, +# output_parser=output_parser, +# stop=["\nObservation:"], +# allowed_tools=tool_names +# ) +agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=True) + +result = agent_executor.run("remove background from yolo.txt") +print(result) \ No newline at end of file diff --git a/chatbot/with_background_2.png b/chatbot/with_background_2.png new file mode 100644 index 0000000..b4dd422 Binary files /dev/null and b/chatbot/with_background_2.png differ diff --git a/chatbot/yolo.txt b/chatbot/yolo.txt new file mode 100644 index 0000000..182a3e9 --- /dev/null +++ b/chatbot/yolo.txt @@ -0,0 +1 @@ +Sup \ No newline at end of file diff --git a/config.py b/config.py index d42ed9a..1080bd2 100644 --- a/config.py +++ b/config.py @@ -53,3 +53,9 @@ VECTORIZER_AI_URL = ( os.environ.get("VECTORIZER_AI_URL") or "https://vectorizer.ai/api/v1/vectorize" ) +TRAINING_DB_URL = ( + os.environ.get("TRAINING_DB_URL") or "http://localhost:1337/" +) +BASHFUL_STUDIO_API_URL = ( + os.environ.get("BASHFUL_STUDIO_API_URL") or "http://localhost:1337/" +) diff --git a/local_dev_server.py b/local_dev_server.py index a7f25f1..96d2c12 100644 --- a/local_dev_server.py +++ b/local_dev_server.py @@ -22,4 +22,4 @@ def route_func(*args, **kwargs): app.add_url_rule(f'/{name}', name, route_func, methods=['POST', 'GET', 'PUT', 'DELETE', 'PATCH', 'OPTIONS', 'HEAD']) if __name__ == "__main__": - app.run(host="0.0.0.0", port=8080) + app.run(host="0.0.0.0", debug=True, port=8080) diff --git a/main.py b/main.py index 4674386..86b8fca 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,6 @@ import functions_framework from pydantic import ValidationError -from models.http_models.bashful_http import Txt2ImgModel +from models.http_models.bashful_http import RemoveBGRequest, Txt2ImgModel, GetModelConfigsModel, Txt2SVGModel from utils.http_utils import cors_decorator import bashful_marketplace @@ -26,8 +26,8 @@ def get_model_configs(request): focus on what they are doing and not the AI generation. """ try: - data = request.json() - model = Txt2ImgModel(**data) + data = request.json + model = GetModelConfigsModel(**data) except ValidationError as e: return ({"error": str(e)}, 400) # Bad Request @@ -40,6 +40,11 @@ def txt2img(request): """ Generate an image from a prompt and model configuration. """ + try: + data = request.json + model = Txt2ImgModel(**data) + except ValidationError as e: + return ({"error": str(e)}, 400) return bashful.txt2img(request) @@ -51,6 +56,32 @@ def img2img(request): """ return bashful.img2img(request) +@functions_framework.http +@cors_decorator +def txt2svg(request): + """ + Generate an svg from the prompt and model configuration. + """ + try: + data = request.json + model = Txt2SVGModel(**data) + except ValidationError as e: + return ({"error": str(e)}, 400) + return bashful.txt2svg(request) + +@functions_framework.http +@cors_decorator +def remove_bg(request): + """ + Generate an svg from the prompt and model configuration. + """ + try: + data = request.json + model = RemoveBGRequest(**data) + except ValidationError as e: + return ({"error": str(e)}, 400) + return bashful.remove_background(request) + ##################### PROJECT: Bashful-Marketplace ##################### @functions_framework.http diff --git a/models/http_models/bashful_http.py b/models/http_models/bashful_http.py index bfbd4e8..8165b76 100644 --- a/models/http_models/bashful_http.py +++ b/models/http_models/bashful_http.py @@ -14,4 +14,8 @@ class Txt2SVGModel(Txt2ImgModel): ... class GetModelConfigsModel(BaseModel): - ... \ No newline at end of file + ... + + +class RemoveBGRequest(BaseModel): + image_url: str \ No newline at end of file diff --git a/photoroom-sdk-result.png b/photoroom-sdk-result.png new file mode 100644 index 0000000..2574d19 Binary files /dev/null and b/photoroom-sdk-result.png differ diff --git a/services/gcp_provider.py b/services/gcp_provider.py index 8fd4a7c..f23c473 100644 --- a/services/gcp_provider.py +++ b/services/gcp_provider.py @@ -30,3 +30,28 @@ def upload_image_to_storage( blob.content_type = "image/png" blob.upload_from_string(image_bytes, content_type="image/png") return blob.public_url + + def upload_svg_to_storage( + self, svg_bytes, gcp_bucket_name, gcp_folder_and_name_destination + ): + """ + Uploads an image to Google Cloud Storage from disk. + + Args: + image_path (str): The path to the image file on disk. + bucket_name (str): The name of the Google Cloud Storage bucket. + destination_blob_name (str): The name of the blob to create. + """ + # Create a new client using Google Cloud Storage credentials + client = storage.Client.from_service_account_json( + "bashful_api_service_account.json" + ) + + # Get the bucket to upload to + bucket = client.get_bucket(gcp_bucket_name) + + # Create a new blob and upload the file's content + blob = bucket.blob(gcp_folder_and_name_destination) + blob.content_type = "svg+xml" + blob.upload_from_string(svg_bytes, content_type="svg+xml") + return blob.public_url \ No newline at end of file