Update app.py

app.py

@@ -5,8 +5,62 @@ import numpy as np
 from moviepy.editor import *
 from share_btn import community_icon_html, loading_icon_html, share_js
 
-token = os.environ.get('HF_TOKEN')
-pix2pix = gr.Blocks.load(name="spaces/fffiloni/instruct-pix2pix-clone", api_key=token)
+from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
+import torch
+from PIL import Image
+import time
+import psutil
+import random
+
+#token = os.environ.get('HF_TOKEN')
+#pix2pix = gr.Blocks.load(name="spaces/fffiloni/instruct-pix2pix-clone", api_key=token)
+
+pipe = DiffusionPipeline.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 🥶"
+
+if torch.cuda.is_available():
+    pipe = pipe.to("cuda")
+
+def pix2pix(
+    prompt,
+    text_guidance_scale,
+    image_guidance_scale,
+    image,
+    steps,
+    neg_prompt="",
+    width=512,
+    height=512,
+    seed=0,
+):
+    print(psutil.virtual_memory())  # print memory usage
+
+    if seed == 0:
+        seed = random.randint(0, 2147483647)
+
+    generator = torch.Generator("cuda").manual_seed(seed)
+
+    try:
+        ratio = min(height / image.height, width / image.width)
+        image = image.resize((int(image.width * ratio), int(image.height * ratio)), Image.LANCZOS)
+
+        result = pipe(
+            prompt,
+            negative_prompt=neg_prompt,
+            image=image,
+            num_inference_steps=int(steps),
+            image_guidance_scale=image_guidance_scale,
+            guidance_scale=text_guidance_scale,
+            generator=generator,
+        )
+
+        # return replace_nsfw_images(result)
+        return result.images, result.nsfw_content_detected, seed
+    except Exception as e:
+        return None, None, error_str(e)
 
 def get_frames(video_in):
     frames = []
@@ -71,7 +125,7 @@ def infer(prompt,video_in, seed_in, trim_value):
     
     for i in frames_list[0:int(n_frame)]:
         pix2pix_img = pix2pix(prompt,5.5,1.5,i,15,"",512,512,seed_in,fn_index=0)
-        images = [os.path.join(pix2pix_img[0], img) for img in os.listdir(pix2pix_img[0])]
+        images = pix2pix_img[0]
         result_frames.append(images[0])
         print("frame " + i + ": done;")
 
@@ -132,10 +186,13 @@ with gr.Blocks(css='style.css') as demo:
                 prompt = gr.Textbox(label="Prompt", placeholder="enter prompt", show_label=False, elem_id="prompt-in")
                 video_inp = gr.Video(label="Video source", source="upload", type="filepath", include_audio=False, elem_id="input-vid")
                 with gr.Row():
-                    seed_inp = gr.Slider(label="Seed", minimum=0, maximum=10000, step=1, value=123456)
+                    seed_inp = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, value=123456)
                     trim_in = gr.Slider(label="Cut video at (s)", minimun=1, maximum=3, step=1, value=1)
             with gr.Column():
-                
+                gr.HTML("""
+                <a style="display:inline-block" href="https://huggingface.co/spaces/fffiloni/Pix2Pix-Video?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a></p>
+            
+                """)
                 video_out = gr.Video(label="Pix2pix video result", elem_id="video-output")
                 submit_btn = gr.Button("Generate Pix2Pix video")