update

app.py

@@ -13,12 +13,12 @@ from wan_pipeline import WanPipeline
 from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
 from PIL import Image
 from diffusers.utils import export_to_video
-
-
 from huggingface_hub import login
-login(token=os.getenv('HF_TOKEN'))
 
+# Authenticate with HF
+login(token=os.getenv('HF_TOKEN'))
 
+# Set seed
 def set_seed(seed):
     random.seed(seed)
     os.environ['PYTHONHASHSEED'] = str(seed)
@@ -33,20 +33,17 @@ model_paths = {
     "wan-t2v": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 }
 
-# Global variable for current model
 current_model = None
-
-# Folder to save video outputs
 OUTPUT_DIR = "generated_videos"
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 
 def load_model(model_name):
     global current_model
     if current_model is not None:
-        del current_model  # Delete the old model
-        torch.cuda.empty_cache()  # Free GPU memory
-        gc.collect()  # Force garbage collection
-    
+        del current_model
+        torch.cuda.empty_cache()
+        gc.collect()
+
     if "wan-t2v" in model_name:
         vae = AutoencoderKLWan.from_pretrained(model_paths[model_name], subfolder="vae", torch_dtype=torch.bfloat16)
         scheduler = UniPCMultistepScheduler(prediction_type='flow_prediction', use_flow_sigmas=True, num_train_timesteps=1000, flow_shift=8.0)
@@ -54,9 +51,8 @@ def load_model(model_name):
         current_model.scheduler = scheduler
     else:
         current_model = StableDiffusion3Pipeline.from_pretrained(model_paths[model_name], torch_dtype=torch.bfloat16).to("cuda")
-    
-    return current_model.to('cuda')
 
+    return current_model.to("cuda")
 
 @spaces.GPU(duration=500)
 def generate_content(prompt, model_name, guidance_scale=7.5, num_inference_steps=50, use_cfg_zero_star=True, use_zero_init=True, zero_steps=0, seed=None, compare_mode=False):
