lunarflu HF Staff commited on
Commit
dc2fc67
Β·
1 Parent(s): e998610

[tasks] testing something

Browse files
Files changed (1) hide show
  1. app.py +19 -406
app.py CHANGED
@@ -1,39 +1,11 @@
1
  import discord
2
  import os
3
- import threading
4
  import gradio as gr
5
- import requests
6
- import json
7
- import random
8
- import time
9
- import re
10
- from discord import Embed, Color
11
  from discord.ext import commands
12
- # test # fix? # unstick # fix
13
  from gradio_client import Client
14
  from PIL import Image
15
- #from ratelimiter import RateLimiter
16
-
17
  import asyncio
18
 
19
- import shutil # for doing image movement magic
20
-
21
- #todo
22
- # 4 -> combined image
23
- # loading emoji? animated emojis?
24
- # make success / fail emojis more consistent across both painter + dfif
25
- # tasks for concurrent coroutines
26
- # ratelimits
27
-
28
- # enlarge each of 4 images?
29
- # Error: [Errno 104] Connection reset by peer?
30
-
31
- # clean up old threads
32
- # safety for on_reaction_add?
33
- # could use one channel, use threads to organize it. Otherwise may be too split and harder to keep track of
34
- # lock generation after ~120s, can change
35
- # restructure using slash commands? generate -> deepfloydif -> prompt -> thread -> combined -> upscale -> thread
36
-
37
  DFIF_TOKEN = os.getenv('HF_TOKEN')
38
  DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
39
 
@@ -46,392 +18,33 @@ intents.message_content = True
46
 
47
  bot = commands.Bot(command_prefix='!', intents=intents)
48
 
