large code refactoring, no major changes to functionality

extensions/stable_diffusion.py

@@ -186,47 +186,37 @@ class StableDiffusion(commands.Cog):
         output += lora_options
         await ctx.send(output)
 
-    @commands.command(
+    async def get_image_from_ctx(self, ctx):
-        description="Imagine", 
+        if ctx.message.attachments:
-        help="Generate an image using stable diffusion. You can add keyword arguments to your prompt and they will be treated as stable diffusion options. Usage !imagine (topic)", 
+            file_url = ctx.message.attachments[0].url
-        brief="Generate an image"
+            return file_url
-        ) 
-    async def imagine(self, ctx):
-        url = self.stable_diffusion_url
-        if url == "disabled":
-            await ctx.send("Command is currently disabled.")
-            return
-        else:
-            url=f"{url}/sdapi/v1/txt2img"
-        prompt = self.get_prompt_from_ctx(ctx)
-        key_value_pairs = self.get_kv_from_ctx(ctx)
-        if prompt == None:
-            prompt = await self.generate_prompt()
         try:
-            neg_prompt_file = f"{self.data_dir}negative_prompt.txt"
+            file_url = ctx.message.content.split(" ", maxsplit=1)[1]
-            with open(neg_prompt_file, 'r') as f:
+            return file_url
-                negative_prompt = f.readline()
         except:
-            neg_prompt_file = f"{self.data_dir}negative_prompt.txt"
+            self.bot.logger.info("Couldn't find image.")
-            with open(neg_prompt_file, 'w') as f:
+            return None
-                f.writelines(self.default_neg_prompt)
+
-                negative_prompt = self.default_neg_prompt
+    async def txt2img(self, ctx, prompt):
-        await ctx.send(f"Please be patient this may take some time! Generating: {prompt}.")
+        url = f"{self.stable_diffusion_url}/sdapi/v1/txt2img"
+        key_value_pairs = self.get_kv_from_ctx(ctx)
+        headers = {'Content-Type': 'application/json'}  
         payload = {
             "prompt": prompt,
             "steps": 25,
-            "negative_prompt": negative_prompt
+            "negative_prompt": self.get_negative_prompt() 
-        }
-        headers = {
-            'Content-Type': 'application/json'
         }
         if key_value_pairs:
             payload.update(key_value_pairs)
         try:
             async with self.bot.http_session.post(url, headers=headers, json=payload) as resp:
+                if resp.status != 200:
+                    await ctx.send(f"{resp.status} {resp.reason}")
+                    self.bot.logger.exception(f"{resp.status} {resp.reason}")
+                    return
                 r = await resp.json()
         except ConnectionRefusedError:
-            await ctx.send("Failed to conenct to image generation service")
+            await ctx.send("Failed to connect to image generation service")
             self.bot.logger.exception("Failed to connect to image generation service")
             return
         except:
@@ -234,7 +224,57 @@ class StableDiffusion(commands.Cog):
             self.bot.logger.exception("Failed to generate image")
             return
         
-        for i in r['images']:
+        await self.send_generated_image(ctx, r['images'], prompt)
+    
+    async def save_image(self, url):
+        async with self.bot.http_session.get(url) as response:
+            image_name = self.working_dir + str(time.time_ns()) + ".png"
+            with open(image_name, 'wb') as out_file:
+                self.bot.logger.debug(f"Saving image: {image_name}")
+                while True:
+                    chunk = await response.content.read(1024)
+                    if not chunk:
+                        break
+                    out_file.write(chunk)
+        return image_name
+
+    async def img2img(self, ctx, prompt):
+        url = f"{self.stable_diffusion_url}/sdapi/v1/img2img"
+        file_url = await self.get_image_from_ctx(ctx)
+        image_name = await self.save_image(file_url)
+        file_url = await self.my_open_img_file(image_name)
+        key_value_pairs = self.get_kv_from_ctx(ctx)
+        headers = {'Content-Type': 'application/json'}  
+        payload = {
+            "init_images": [file_url],
+            "prompt": prompt,
+            "steps": 40,
+            "negative_prompt": self.get_negative_prompt(),
+            "denoising_strength": 0.5,
+        }
+        if key_value_pairs:
+            payload.update(key_value_pairs)
+        try:
+            async with self.bot.http_session.post(url, headers=headers, json=payload) as resp:
+                if resp.status != 200:
+                    await ctx.send(f"{resp.status} {resp.reason}")
+                    self.bot.logger.error(f"{resp.status} {resp.reason}")
+                    return
+                r = await resp.json()
+        except ConnectionRefusedError:
+            await ctx.send("Failed to connect to image generation service")
+            self.bot.logger.exception("Failed to connect to image generation service")
+            return
+        except:
+            await ctx.send("Failed to generate image")
+            self.bot.logger.exception("Failed to generate image")
+            return
+        
+        await self.send_generated_image(ctx, r['images'], prompt)
+
+
+    async def send_generated_image(self, ctx, images, prompt):
+        for i in images:
             image = Image.open(io.BytesIO(base64.b64decode(i.split(",", 1)[0])))
             try:
                 if ctx.channel.is_nsfw():
