image generation

.gitignore

@@ -1,3 +1,4 @@
 
 *.pyc
 /autogen_cache/42
+*.jpeg

app.py

@@ -3,6 +3,7 @@ from configfile import Config
 from src.hf_autogen.hfautogen import hf_llmconfig
 from src.streamlitui.loadui import LoadStreamlitUI
 from src.usecases.textgen import TexGeneration
+from src.usecases.imggen import ImageGeneration
 
 
 
@@ -37,16 +38,26 @@ if __name__ == "__main__":
                 st.write(problem)
             
             
-            obj_txtgen = TexGeneration(assistant_name="Assistant", user_proxy_name='Userproxy',
+            obj_txt_gen = TexGeneration(assistant_name="Assistant", user_proxy_name='Userproxy',
                                                     llm_config=llm_config,
                                                     problem=problem)
-            obj_txtgen.run()
+            obj_txt_gen.run()
     
     elif user_input['selected_usecase'] == "Image Generation": 
         st.subheader("Image generation")
+        
         if problem:
             with st.chat_message("user"):
                 st.write(problem)
+                
+           
+            obj_img_gen = ImageGeneration(assistant_name="Image_Assistant", user_proxy_name='Userproxy',
+                                                    llm_config=llm_config,
+                                                    problem=problem)
+            obj_img_gen.run()
+            
+            # with st.chat_message('ai'):
+            #     st.image(image.open('./imagegen/response.jpeg'))
                
             
             

src/agents/assistantagent.py

@@ -1,5 +1,7 @@
 from autogen import AssistantAgent
 import streamlit as st
+import base64
+from io import BytesIO
 
 
 class TrackableAssistantAgent(AssistantAgent):
@@ -7,8 +9,13 @@ class TrackableAssistantAgent(AssistantAgent):
         if message and type(message)== str and sender.name =="Userproxy":
             with st.chat_message("user"):
                 st.write(message)
-        
-                
         return super()._process_received_message(message, sender, silent)
 
+class TrackableImageAssistantAgent(AssistantAgent):
+    def _process_received_message(self, message, sender, silent):
+        # with st.chat_message('ai'):
+        #     st.image('./imagegen/response.jpeg')
+        return super()._process_received_message(message, sender, silent)
+
+
 

src/agents/userproxyagent.py

@@ -1,12 +1,17 @@
 from autogen import UserProxyAgent
 import streamlit as st
+import base64
+from io import BytesIO
 
 
 class TrackableUserProxyAgent(UserProxyAgent):
     def _process_received_message(self, message, sender, silent):
-        with st.chat_message(sender.name.lower()):
-            if type(message)==str:
+        
+        if type(message)==str and sender.name == 'Image_Assistant':
+            with st.chat_message('ai'):
+                st.image('./imagegen/response.jpeg')
+
+        else :
+            with st.chat_message('ai'):
                 st.write(message)
-            else :
-                st.write(message['content'])
         return super()._process_received_message(message, sender, silent)

src/hf_autogen/hfautogen.py

@@ -49,7 +49,7 @@ class APIModelClient:
 
         input_data = {
             "inputs": conversation_history,
-            "parameters": {"max_new_tokens": 1000, "return_full_text": False, "do_sample": False},
+            "parameters": {"return_full_text": False, "do_sample": False},
             "options": {"wait_for_model": True, "use_cache": False}
             # Include any other parameters required by your API
         }

src/hf_autogen/imghfautogen.py

@@ -0,0 +1,226 @@
+import autogen
+from autogen import AssistantAgent, UserProxyAgent, GroupChatManager, GroupChat, ConversableAgent
+from types import SimpleNamespace
+import requests
+import json
+import os
+import shutil
+import random
+import streamlit as st
+import base64
+from io import BytesIO
+
+from src.agents.assistantagent import TrackableImageAssistantAgent
+from src.agents.userproxyagent import TrackableUserProxyAgent
+
+class APIModelClient:
+    def __init__(self, config, **kwargs):
+        self.device = config.get("device", "cpu")
+        self.api_url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large"
+        #self.api_url = "https://api-inference.huggingface.co/models/google/gemma-7b-it"  # Add the API URL to the config
+        self.headers = {"Authorization": "Bearer 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"}  # Example: Add any required headers
+
+        self.model_name = config.get("model")
+        self.chat_index = 0
+
+        self.conversion_mem = ""
+
+        # self.tokenizer and self.model lines are removed or modified
+
+    def create(self, params):
+        conversation_history = ""
+
+        for message in params["messages"]:
+            prefix = ""
+            if message["role"] == "system":
+                prefix = f'Bot Description:\n'
+            elif message["role"] == "user":
+                prefix = f'User____:\n'
+            else:
+                prefix = f'Agent ({message["role"]}):\n'
+            conversation_history += prefix + f'{message["content"]}\n\n'
+        
+
+        #try:
+        #_input = f'Given the context of the last message: {params["messages"][-2]["content"]}\n\n\nHere is input on the context: {params["messages"][-1]["content"]}'
+
+
+        #except Exception as e:
+        #    print(e)
+        #    _input = params["messages"][-1]["content"]
+
+        input_data = {
+            "inputs": conversation_history,
+            # "parameters": {"return_full_text": False, "do_sample": False},
+            # "options": {"wait_for_model": True, "use_cache": False}
+            # Include any other parameters required by your API
+        }
+
+        # Sending the request to your model's API
+        response = requests.post(self.api_url, json=input_data, headers=self.headers)
+
+        if response.status_code == 200:
+            return response
+        else:
+            raise Exception(f"API request failed with status code {response.status_code}: {response.text}")
+
+
+    def message_retrieval(self, response):
+        """Retrieve the messages from the response."""
+        import io
+        from PIL import Image
+        image = Image.open(io.BytesIO(response.content))
+        image_path = './imagegen/response.jpeg'
+        image.save(image_path)
+        # Open the image using PIL
+        image = Image.open(image_path)
+
+        # Display the image in Streamlit
+        st.image(image, caption="Loaded Image", use_column_width=True)
+        return [str(response.content)]
+
+    def cost(self, response) -> float:
+        """Calculate the cost of the response."""
+        response.cost = 0
+        return 0
+
+    @staticmethod
+    def get_usage(response):
+        # returns a dict of prompt_tokens, completion_tokens, total_tokens, cost, model
+        # if usage needs to be tracked, else None
+        return {}
+
+
+class APIModelClientWithArguments(APIModelClient):
+    def __init__(self, config, hf_key, hf_url="https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large", **kwargs):
+        self.device = config.get("device", "cpu")
+        self.api_url = hf_url
+        # self.api_url = "https://api-inference.huggingface.co/models/deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"  # Add the API URL to the config
+
+        self.headers = {"Authorization": f"Bearer {hf_key}"}  # Example: Add any required headers
+
+        self.model_name = config.get("model")
+        self.chat_index = 0
+
+        self.conversion_mem = ""
+        
+
+def hf_llmconfig(selected_model):
+    llm_config = {
+        "config_list": [{
+            "model": selected_model,
+            "model_client_cls": "APIModelClientWithArguments",
+            "device": ""
+        }]
+    }
+    st.session_state['llm_config'] = llm_config
+    return llm_config
+def UserAgent(name, llm_config, max_consecutive_auto_reply=1, code_dir="coding", use_docker=False, system_message="You are a helpful AI assistant"):
+    llm_config = {
+        "config_list": [{
+            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            "model_client_cls": "APIModelClientWithArguments",
+            "device": ""
+        }]
+    }
+    user_agent = TrackableUserProxyAgent(
+        name=name,
+        max_consecutive_auto_reply=max_consecutive_auto_reply,
+        llm_config=llm_config,
+        is_termination_msg=lambda x: x.get("content", "").rstrip().endswith("TERMINATE"),
+        code_execution_config={
+            "work_dir": code_dir,
+            "use_docker": use_docker,
+        },
+        system_message=system_message,
+        human_input_mode="NEVER"
+    )
+
+    user_agent.register_model_client(model_client_cls=APIModelClientWithArguments, hf_key=st.session_state["api_key"])
+
+    return user_agent
+
+def ModelAgent(name, llm_config, hf_url="https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large", system_message="", code_execution=False):
+    default_system_message = """You are a helpful AI assistant for generating and manipulating images.
+    """
+
+    if system_message == "":
+        system_message = default_system_message
+
+    # llm_config = {
+    #     "config_list": [{
+    #         "model": "",
+    #         "model_client_cls": "APIModelClientWithArguments",
+    #         "device": ""
+    #     }]
+    # }
+    llm_config =llm_config
+
+
+
+    agent = TrackableImageAssistantAgent(
+        name=name,
+        llm_config=llm_config,
+        system_message=system_message,
+        code_execution_config=code_execution,
+        
+    )
+    agent.register_model_client(model_client_cls=APIModelClientWithArguments, hf_key=st.session_state["api_key"], hf_url=hf_url)
+
+    return agent
+
+
+async def InitChat(user, agent, _input, summary_method="reflection_with_llm"):
+    def clear_directory_contents(dir_path):
+        try:
+            for item in os.listdir(dir_path):
+                item_path = os.path.join(dir_path, item)
+                if os.path.isfile(item_path) or os.path.islink(item_path):
+                    os.remove(item_path)  # Remove files and links
+                elif os.path.isdir(item_path):
+                    shutil.rmtree(item_path)  # Remove directories
+            shutil.rmtree(dir_path)
+            print(f"All contents of '{dir_path}' have been removed.")
+        except FileNotFoundError:
+            pass
+
+    #seed = random.randint(0, 99999)
+    seed = 42
+    #clear_directory_contents(f'./autogen_cache/{seed}')
+
+    custom_cache = autogen.Cache({"cache_seed": seed, "cache_path_root": "autogen_cache"})
+
+    await user.a_initiate_chat(
+        agent,
+        max_turns=1,
+        message=_input,
+        summary_method=summary_method,
+        cache=custom_cache,
+        
+    )
+
+    #clear_directory_contents(f'./autogen_cache/{seed}')    
+
+def GroupChat(user, agents, _input, hf_key, hf_url="https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large", max_round=5):
+    llm_config = {
+            "config_list": [{
+                "model": "",
+                "model_client_cls": "APIModelClientWithArguments",
+                "device": ""
+            }]
+        }
+
+    groupchat = autogen.GroupChat(agents=agents, messages=[], max_round=max_round, speaker_selection_method="round_robin", allow_repeat_speaker=False)
+    manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=llm_config)
+
+    manager.register_model_client(model_client_cls=APIModelClientWithArguments, hf_key=hf_key, hf_url=hf_url)
+    InitChat(user, manager, _input)
+
+#Write me a script to save the BTC chart from the past year to an image.
+
+# if __name__ == "__main__":
+#     print("Running as main")
+
+
+ 
+ 