49
-
50
- #---------------------------------------------------------------------------------------------------------------------------------------------
51
- @bot.event
52
- async def on_ready():
53
- print('Logged on as', bot.user)
54
- bot.log_channel = bot.get_channel(1100458786826747945) # 1100458786826747945 = bot-test, 1107006391547342910 = lunarbot server
55
- #---------------------------------------------------------------------------------------------------------------------------------------------
56
- async def safetychecks(ctx):
57
- try:
58
- if ctx.author.bot:
59
- print(f"Error: The bot is not allowed to use its own commands.")
60
- return False
61
-
62
- #βœ…βœ… check if the bot is offline
63
- offline_bot_role_id = 1103676632667017266
64
- bot_member = ctx.guild.get_member(bot.user.id)
65
- if any(role.id == offline_bot_role_id for role in bot_member.roles):
66
- print(f"Error: The bot is offline or under maintenance. (Remove the offline-bot role to bring it online)")
67
- return False
68
-
69
- #βœ…βœ… check if the command is in the allowed channel(s)
70
- channel_id = 1100458786826747945
71
- if ctx.channel.id != 1100458786826747945:
72
- print(f"Error: This is not a permitted channel for that command. (bot-test is the correct channel)")
73
- return False
74
-
75
- #βœ…βœ… check if the user has the required role(s)
76
- guild_id = 879548962464493619
77
- required_role_id = 900063512829755413 # @verified for now
78
- guild = bot.get_guild(guild_id)
79
- required_role = guild.get_role(required_role_id)
80
- if required_role not in ctx.author.roles:
81
- print(f"Error: The user does not have the required role to use that command. ({required_role} is the correct role)")
82
- return False
83
-
84
- return True
85
-
86
- except Exception as e:
87
- print(f"Error: safetychecks failed somewhere, command will not continue.")
88
- await ctx.message.reply(f"<@811235357663297546> SC failed somewhere") # this will always ping, as long as the bot has access to the channel
89
- #----------------------------------------------------------------------------------------------------------------------------------------------
90
- # jojo
91
- @bot.command()
92
- async def jojo(ctx):
93
- # img + face βœ…
94
- # img + no face βœ…
95
- # no image βœ…
96
- # no generation βœ…
97
- # responsive? βœ…
98
- # ratelimits? ❌
99
- # safety checks?βœ…
100
- # bot no crash βœ…
101
- try:
102
- if await safetychecks(ctx): #βœ…
103
- await ctx.message.add_reaction('πŸ€–')
104
- thread = await ctx.message.create_thread(name=f'{ctx.author} Jojo Thread')
105
- if ctx.message.attachments:
106
- await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...yare yare, daze ...')
107
- attachment = ctx.message.attachments[0]
108
- style = 'JoJo'
109
- im = jojogan.predict(attachment.url, style)
110
- #await ctx.message.reply(f'Here is the {style} version of it', file=discord.File(im))
111
- await thread.send(f'{ctx.author.mention}Here is the {style} version of it', file=discord.File(im))
112
- await ctx.message.add_reaction('βœ…') # img + face
113
- else: # no image
114
- await thread.send(f"{ctx.author.mention}No attachments to be found...Can't animify dat! Try sending me an image πŸ˜‰")
115
- await ctx.message.add_reaction('❌')
116
- except Exception as e: # no generation / img + no face
117
- print(f"Error: {e}")
118
- await thread.send(f"{ctx.author.mention}Error: {e}")
119
- await ctx.message.add_reaction('❌')
120
-
121
- #----------------------------------------------------------------------------------------------------------------------------------------------
122
- # Disney
123
- @bot.command()
124
- async def disney(ctx):
125
- try:
126
- if await safetychecks(ctx): #βœ…
127
- await ctx.message.add_reaction('πŸ€–')
128
- thread = await ctx.message.create_thread(name=f'{ctx.author} disney Thread')
129
- if ctx.message.attachments:
130
- await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
131
- attachment = ctx.message.attachments[0]
132
- style = 'disney'
133
- im = jojogan.predict(attachment.url, style)
134
- await thread.send(f'{ctx.author.mention}Here is the {style} version of it', file=discord.File(im))
135
- await ctx.message.add_reaction('βœ…') # img + face
136
- else: # no image
137
- await thread.send(f"{ctx.author.mention}No attachments to be found...Can't animify dat! Try sending me an image πŸ˜‰")
138
- await ctx.message.add_reaction('❌')
139
- except Exception as e: # no generation / img + no face
140
- print(f"Error: {e}")
141
- await thread.send(f"{ctx.author.mention}Error: {e}")
142
- await ctx.message.add_reaction('❌')
143
-
144
- #----------------------------------------------------------------------------------------------------------------------------------------------
145
- # Spider-Verse
146
- @bot.command()
147
- async def spiderverse(ctx):
148
- try:
149
- if await safetychecks(ctx): #βœ…
150
- await ctx.message.add_reaction('πŸ€–')
151
- thread = await ctx.message.create_thread(name=f'{ctx.author} spider-verse Thread')
152
- if ctx.message.attachments:
153
- await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
154
- attachment = ctx.message.attachments[0]
155
- style = 'Spider-Verse'
156
- im = jojogan.predict(attachment.url, style)
157
- await thread.send(f'{ctx.author.mention}Here is the {style} version of it', file=discord.File(im))
158
- await ctx.message.add_reaction('βœ…') # img + face
159
- else: # no image
160
- await thread.send(f"{ctx.author.mention}No attachments to be found...Can't animify dat! Try sending me an image πŸ˜‰")
161
- await ctx.message.add_reaction('❌')
162
- except Exception as e: # no generation / img + no face
163
- print(f"Error: {e}")
164
- await thread.send(f"{ctx.author.mention}Error: {e}")
165
- await ctx.message.add_reaction('❌')
166
-
167
- #----------------------------------------------------------------------------------------------------------------------------------------------
168
- # sketch
169
- @bot.command()
170
- async def sketch(ctx):
171
- try:
172
- if await safetychecks(ctx): #βœ…
173
- await ctx.message.add_reaction('πŸ€–')
174
- thread = await ctx.message.create_thread(name=f'{ctx.author} sketch Thread')
175
- if ctx.message.attachments:
176
- await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
177
- attachment = ctx.message.attachments[0]
178
- style = 'sketch'
179
- im = jojogan.predict(attachment.url, style)
180
- await thread.send(f'{ctx.author.mention}Here is the {style} version of it', file=discord.File(im))
181
- await ctx.message.add_reaction('βœ…') # img + face
182
- else: # no image
183
- await thread.send(f"{ctx.author.mention}No attachments to be found...Can't animify dat! Try sending me an image πŸ˜‰")
184
- await ctx.message.add_reaction('❌')
185
- except Exception as e: # no generation / img + no face
186
- print(f"Error: {e}")
187
- await thread.send(f"{ctx.author.mention}Error: {e}")
188
- await ctx.message.add_reaction('❌')
189
- #----------------------------------------------------------------------------------------------------------------------------------------------
190
- # Stage 1
191
- @bot.command()
192
- async def deepfloydif(ctx, *, prompt: str):
193
- try:
194
- try:
195
- if await safetychecks(ctx): #βœ…
196
- await ctx.message.add_reaction('πŸ€–') # loading emoji?
197
- dfif_command_message_id = ctx.message.id # we will use this in some magic later on
198
- thread = await ctx.message.create_thread(name=f'{ctx.author} DeepfloydIF Image Upscaling Thread ')
199
- # create thread -> send new message inside thread + combined_image -> add reactions -> dfif2
200
-
201
- #current_time = int(time.time())
202
- #random.seed(current_time)
203
-
204
- negative_prompt = ''
205
- seed = random.randint(0, 1000)
206
- #seed = 1
207
- number_of_images = 4
208
- guidance_scale = 7
209
- custom_timesteps_1 = 'smart50'
210
- number_of_inference_steps = 50
211
-
212
- await thread.send(f'{ctx.author.mention}Generating images in thread, can take ~1 minute...')
213
-
214
- except Exception as e:
215
- print(f"Error: {e}")
216
- await ctx.reply('stage 1 error -> pre generation')
217
- await ctx.message.add_reaction('❌')
218
-
219
- try:
220
- stage_1_results, stage_1_param_path, stage_1_result_path = df.predict(
221
- prompt, negative_prompt, seed, number_of_images, guidance_scale, custom_timesteps_1, number_of_inference_steps, api_name='/generate64')
222
-
223
- partialpath = stage_1_result_path[5:] #magic for later
224
- except Exception as e:
225
- print(f"Error: {e}")
226
- await ctx.reply('stage 1 error -> during generation')
227
- await ctx.message.add_reaction('❌')
228
-
229
- try:
230
- png_files = [f for f in os.listdir(stage_1_results) if f.endswith('.png')]
231
-
232
- if png_files:
233
- first_png = png_files[0]
234
- second_png = png_files[1]
235
- third_png = png_files[2]
236
- fourth_png = png_files[3]
237
-
238
- first_png_path = os.path.join(stage_1_results, first_png)
239
- second_png_path = os.path.join(stage_1_results, second_png)
240
- third_png_path = os.path.join(stage_1_results, third_png)
241
- fourth_png_path = os.path.join(stage_1_results, fourth_png)
242
-
243
- img1 = Image.open(first_png_path)
244
- img2 = Image.open(second_png_path)
245
- img3 = Image.open(third_png_path)
246
- img4 = Image.open(fourth_png_path)
247
-
248
- combined_image = Image.new('RGB', (img1.width * 2, img1.height * 2))
249
 
