Tonic commited on
Commit
89f5b9b
1 Parent(s): d32b65d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +169 -0
app.py ADDED
@@ -0,0 +1,169 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from autogen import AssistantAgent, UserProxyAgent, config_list_from_json
2
+ import autogen
3
+ import replicate
4
+ import requests
5
+ from datetime import datetime
6
+ import http.client
7
+ import json
8
+ import base64
9
+
10
+ config_list = config_list_from_json(env_or_file="OAI_CONFIG_LIST")
11
+ llm_config = {"config_list": config_list, "request_timeout": 120}
12
+
13
+ # function to use llava model to review image
14
+
15
+
16
+ def img_review(image_url, prompt):
17
+ data = {
18
+ "data": [
19
+ {
20
+ "image": "https://picsum.photos/200",
21
+ "features": [],
22
+ },
23
+ ]}
24
+
25
+ headers = {
26
+ "x-api-key": "token 8uOw4ntevc8JKo0Q3tQq:2975e2827ebeb4e103f7b58c1410ba58fa47bc27b1302de614a000bf51bd2114",
27
+ "content-type": "application/json",
28
+ }
29
+
30
+ connection = http.client.HTTPSConnection("api.scenex.jina.ai")
31
+ connection.request("POST", "/v1/describe", json.dumps(data), headers)
32
+ response = connection.getresponse()
33
+
34
+ print(response.status, response.reason)
35
+ response_data = response.read().decode("utf-8")
36
+ print(response_data)
37
+
38
+ connection.close()
39
+
40
+ return response_data
41
+
42
+
43
+ result = img_review(
44
+ "https://cdn.discordapp.com/attachments/1083723388712919182/1089909178266558554/HannaD_A_captivating_digital_artwork_features_a_red-haired_girl_664d73dc-b537-490e-b044-4fbf22733559.png", "a llama driving a car")
45
+ print(result)
46
+
47
+ # def img_review(image_path, prompt):
48
+ # output = replicate.run(
49
+ # "yorickvp/llava-13b:6bc1c7bb0d2a34e413301fee8f7cc728d2d4e75bfab186aa995f63292bda92fc",
50
+ # input={
51
+ # "image": open(image_path, "rb"),
52
+ # "prompt": f"What is happening in the image? From scale 1 to 10, decide how similar the image is to the text prompt {prompt}?",
53
+ # }
54
+ # )
55
+
56
+ # result = ""
57
+ # for item in output:
58
+ # result += item
59
+
60
+ # return result
61
+
62
+
63
+ # function to use stability-ai model to generate image
64
+ def text_to_image_generation(prompt):
65
+ output = replicate.run(
66
+ "stability-ai/sdxl:c221b2b8ef527988fb59bf24a8b97c4561f1c671f73bd389f866bfb27c061316",
67
+ input={
68
+ "prompt": prompt
69
+ }
70
+ )
71
+
72
+ if output and len(output) > 0:
73
+ # Get the image URL from the output
74
+ image_url = output[0]
75
+ print(f"generated image for {prompt}: {image_url}")
76
+
77
+ # Download the image and save it with a filename based on the prompt and current time
78
+ current_time = datetime.now().strftime("%Y%m%d%H%M%S")
79
+ shortened_prompt = prompt[:50]
80
+ filename = f"imgs/{shortened_prompt}_{current_time}.png"
81
+
82
+ response = requests.get(image_url)
83
+ if response.status_code == 200:
84
+ with open(filename, "wb") as file:
85
+ file.write(response.content)
86
+ return f"Image saved as '{filename}'"
87
+ else:
88
+ return "Failed to download and save the image."
89
+ else:
90
+ return "Failed to generate the image."
91
+
92
+
93
+ # Create llm config
94
+ llm_config_assistants = {
95
+ "functions": [
96
+ {
97
+ "name": "text_to_image_generation",
98
+ "description": "use latest AI model to generate image based on a prompt, return the file path of image generated",
99
+ "parameters": {
100
+ "type": "object",
101
+ "properties": {
102
+ "prompt": {
103
+ "type": "string",
104
+ "description": "a great text to image prompt that describe the image",
105
+ }
106
+ },
107
+ "required": ["prompt"],
108
+ },
109
+ },
110
+ {
111
+ "name": "image_review",
112
+ "description": "review & critique the AI generated image based on original prompt, decide how can images & prompt can be improved",
113
+ "parameters": {
114
+ "type": "object",
115
+ "properties": {
116
+ "prompt": {
117
+ "type": "string",
118
+ "description": "the original prompt used to generate the image",
119
+ },
120
+ "image_path": {
121
+ "type": "string",
122
+ "description": "the image file path, make sure including the full file path & file extension",
123
+ }
124
+ },
125
+ "required": ["prompt", "image_path"],
126
+ },
127
+ },
128
+ ],
129
+ "config_list": config_list,
130
+ "request_timeout": 120}
131
+
132
+ # Create assistant agent
133
+ img_gen_assistant = AssistantAgent(
134
+ name="text_to_img_prompt_expert",
135
+ system_message="You are a text to image AI model expert, you will use text_to_image_generation function to generate image with prompt provided, and also improve prompt based on feedback provided until it is 10/10.",
136
+ llm_config=llm_config_assistants,
137
+ function_map={
138
+ "image_review": img_review,
139
+ "text_to_image_generation": text_to_image_generation
140
+ }
141
+ )
142
+
143
+ img_critic_assistant = AssistantAgent(
144
+ name="img_critic",
145
+ system_message="You are an AI image critique, you will use img_review function to review the image generated by the text_to_img_prompt_expert against the original prompt, and provide feedback on how to improve the prompt.",
146
+ llm_config=llm_config_assistants,
147
+ function_map={
148
+ "image_review": img_review,
149
+ "text_to_image_generation": text_to_image_generation
150
+ }
151
+ )
152
+
153
+ # Create user proxy agent
154
+ user_proxy = UserProxyAgent(
155
+ name="user_proxy",
156
+ human_input_mode="ALWAYS",
157
+ )
158
+
159
+ # Create groupchat
160
+ groupchat = autogen.GroupChat(
161
+ agents=[user_proxy, img_gen_assistant, img_critic_assistant], messages=[], max_round=50)
162
+
163
+ manager = autogen.GroupChatManager(
164
+ groupchat=groupchat,
165
+ llm_config=llm_config)
166
+
167
+ # # Start the conversation
168
+ # user_proxy.initiate_chat(
169
+ # manager, message="Generate a photo realistic image of llama driving a car")