support text to image

app.py

@@ -5,7 +5,6 @@ import gradio as gr
 import tempfile
 import spaces
 import numpy as np
-from PIL import Image
 import random
 
 MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers"
@@ -63,20 +62,23 @@ def handle_image_upload_for_dims_wan(uploaded_pil_image, current_h_val, current_
     except Exception as e:
         gr.Warning("Error attempting to calculate new dimensions")
         return gr.update(value=DEFAULT_H_SLIDER_VALUE), gr.update(value=DEFAULT_W_SLIDER_VALUE)
-        
-def get_duration(input_image, prompt, height, width, 
-                   negative_prompt, duration_seconds,
-                   guidance_scale, steps,
-                   seed, randomize_seed, 
-                   progress):
+
+def get_duration_video(input_image, prompt, height, width, 
+                     negative_prompt, duration_seconds,
+                     guidance_scale, steps,
+                     seed, randomize_seed, 
+                     progress):
     if steps > 4 and duration_seconds > 4:
         return 90
     elif steps > 4 or duration_seconds > 4:
         return 75
     else:
         return 60
+        
+def get_duration_image(prompt, height, width, negative_prompt, guidance_scale, steps, seed, randomize_seed, progress):
+    return 30 if steps > 4 else 20
 
-@spaces.GPU(duration=get_duration)
+@spaces.GPU(duration=get_duration_video)
 def generate_video(input_image, prompt, height, width, negative_prompt=default_negative_prompt, duration_seconds=2, guidance_scale=0, steps=4, seed=44, randomize_seed=False, progress=gr.Progress(track_tqdm=True)):
     target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
     target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
@@ -108,58 +110,80 @@ def generate_video(input_image, prompt, height, width, negative_prompt=default_n
     export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
     return video_path, current_seed
 
+@spaces.GPU(duration=get_duration_image)
+def generate_image(prompt, height, width, negative_prompt=default_negative_prompt, guidance_scale=0, steps=4, seed=44, randomize_seed=False, progress=gr.Progress(track_tqdm=True)):
+    """Generates a single image using the text-to-video pipeline by requesting only one frame."""
+    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
+    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)
+
+    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+
+    with torch.inference_mode():
+        output_frame = text_to_video_pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=target_h,
+            width=target_w,
+            num_frames=1, 
+            guidance_scale=float(guidance_scale),
+            num_inference_steps=int(steps),
+            generator=torch.Generator(device="cuda").manual_seed(current_seed)
+        ).frames[0][0]  
+
+    return output_frame, current_seed
+
 with gr.Blocks() as demo:
     gr.Markdown("# Fast Wan 2.2 TI2V 5B Demo")
     gr.Markdown("""This Demo is using [FastWan2.2-TI2V-5B](https://huggingface.co/FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers) which is fine-tuned with Sparse-distill method which allows wan to generate high quality videos in 3-5 steps.""")
 
