added custom model option

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import os
+import shutil
 import subprocess
 from share_btn import community_icon_html, loading_icon_html, share_js
 import cv2
@@ -22,7 +23,17 @@ for model_id in model_ids:
     model_name = model_id.split('/')[-1]
     snapshot_download(model_id, local_dir=f'checkpoints/{model_name}')
 
+def load_model(model_id):
+    local_dir = f'checkpoints/stable-diffusion-v1-5'
+    # Check if the directory exists
+    if os.path.exists(local_dir):
+        # Delete the directory if it exists
+        shutil.rmtree(local_dir)
 
+    model_name = model_id.split('/')[-1]
+    snapshot_download(model_id, local_dir=f'checkpoints/{model_name}')
+    os.rename(f'checkpoints/{model_name}', f'checkpoints/stable-diffusion-v1-5')
+    return "model loaded"
 
 def get_frame_count(filepath):
     if filepath is not None:
@@ -259,38 +270,45 @@ with gr.Blocks(css=css) as demo:
             
             [![Duplicate this Space](https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-sm.svg#center)](https://huggingface.co/spaces/fffiloni/ControlVideo?duplicate=true)
         """)
-        with gr.Row():
-            with gr.Column():
-                #video_in = gr.Video(source="upload", type="filepath", visible=True)
+        
+        with gr.Column():
+            with gr.Row():
                 video_path = gr.Video(label="Input video", source="upload", type="filepath", visible=True, elem_id="video-in")
-                prompt = gr.Textbox(label="prompt", elem_id="prompt-in")
                 with gr.Column():
-                    video_length = gr.Slider(label="Video length", info="How many frames do you want to process ? For demo purpose, max is set to 24", minimum=1, maximum=12, step=1, value=2)
-                    with gr.Row():
-                        condition = gr.Dropdown(label="Condition", choices=["depth", "canny", "pose"], value="depth")
-                        seed = gr.Number(label="seed", value=42)
-                submit_btn = gr.Button("Submit")
+                    video_res = gr.Video(label="result", elem_id="video-out")
+                    with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
+                        community_icon = gr.HTML(community_icon_html)
+                        loading_icon = gr.HTML(loading_icon_html)
+                        share_button = gr.Button("Share to community", elem_id="share-btn")
+            with gr.Row():
+                    chosen_model = gr.Textbox(label="Custom model (*1.5)", placeholder="E.g: nitrosocke/Ghibli-Diffusion")
+                    model_status = gr.Textbox(label="status")
+                    load_model_btn = gr.Button("load model (optional)")
+            prompt = gr.Textbox(label="prompt", info="If you loaded a custom model, do not forget to include Prompt trigger", elem_id="prompt-in")
             with gr.Column():
-                video_res = gr.Video(label="result", elem_id="video-out")
-                with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
-                    community_icon = gr.HTML(community_icon_html)
-                    loading_icon = gr.HTML(loading_icon_html)
-                    share_button = gr.Button("Share to community", elem_id="share-btn")
-                gr.Examples(
-                    examples=[["Indiana Jones moonwalk in the wild jungle", "./examples/moonwalk.mp4", 'depth', 24, 192837465]],
-                    fn=run_inference,
-                    inputs=[prompt,
-                             video_path,
-                             condition,
-                             video_length,
-                             seed
-                            ],
-                    outputs=[video_res, share_group],
-                    cache_examples=True
-                )
+                video_length = gr.Slider(label="Video length", info="How many frames do you want to process ? For demo purpose, max is set to 24", minimum=1, maximum=12, step=1, value=2)
+                with gr.Row():
+                    condition = gr.Dropdown(label="Condition", choices=["depth", "canny", "pose"], value="depth")
+                    seed = gr.Number(label="seed", value=42)
+            submit_btn = gr.Button("Submit")
+        
+            
+            
+            gr.Examples(
+                examples=[["Indiana Jones moonwalk in the wild jungle", "./examples/moonwalk.mp4", 'depth', 24, 192837465]],
+                fn=run_inference,
+                inputs=[prompt,
+                         video_path,
+                         condition,
+                         video_length,
+                         seed
+                        ],
+                outputs=[video_res, share_group],
+                cache_examples=False
+            )
                 
     share_button.click(None, [], [], _js=share_js)
-    
+    load_model_btn.click(fn=load_model, inputs=[chosen_model], outputs=[model_status], queue=False)
     video_path.change(fn=get_frame_count,
                       inputs=[video_path],
                       outputs=[video_length],