250
- combined_image.paste(img1, (0, 0))
251
- combined_image.paste(img2, (img1.width, 0))
252
- combined_image.paste(img3, (0, img1.height))
253
- combined_image.paste(img4, (img1.width, img1.height))
254
-
255
- combined_image_path = os.path.join(stage_1_results, f'{partialpath}{dfif_command_message_id}.png')
256
- combined_image.save(combined_image_path)
257
-
258
- with open(combined_image_path, 'rb') as f:
259
- combined_image_dfif = await thread.send(f'{ctx.author.mention}React with the image number you want to upscale!', file=discord.File(
260
- f, f'{partialpath}{dfif_command_message_id}.png')) # named something like: tmpgtv4qjix1111269940599738479.png
261
 
262
- async def react1234(emoji, combined_image_dfif):
263
- for emoji in ['1️⃣', '2️⃣', '3️⃣', '4️⃣']:
264
- await combined_image_dfif.add_reaction(emoji)
265
 
266
- await react1234(emoji, combined_image_dfif)
267
-
268
- ''' individual images
269
- if png_files:
270
- for i, png_file in enumerate(png_files):
271
- png_file_path = os.path.join(stage_1_results, png_file)
272
- img = Image.open(png_file_path)
273
- image_path = os.path.join(stage_1_results, f'{i+1}{partialpath}.png')
274
- img.save(image_path)
275
- with open(image_path, 'rb') as f:
276
- await thread.send(f'{ctx.author.mention}Image {i+1}', file=discord.File(f, f'{i+1}{partialpath}.png'))
277
- await asyncio.sleep(1)
278
-
279
- '''
280
-
281
- except Exception as e:
282
- print(f"Error: {e}")
283
- await ctx.reply('stage 1 error -> posting images in thread')
284
- await ctx.message.add_reaction('❌')
285
-
286
- #deepfloydif try/except
287
- except Exception as e:
288
- print(f"Error: {e}")
289
- await ctx.reply('An error occurred in stage 1 for deepfloydif')
290
  await ctx.message.add_reaction('❌')
