Update app.py

app.py

@@ -40,8 +40,6 @@ def check_nsfw_images(images: list[Image.Image]) -> list[bool]:
     has_nsfw_concepts = safety_checker(images=[images], clip_input=safety_checker_input.pixel_values.to(device))
     return has_nsfw_concepts
 
-# Function 
-@spaces.GPU(enable_queue=True)
 def generate_image(prompt, base, motion, step, progress=gr.Progress()):
     global step_loaded
     global base_loaded
@@ -71,36 +69,41 @@ def generate_image(prompt, base, motion, step, progress=gr.Progress()):
 
     output = pipe(
         prompt=prompt,
-        width=1024,
-        height=576,
+
+        # this corresponds roughly to 16:9
+        # which is the aspect ratio video used by AiTube
+        width=910, # 1024,
+        height=512, # 576,
+        
         guidance_scale=1.0,
         num_inference_steps=step,
         callback=progress_callback,
         callback_steps=1
     )
 
-    has_nsfw_concepts = check_nsfw_images([output.frames[0][0]])
-    if has_nsfw_concepts[0]:
-        gr.Warning("NSFW content detected.")
-        return None
+    # AiTube aims for real time, but we are loosing FPS if we perform this step
+    #has_nsfw_concepts = check_nsfw_images([output.frames[0][0]])
+    #if has_nsfw_concepts[0]:
+    #    gr.Warning("NSFW content detected.")
+    #    return None
 
     name = str(uuid.uuid4()).replace("-", "")
     path = f"/tmp/{name}.mp4"
+    
+    # I think we are looking time here too, converting to mp4 is too slow, we should return
+    # the frames unencoded to the frontend renderer
     export_to_video(output.frames[0], path, fps=10)
+    
     return path
 
 
 # Gradio Interface
-with gr.Blocks(css="style.css") as demo:
-    gr.HTML(
-        "<h1><center>AnimateDiff-Lightning ⚡</center></h1>" +
-        "<p><center>Lightning-fast text-to-video generation</center></p>" +
-        "<p><center><a href='https://huggingface.co/ByteDance/AnimateDiff-Lightning'>https://huggingface.co/ByteDance/AnimateDiff-Lightning</a></center></p>"
-    )
+with gr.Blocks() as demo:
+
     with gr.Group():
         with gr.Row():
             prompt = gr.Textbox(
-                label='Prompt (English)'
+                label='Prompt'
             )
         with gr.Row():
             select_base = gr.Dropdown(
@@ -135,7 +138,7 @@ with gr.Blocks(css="style.css") as demo:
                     ('2-Step', 2),
                     ('4-Step', 4),
                     ('8-Step', 8)],
-                value=4,
+                value=2,
                 interactive=True
             )
             submit = gr.Button(
@@ -145,8 +148,8 @@ with gr.Blocks(css="style.css") as demo:
     video = gr.Video(
         label='AnimateDiff-Lightning',
         autoplay=True,
-        height=1024,
-        width=576,
+        height=512,
+        width=910,
         elem_id="video_output"
     )