@@ -249,6 +289,7 @@ class StableDiffusion(commands.Cog):
                     
             with open(filepath, "rb") as fh:
                 f = discord.File(fh, filename=filepath)
+            
             prompt = prompt.replace('\n',' ')
             log_data = f'Author: {ctx.author.name}, Prompt: {prompt}, Filename: {my_filename}\n'
             with open(f"{self.data_dir}stable_diffusion.log", 'a') as log_file:
@@ -257,6 +298,38 @@ class StableDiffusion(commands.Cog):
             await ctx.send(f'Generated by: {ctx.author.name}\nPrompt: {prompt}', file=f)
 
 
+    def get_negative_prompt(self):
+        try:
+            neg_prompt_file = f"{self.data_dir}negative_prompt.txt"
+            with open(neg_prompt_file, 'r') as f:
+                negative_prompt = f.readline()
+        except:
+            neg_prompt_file = f"{self.data_dir}negative_prompt.txt"
+            with open(neg_prompt_file, 'w') as f:
+                f.writelines(self.default_neg_prompt)
+                negative_prompt = self.default_neg_prompt
+        return negative_prompt
+
+
+    @commands.command(
+    description="Imagine",
+    help="Generate an image using stable diffusion. You can add keyword arguments to your prompt and they will be treated as stable diffusion options. Usage !imagine (topic)",
+    brief="Generate an image"
+    )
+    async def imagine(self, ctx):
+        url = self.stable_diffusion_url
+        if url == "disabled":
+            await ctx.send("Command is currently disabled")
+            return
+        prompt = self.get_prompt_from_ctx(ctx)
+        if prompt == None:
+            prompt = await self.generate_prompt()
+        
+        await ctx.send(f"Please be patient this may take some time! Generating: {prompt}.")
+
+        await self.txt2img(ctx, prompt)
+
+        
     @commands.command(
         description="Describe", 
         help="Get better understanding of what the bot \"sees\" when you post an image! (Runs it through CLIP) Usage !describe (image link)", 
@@ -269,28 +342,9 @@ class StableDiffusion(commands.Cog):
             return
         else:
             url=f"{url}/sdapi/v1/interrogate"
-        try:
+        file_url = await self.get_image_from_ctx(ctx)
-            if ctx.message.content.startswith("!describe "):
+        image_name = await self.save_image(file_url)
-                file_url = ctx.message.content.split(" ", maxsplit=1)[1]
+        img_link = await self.my_open_img_file(image_name)
-            elif ctx.message.attachments:
-                file_url = ctx.message.attachments[0].url
-            else:
-                self.bot.logger.debug("No image linked or attached.")
-                return
-        except:
-            self.bot.logger.exception("Couldn't find image.")
-            return   
-        async with self.bot.http_session.get(file_url) as response:
-            imageName = self.working_dir + str(time.time_ns()) + ".png"
-            with open(imageName, 'wb') as out_file:
-                self.bot.logger.debug(f"Saving image: {imageName}")
-                while True:
-                    chunk = await response.content.read(1024)
-                    if not chunk:
-                        break
-                    out_file.write(chunk)
-
-        img_link = await self.my_open_img_file(imageName)
         try:
             payload = {"image": img_link}
             async with self.bot.http_session.post(url, json=payload) as response:
@@ -310,65 +364,11 @@ class StableDiffusion(commands.Cog):
         if url == "disabled":
             await ctx.send("Command is currently disabled")
             return
-        try:
-            if ctx.message.attachments:
-                file_url = ctx.message.attachments[0].url
-            elif ctx.message.content.startswith("!reimagine "):
-                file_url = ctx.message.content.split(" ", maxsplit=2)[1]
-            else:
-                await ctx.send("No image linked or attached.")
-                return
-        except:
-            self.bot.logger.exception("Couldn't find image.")
-            return
         prompt = self.get_prompt_from_ctx(ctx)
         if not prompt:
             prompt = ""
-        key_value_pairs = self.get_kv_from_ctx(ctx)
+        await ctx.send(f"Please be patient this may take some time! Generating: {prompt}.")
-        try:
+        await self.img2img(ctx, prompt)
-            async with self.bot.http_session.get(file_url) as response:
-                imageName = self.working_dir + str(time.time_ns()) + ".png"
-                with open(imageName, 'wb') as out_file:
-                    self.bot.logger.debug(f"Saving image: {imageName}")
-                    while True:
-                        chunk = await response.content.read(1024)
-                        if not chunk:
-                            break
-                        out_file.write(chunk)
-                            
-        except ConnectionRefusedError:
-            await ctx.send("My image generation service may not be running.")
-            self.bot.logger.exception("Couldn't connect to image generation service")
-        except:
-            await ctx.send("Failed to download image")
-            self.bot.logger.exception("Failed to download image")
-            return
-
-        img_link = await self.my_open_img_file(imageName)
-
-        negative_prompt = "badhandsv4, worst quality, lowres, EasyNegative, hermaphrodite, cropped, not in the frame, additional faces, jpeg large artifacts, jpeg small artifacts, ugly, tiling, poorly drawn hands, poorly drawn feet, poorly drawn face, out of frame, extra limbs, disfigured, deformed, body out of frame, blurry, bad anatomy, blurred, watermark, grainy, signature, cut off, draft, not finished drawing, unfinished image, bad eyes, doll, 3d, cartoon, (bad eyes:1.2), (worst quality:1.2), (low quality:1.2), bad-image-v2-39000, (bad_prompt_version2:0.8), nude, badhandv4 By bad artist -neg easynegative ng_deepnegative_v1_75t verybadimagenegative_v1.3, (Worst Quality, Low Quality:1.4), Poorly Made Bad 3D, Lousy Bad Realistic, nsfw,(worst quality, low quality:1.4), (lip, nose, tooth, rouge, lipstick, eyeshadow:1.4), ( jpeg artifacts:1.4), (depth of field, bokeh, blurry, film grain, chromatic aberration, lens flare:1.0), (1boy, abs, muscular, rib:1.0), greyscale, monochrome, dusty sunbeams, trembling, motion lines, motion blur, emphasis lines, text, title, logo, signature, child, childlike, young, easynegative, (bad-hands-5:0.8), plain background, monochrome, poorly drawn face, poorly drawn hands, watermark, censored, (mutated hands and fingers), ugly, worst quality, low quality,, nsfw,(worst quality, low quality:1.4), (lip, nose, tooth, rouge, lipstick, eyeshadow:1.4), ( jpeg artifacts:1.4), (depth of field, bokeh, blurry, film grain, chromatic aberration, lens flare:1.0), (1boy, abs, muscular, rib:1.0), greyscale, monochrome, dusty sunbeams, trembling, motion lines, motion blur, emphasis lines, text, title, logo, signature, child, childlike, young"
-
-        await ctx.send("Please be patient this may take some time! Generating: " + prompt + ".")
-
-        payload = {"init_images": [img_link], "prompt": prompt, "steps": 40, "negative_prompt": negative_prompt, "denoising_strength": 0.5}
-        if key_value_pairs:
-            payload.update(key_value_pairs)
-
-        try:
-            async with self.bot.http_session.post(url=f'{url}/sdapi/v1/img2img', json=payload) as response:
-                data = await response.json()
-                for i in data['images']:
-                    if not os.path.isdir(f"{self.working_dir}reimagined/"+ str(ctx.author.id)):
-                        os.makedirs(f"{self.working_dir}reimagined/"+ str(ctx.author.id))
-                    image = Image.open(io.BytesIO(base64.b64decode(i.split(",",1)[0])))
-                    my_filename = self.working_dir + str(time.time_ns()) + ".png"
-                    image.save(my_filename)
-                    with open(my_filename, "rb") as fh:
-                        f = discord.File(fh, filename=my_filename)
-                    await ctx.send(file=f)
-        except ConnectionRefusedError:
-            await ctx.send("My image generation service may not be running.")
-            self.bot.logger.exception("Couldn't connect to image generation service")
             
     @commands.command(
         description="Negative Prompt",