291
 
292
  #----------------------------------------------------------------------------------------------------------------------------
293
- # Stage 2
294
- async def dfif2(index: int, stage_1_result_path, thread, dfif_command_message_id): # add safetychecks
295
- try:
296
- await thread.send(f"inside dfif2, upscaling")
297
- selected_index_for_stage_2 = index
298
- seed_2 = 0
299
- guidance_scale_2 = 4
300
- custom_timesteps_2 = 'smart50'
301
- number_of_inference_steps_2 = 50
302
- result_path = df.predict(stage_1_result_path, selected_index_for_stage_2, seed_2,
303
- guidance_scale_2, custom_timesteps_2, number_of_inference_steps_2, api_name='/upscale256')
304
-
305
- await thread.send(f"upscale done")
306
- with open(result_path, 'rb') as f:
307
- await thread.send(f'Here is the upscaled image! :) ', file=discord.File(f, 'result.png'))
308
-
309
- # using custom emoji that looks nicer
310
- emoji_guild = thread.guild
311
- confirm_emoji_id = 1098629085955113011
312
- confirm_emoji = discord.utils.get(emoji_guild.emojis, id=confirm_emoji_id)
313
-
314
- # assuming dfif2 is always inside a thread, we can always exit the thread to find the channel with the original command,
315
- # which allows us to react confirm on that message.
316
- parent_channel = thread.parent
317
- dfif_command_message = await parent_channel.fetch_message(dfif_command_message_id)
318
-
319
- # reacting to original !deepfloydif command + using a custom emoji to do it
320
- await dfif_command_message.add_reaction(confirm_emoji)
321
- await thread.send(f"upscale posted")
322
- #await ctx.reply('Here is the result of the second stage', file=discord.File(f, 'result.png'))
323
- #await ctx.message.add_reaction('βœ…') need to fix this
324
 
325
- '''
326
- try:
327
- dfif_command_message = await channel.fetch_message(dfif_command_message_id)
328
- await dfif_command_message.add_reaction('βœ…')
329
- '''
330
-
331
- except Exception as e:
332
- print(f"Error: {e}")
333
-
334
- #await ctx.reply('An error occured in stage 2') need to fix
335
- #await ctx.message.add_reaction('❌')
336
  #----------------------------------------------------------------------------------------------------------------------------