src/usecases/imggen.py

@@ -0,0 +1,40 @@
+import asyncio
+from src.hf_autogen.imghfautogen import APIModelClientWithArguments,ModelAgent, UserAgent, InitChat
+from src.agents.assistantagent import TrackableAssistantAgent
+from src.agents.userproxyagent import TrackableUserProxyAgent
+import streamlit as st
+
+
+class ImageGeneration:
+    def __init__(self, assistant_name, user_proxy_name, llm_config, problem):
+        # self.assistant = TrackableAssistantAgent(name=assistant_name,
+        #                                          system_message="""you are helpful assistant. Reply "TERMINATE" in 
+        #                                          the end when everything is done """,
+        #                                          human_input_mode="NEVER",
+        #                                          llm_config=llm_config,
+        #                                          )
+        
+        # self.user_proxy = TrackableUserProxyAgent(name=user_proxy_name,
+        #                                           system_message="You are Admin",
+        #                                           human_input_mode="NEVER",
+        #                                           llm_config=llm_config,
+        #                                           code_execution_config=False,
+        #                                           is_termination_msg=lambda x: x.get("content", "").strip().endswith(
+        #                                               "TERMINATE"))
+        
+        
+        self.user = UserAgent(name=user_proxy_name,llm_config=llm_config)
+        self.assistant = ModelAgent(name=assistant_name,
+                                llm_config=llm_config,
+                                hf_url="https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large",
+                                system_message="You are a friendly AI assistant. Your job is to generate image with HD quality")
+
+        self.problem = problem
+        self.loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self.loop)
+
+    # async def initiate_chat(self):
+    #     await InitChat(self.user, self.assistant, self.problem)
+
+    def run(self):
+        self.loop.run_until_complete(InitChat(self.user, self.assistant, self.problem))