-    with gr.Row():
-        with gr.Column():
-            input_image_component = gr.Image(type="pil", label="Input Image (optional, auto-resized to target H/W)")
-            prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
-            duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), step=0.1, value=2, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
-
-            with gr.Accordion("Advanced Settings", open=False):
-                negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
-                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
-                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
-                with gr.Row():
-                    height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
-                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
-                steps_slider = gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Inference Steps")
-                guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.0, step=0.01, value=0.0, label="Guidance Scale")
-            generate_button = gr.Button("Generate Video", variant="primary")
-        with gr.Column():
-            video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
-
-    input_image_component.upload(
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
-
-    input_image_component.clear(
-        fn=handle_image_upload_for_dims_wan,
-        inputs=[input_image_component, height_input, width_input],
-        outputs=[height_input, width_input]
-    )
-
-    ui_inputs = [
-        input_image_component, prompt_input, height_input, width_input,
-        negative_prompt_input, duration_seconds_input,
-        guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox
-    ]
-    generate_button.click(fn=generate_video, inputs=ui_inputs, outputs=[video_output, seed_input])
-
-    gr.Examples(
-        examples=[ 
-            [None, "A person eating spaghetti", 1024, 720],
-            ["cat.png", "The cat removes the glasses from its eyes.", 1088, 800],
-            [None, "a penguin playfully dancing in the snow, Antarctica", 1024, 720],
-            ["peng.png", "a penguin running towards camera joyfully, Antarctica", 896, 512],
-        ],
-        
-        inputs=[input_image_component, prompt_input, height_input, width_input], outputs=[video_output, seed_input], fn=generate_video, cache_examples="lazy"
-    )
+    with gr.Tabs():
+        with gr.TabItem("Text/Image-to-Video"):
+            with gr.Row():
+                with gr.Column():
+                    input_image_component = gr.Image(type="pil", label="Input Image (optional, auto-resized to target H/W)")
+                    prompt_input = gr.Textbox(label="Prompt", value=default_prompt_i2v)
+                    duration_seconds_input = gr.Slider(minimum=round(MIN_FRAMES_MODEL/FIXED_FPS,1), maximum=round(MAX_FRAMES_MODEL/FIXED_FPS,1), step=0.1, value=2, label="Duration (seconds)", info=f"Clamped to model's {MIN_FRAMES_MODEL}-{MAX_FRAMES_MODEL} frames at {FIXED_FPS}fps.")
+                      
+                    with gr.Accordion("Advanced Settings", open=False):
+                        negative_prompt_input = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
+                        seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
+                        randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
+                        with gr.Row():
+                            height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
+                            width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
+                        steps_slider = gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Inference Steps")
+                        guidance_scale_input = gr.Slider(minimum=0.0, maximum=5.0, step=0.01, value=0.0, label="Guidance Scale")
+                    
+                    generate_button = gr.Button("Generate Video", variant="primary")
+                
+                with gr.Column():
+                    video_output = gr.Video(label="Generated Video", autoplay=True, interactive=False)
+
+            input_image_component.upload(fn=handle_image_upload_for_dims_wan, inputs=[input_image_component, height_input, width_input], outputs=[height_input, width_input])
+            input_image_component.clear(fn=handle_image_upload_for_dims_wan, inputs=[input_image_component, height_input, width_input], outputs=[height_input, width_input])
+            
+            ui_inputs_video = [input_image_component, prompt_input, height_input, width_input, negative_prompt_input, duration_seconds_input, guidance_scale_input, steps_slider, seed_input, randomize_seed_checkbox]
+            generate_button.click(fn=generate_video, inputs=ui_inputs_video, outputs=[video_output, seed_input])
+
+        with gr.TabItem("Text-to-Image"):
+            with gr.Row():
+                with gr.Column():
+                    prompt_input_img = gr.Textbox(label="Prompt", value="A majestic lion in the savanna, cinematic lighting, 4k, high detail")
+                    with gr.Accordion("Advanced Settings", open=False):
+                        negative_prompt_input_img = gr.Textbox(label="Negative Prompt", value=default_negative_prompt, lines=3)
+                        seed_input_img = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42, interactive=True)
+                        randomize_seed_checkbox_img = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
+                        with gr.Row():
+                            height_input_img = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE, value=DEFAULT_H_SLIDER_VALUE, label=f"Output Height (multiple of {MOD_VALUE})")
+                            width_input_img = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE, value=DEFAULT_W_SLIDER_VALUE, label=f"Output Width (multiple of {MOD_VALUE})")
+                        steps_slider_img = gr.Slider(minimum=1, maximum=8, step=1, value=4, label="Inference Steps")
+                        guidance_scale_input_img = gr.Slider(minimum=0.0, maximum=5.0, step=0.01, value=0.0, label="Guidance Scale")
+                    generate_button_img = gr.Button("Generate Image", variant="primary")
+                with gr.Column():
+                    image_output = gr.Image(label="Generated Image", interactive=False)
+
+            ui_inputs_img = [prompt_input_img, height_input_img, width_input_img, negative_prompt_input_img, guidance_scale_input_img, steps_slider_img, seed_input_img, randomize_seed_checkbox_img]
+            generate_button_img.click(fn=generate_image, inputs=ui_inputs_img, outputs=[image_output, seed_input_img])
 
 if __name__ == "__main__":
     demo.queue().launch()