337
- @bot.event
338
- async def on_reaction_add(reaction, user): # ctx = await bot.get_context(reaction.message)? could try later, might simplify
339
- try:
340
- # safety checks first ❌
341
- '''
342
- if user.bot:
343
- return
344
-
345
- #offline bot check ❌
346
- offline_bot_role_id = 1103676632667017266
347
- bot_member = reaction.message.guild.get_member(bot.user.id)
348
- if any(role.id == offline_bot_role_id for role in bot_member.roles):
349
- return
350
-
351
- # verified role check ❌
352
- guild = reaction.message.guild
353
- required_role_id = 900063512829755413 # @verified for now
354
- required_role = guild.get_role(required_role_id)
355
- if required_role not in user.roles:
356
- return
357
-
358
- #channel check ❌
359
- if reaction.message.channel.id != 1100458786826747945:
360
- return
361
-
362
- '''
363
-
364
- thread = reaction.message.channel
365
- threadparentid = thread.parent.id
366
-
367
- if not user.bot:
368
- if threadparentid == 1100458786826747945: # bot-test
369
- # 811235357663297546 = lunarflu
370
- if reaction.message.attachments:
371
- if user.id == reaction.message.mentions[0].id: # if user.id == reaction.message.mentions[0].id:
372
- # magic begins
373
- await reaction.message.channel.send("reaction detected")
374
- attachment = reaction.message.attachments[0]
375
- image_name = attachment.filename # named something like: tmpgtv4qjix1111269940599738479.png
376
- # remove .png first
377
- partialpathmessageid = image_name[:-4] # should be tmpgtv4qjix1111269940599738479
378
- # extract partialpath, messageid
379
- partialpath = partialpathmessageid[:11] # tmpgtv4qjix
380
- messageid = partialpathmessageid[11:] # 1111269940599738479
381
- # add /tmp/ to partialpath, save as new variable
382
- fullpath = "/tmp/" + partialpath # should be /tmp/tmpgtv4qjix
383
- await reaction.message.channel.send(f"fullpath extracted, {fullpath}")
384
- emoji = reaction.emoji
385
-
386
- if emoji == "1️⃣":
387
- index = 0
388
- elif emoji == "2️⃣":
389
- index = 1
390
- elif emoji == "3️⃣":
391
- index = 2
392
- elif emoji == "4️⃣":
393
- index = 3
394
-
395
- await reaction.message.channel.send(f"index extracted, {index}")
396
- index = index
397
- stage_1_result_path = fullpath
398
- thread = reaction.message.channel
399
- dfif_command_message_id = messageid
400
- await reaction.message.channel.send(f"calling dfif2")
401
- await dfif2(index, stage_1_result_path, thread, dfif_command_message_id)
402
-
403
- '''
404
-
405
- if reaction.message.attachments:
406
- if user.id == reaction.message.mentions[0].id: # all we care about is upscaling whatever image this is
407
-
408
- # magic begins
409
- attachment = reaction.message.attachments[0]
410
- image_name = attachment.filename
411
- # we know image_name will be something like 1tmpgtv4qjix.png
412
- # remove .png first
413
- indexpartialpath = image_name[:-4] # should be 1tmpgtv4qjix
414
- # extract index as an integer (dfif2 needs integer)
415
- index = int(indexpartialpath[0]) - 1# should be 1
416
- # extract partialpath
417
- partialpath = indexpartialpath[1:] # should be tmpgtv4qjix
418
- # add /tmp/ to partialpath, save as new variable
419
- fullpath = "/tmp/" + partialpath # should be /tmp/tmpgtv4qjix
420
-
421
- stage_1_result_path = fullpath
422
- index = index
423
- await dfif2(index, stage_1_result_path, thread)
424
-
425
- '''
426
-
427
-
428
- except Exception as e:
429
- print(f"Error: {e}")
430
-
431
- #await ctx.reply('An error occured in stage 2') need to fix
432
- #await ctx.message.add_reaction('❌')
433
- #----------------------------------------------------------------------------------------------------------------------------
434
-
435
 
436
  def run_bot():
437
  bot.run(DISCORD_TOKEN)
 
1
  import discord
2
  import os
 
3
  import gradio as gr
 
 
 
 
 
 
4
  from discord.ext import commands
 
5
  from gradio_client import Client
6
  from PIL import Image
 
 
7
  import asyncio
8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
  DFIF_TOKEN = os.getenv('HF_TOKEN')
10
  DISCORD_TOKEN = os.environ.get("GRADIOTEST_TOKEN", None)
11
 
 
18
 
19
  bot = commands.Bot(command_prefix='!', intents=intents)
20
 
21
+ #---------------------------------------------------------------------------------------------------------------------------
22
+ async def generate_jojo_image(ctx):
23
+ if ctx.message.attachments:
24
+ attachment = ctx.message.attachments[0]
25
+ await ctx.message.add_reaction('πŸ€–')
26
+ await ctx.message.add_reaction('βŒ›')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27
 
28
+ style = 'JoJo'
29
+ im = jojogan.predict(attachment.url, style)
30
+ await ctx.send(f'{ctx.author.mention} Here is the {style} version of it')
31
+ await ctx.send(file=discord.File(im))
 
 
 
 
 
 
 
32
 
33
+ await ctx.message.remove_reaction('βŒ›', bot.user)
34
+ await ctx.message.add_reaction('βœ…')
 
35
 
36
+ else:
37
+ await ctx.send(f"{ctx.author.mention} No attachments to be found... Can't animify dat! Try sending me an image πŸ˜‰")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
  await ctx.message.add_reaction('❌')
39
 
40
  #----------------------------------------------------------------------------------------------------------------------------
41
+ @bot.command()
42
+ async def jojo(ctx):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
 
44
+ # task test
45
+ task = asyncio.create_task(generate_jojo_image(ctx))
46
+ await task
 
 
 
 
 
 
 
 
47
  #----------------------------------------------------------------------------------------------------------------------------
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
 
49
  def run_bot():
50
  bot.run(DISCORD_TOKEN)