@@ -68,52 +64,26 @@ def generate_content(prompt, model_name, guidance_scale=7.5, num_inference_steps
     is_video_model = "wan-t2v" in model_name
 
     if is_video_model:
-        if True:
-            negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
-            set_seed(seed)
-            video1_frames = model(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                height=480,
-                width=832,
-                num_frames=81,
-                num_inference_steps=num_inference_steps,
-                guidance_scale=guidance_scale,
-                use_cfg_zero_star=True,
-                use_zero_init=True,
-                zero_steps=0
-            ).frames[0]
-            video1_path = os.path.join(OUTPUT_DIR, f"{seed}_CFG-Zero-Star.mp4")
-            export_to_video(video1_frames, video1_path, fps=16)
-
-            return None, None, video1_path, seed
-
-        
-            # set_seed(seed)
-            # video2_frames = model(
-            #     prompt=prompt,
-            #     guidance_scale=guidance_scale,
-            #     num_frames=81,
-            #     use_cfg_zero_star=False,
-            #     use_zero_init=use_zero_init,
-            #     zero_steps=zero_steps
-            # ).frames[0]
-            # video2_path = os.path.join(OUTPUT_DIR,  f"{seed}_CFG.mp4")
-            # export_to_video(video2_frames, video2_path, fps=16)
-
-            # return None, None, video1_path, video2_path, seed
-        # else:
-        #     video_frames = model(
-        #         prompt=prompt,
-        #         guidance_scale=guidance_scale,
-        #         num_frames=81,
-        #         use_cfg_zero_star=use_cfg_zero_star,
-        #         use_zero_init=use_zero_init,
-        #         zero_steps=zero_steps
-        #     ).frames[0]
-        #     video_path = save_video(video_frames, f"{seed}.mp4")
-        #     return None, None, video_path, None, seed
-    print('prompt: ',prompt)
+        negative_prompt = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
+        video1_frames = model(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=480,
+            width=832,
+            num_frames=81,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            use_cfg_zero_star=True,
+            use_zero_init=True,
+            zero_steps=0
+        ).frames[0]
+        video1_path = os.path.join(OUTPUT_DIR, f"{seed}_CFG-Zero-Star.mp4")
+        export_to_video(video1_frames, video1_path, fps=16)
+
+        return None, None, video1_path, seed
+
+    print("prompt:", prompt)
+
     if compare_mode:
         set_seed(seed)
         image1 = model(
@@ -134,8 +104,8 @@ def generate_content(prompt, model_name, guidance_scale=7.5, num_inference_steps
             use_zero_init=use_zero_init,
             zero_steps=zero_steps
         ).images[0]
+
         return image1, image2, None, seed
-        #return image1, image2, None, None, seed
     else:
         image = model(
             prompt,
@@ -145,14 +115,11 @@ def generate_content(prompt, model_name, guidance_scale=7.5, num_inference_steps
             use_zero_init=use_zero_init,
             zero_steps=zero_steps
         ).images[0]
+
         if use_cfg_zero_star:
             return image, None, None, seed
         else:
             return None, image, None, seed
-        # if use_cfg_zero_star:
-        #     return image, None, None, None, seed
-        # else:
-        #     return None, image, None, None, seed
 
 # Gradio UI
 with gr.Blocks() as demo:
@@ -166,28 +133,63 @@ with gr.Blocks() as demo:
         </div>
     """)
 
-    gr.Interface(
+    with gr.Row():
+        prompt = gr.Textbox(value="A spooky haunted mansion on a hill silhouetted by a full moon.", label="Enter your prompt")
+        model_choice = gr.Dropdown(choices=list(model_paths.keys()), label="Choose Model")
+
+    with gr.Row():
+        guidance_scale = gr.Slider(1, 20, value=4.0, step=0.5, label="Guidance Scale")
+        inference_steps = gr.Slider(10, 100, value=28, step=5, label="Inference Steps")
+
+    with gr.Row():
+        use_opt_scale = gr.Checkbox(value=True, label="Use Optimized-Scale")
+        use_zero_init = gr.Checkbox(value=True, label="Use Zero Init")
+        zero_steps = gr.Slider(0, 20, value=0, step=1, label="Zero out steps")
+
+    with gr.Row():
+        seed = gr.Number(value=42, label="Seed (Leave blank for random)")
+        compare_mode = gr.Checkbox(value=True, label="Compare Mode")
+
+    with gr.Row():
+        out1 = gr.Image(type="pil", label="CFG-Zero* Image")
+        out2 = gr.Image(type="pil", label="CFG Image")
+        video = gr.Video(label="Video")
+        used_seed = gr.Textbox(label="Used Seed")
+
+    generate_btn = gr.Button("Generate")
+
+    # Change logic for when "wan-t2v" is selected
+    def update_params(model_name):
+        if model_name == "wan-t2v":
+            return (
+                gr.update(value=5),
+                gr.update(value=50),
+                gr.update(value=True),
+                gr.update(value=True),
+                gr.update(value=1)
+            )
+        else:
+            return (
+                gr.update(value=4.0),
+                gr.update(value=28),
+                gr.update(value=True),
+                gr.update(value=True),
+                gr.update(value=0)
+            )
+
+    model_choice.change(
+        fn=update_params,
+        inputs=[model_choice],
+        outputs=[guidance_scale, inference_steps, use_opt_scale, use_zero_init, zero_steps]
+    )
+
+    generate_btn.click(
         fn=generate_content,
         inputs=[
-            gr.Textbox(value="A spooky haunted mansion on a hill silhouetted by a full moon.", label="Enter your prompt"),
-            gr.Dropdown(choices=list(model_paths.keys()), label="Choose Model"),
-            gr.Slider(1, 20, value=4.0, step=0.5, label="Guidance Scale"),
-            gr.Slider(10, 100, value=28, step=5, label="Inference Steps"),
-            gr.Checkbox(value=True, label="Use Optimized-Scale"),
-            gr.Checkbox(value=True, label="Use Zero Init"),
-            gr.Slider(0, 20, value=0, step=1, label="Zero out steps"),
-            gr.Number(value=42, label="Seed (Leave blank for random)"),
-            gr.Checkbox(value=True, label="Compare Mode")
+            prompt, model_choice, guidance_scale, inference_steps,
+            use_opt_scale, use_zero_init, zero_steps, seed, compare_mode
         ],
-        outputs=[
-            gr.Image(type="pil", label="CFG-Zero* Image"),
-            gr.Image(type="pil", label="CFG Image"),
-            gr.Video(label="Video"),
-            gr.Textbox(label="Used Seed")
-        ],
-        #title="CFG-Zero*: Improved Classifier-Free Guidance for Flow Matching Models",
-        live=False  # optional
+        outputs=[out1, out2, video, used_seed]
     )
 
 demo.launch(ssr_mode=False)
-