Update app.py

app.py

@@ -5,18 +5,16 @@ import numpy as np
 from moviepy.editor import *
 from share_btn import community_icon_html, loading_icon_html, share_js
 
-from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionInstructPix2PixPipeline
 import torch
 from PIL import Image
 import time
 import psutil
+import math
 import random
 
 
 pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained("timbrooks/instruct-pix2pix", torch_dtype=torch.float16, safety_checker=None)
-pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-pipe.enable_xformers_memory_efficient_attention()
-pipe.unet.to(memory_format=torch.channels_last)
 
 device = "GPU 🔥" if torch.cuda.is_available() else "CPU 🥶"
 
@@ -24,37 +22,37 @@ if torch.cuda.is_available():
     pipe = pipe.to("cuda")
 
 
-def pix2pix(prompt,
-            image,
-            steps,
-            image_guidance_scale,
-            seed):
-    print(psutil.virtual_memory())  # print memory usage
-
-    if seed == 0:
-        seed = random.randint(0, 2147483647)
-
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    #try:
-    image = Image.open(image)
-    ratio = min(512 / image.height, 512 / image.width)
-    image = image.resize((int(image.width * ratio), int(image.height * ratio)), Image.LANCZOS)
-
-    result = pipe(
-        prompt,
-        image=image,
-        num_inference_steps=int(steps),
-        image_guidance_scale=image_guidance_scale,
-        #guidance_scale=text_guidance_scale,
-        generator=generator,
-        #negative_prompt=neg_prompt,
-    )
-
-        # return replace_nsfw_images(result)
-    return result[0]
-    #except Exception as e:
-    #    return None
+def pix2pix(
+        input_image: Image.Image,
+        instruction: str,
+        steps: int,
+        randomize_seed: bool,
+        seed: int,
+        randomize_cfg: bool,
+        text_cfg_scale: float,
+        image_cfg_scale: float,
+    ):
+        seed = random.randint(0, 100000) if randomize_seed else seed
+        text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
+        image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale
+
+        width, height = input_image.size
+        factor = 512 / max(width, height)
+        factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
+        width = int((width * factor) // 64) * 64
+        height = int((height * factor) // 64) * 64
+        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
+
+        if instruction == "":
+            return [input_image, seed]
+
+        generator = torch.manual_seed(seed)
+        edited_image = pipe(
+            instruction, image=input_image,
+            guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,
+            num_inference_steps=steps, generator=generator,
+        ).images[0]
+        return edited_image
 
 
 
@@ -120,7 +118,8 @@ def infer(prompt,video_in, seed_in, trim_value):
     print("set stop frames to: " + str(n_frame))
     
     for i in frames_list[0:int(n_frame)]:
-        pix2pix_img = pix2pix(prompt,i,5.5,15,seed_in)
+        pil_i = Image.open(i)
+        pix2pix_img = pix2pix(pil_i, prompt, 50, False, seed_in, False, 7.5, 1.5)
         print(pix2pix_img)
         image = Image.open(pix2pix_img)
         rgb_im = image.convert("RGB")