Update app.py

app.py

@@ -7,17 +7,29 @@ import gradio as gr
 from huggingface_hub import hf_hub_download
 import spaces
 from comfy import model_management
-from nodes import NODE_CLASS_MAPPINGS
 
-sys.path.append(os.path.abspath(os.path.dirname(__file__)))
+from huggingface_hub import hf_hub_download
 
-print("Current working directory:", os.getcwd())
-print("Absolute path of app.py:", os.path.abspath(os.path.dirname(__file__)))
 
 hf_hub_download(
-    repo_id="Comfy-Org/stable-diffusion-v1-5-archive",
-    filename="v1-5-pruned-emaonly-fp16.safetensors",
-    local_dir="models/checkpoints"
+    repo_id="Madespace/clip", 
+    filename="google_t5-v1_1-xxl_encoderonly-fp8_e4m3fn.safetensors", 
+    local_dir="models/clip"              
+)
+hf_hub_download(
+    repo_id="ezioruan/inswapper_128.onnx",        
+    filename="inswapper_128.onnx",       
+    local_dir="models/insightface"              
+)
+hf_hub_download(
+    repo_id="gmk123/GFPGAN",           
+    filename="GFPGANv1.4.pth",           
+    local_dir="models/facerestore_models" 
+)
+hf_hub_download(
+    repo_id="gemasai/4x_NMKD-Superscale-SP_178000_G",          
+    filename="4x_NMKD-Superscale-SP_178000_G.pth",  
+    local_dir="models/upscale_models" 
 )
 
 
@@ -89,18 +101,18 @@ def add_extra_model_paths() -> None:
         from main import load_extra_path_config
     except ImportError:
         print(
-            "Could not import load_extra_path_config from main.py. Looking in util.extra_config instead."
+            "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
         )
         from ut.extra_config import load_extra_path_config
-        # sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "ut")))
-        # from extra_config import load_extra_path_config
 
     extra_model_paths = find_path("extra_model_paths.yaml")
+
     if extra_model_paths is not None:
         load_extra_path_config(extra_model_paths)
     else:
         print("Could not find the extra_model_paths config file.")
 
+
 add_comfyui_directory_to_sys_path()
 add_extra_model_paths()
 
@@ -127,86 +139,193 @@ def import_custom_nodes() -> None:
     # Initializing custom nodes
     init_extra_nodes()
 
-import_custom_nodes()
+from nodes import NODE_CLASS_MAPPINGS
+
+#TO be added to "model_loaders" as it loads a model
+downloadandloadcogvideomodel = NODE_CLASS_MAPPINGS[
+    "DownloadAndLoadCogVideoModel"
+]()
+downloadandloadcogvideomodel_1 = downloadandloadcogvideomodel.loadmodel(
+    model="THUDM/CogVideoX-5b",
+    precision="bf16",
+    quantization="disabled",
+    enable_sequential_cpu_offload=True,
+    attention_mode="sdpa",
+    load_device="main_device",
+)
+loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+cliploader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
+cliploader_20 = cliploader.load_clip(
+    clip_name="t5/google_t5-v1_1-xxl_encoderonly-fp8_e4m3fn.safetensors",
+    type="sd3",
+    device="default",
+)
+emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
+
+cogvideotextencode = NODE_CLASS_MAPPINGS["CogVideoTextEncode"]()
+cogvideosampler = NODE_CLASS_MAPPINGS["CogVideoSampler"]()
+cogvideodecode = NODE_CLASS_MAPPINGS["CogVideoDecode"]()
+reactorfaceswap = NODE_CLASS_MAPPINGS["ReActorFaceSwap"]()
+cr_upscale_image = NODE_CLASS_MAPPINGS["CR Upscale Image"]()
+vhs_videocombine = NODE_CLASS_MAPPINGS["VHS_VideoCombine"]()
+
+#Add all the models that load a safetensors file
+model_loaders = [downloadandloadcogvideomodel_1, cliploader_20]
+
+# Check which models are valid and how to best load them
+valid_models = [
+    getattr(loader[0], 'patcher', loader[0]) 
+    for loader in model_loaders
+    if not isinstance(loader[0], dict) and not isinstance(getattr(loader[0], 'patcher', None), dict)
+]
+
+#Finally loads the models
+model_management.load_models_gpu(valid_models)
 
+#Run ComfyUI Workflow
 @spaces.GPU(duration=60)
-def generate_image(positive_prompt, negative_prompt):
+def generate_video(positive_prompt, num_frames, input_image):
+    
+    print("Positive Prompt:", positive_prompt)
+    print("Number of Frames:", num_frames)
+    print("Input Image:", input_image)
+
     import_custom_nodes()
     with torch.inference_mode():
-        checkpointloadersimple = NODE_CLASS_MAPPINGS["CheckpointLoaderSimple"]()
-        checkpointloadersimple_4 = checkpointloadersimple.load_checkpoint(
-            ckpt_name="v1-5-pruned-emaonly-fp16.safetensors"
+        # downloadandloadcogvideomodel = NODE_CLASS_MAPPINGS[
+        #     "DownloadAndLoadCogVideoModel"
+        # ]()
+        # downloadandloadcogvideomodel_1 = downloadandloadcogvideomodel.loadmodel(
+        #     model="THUDM/CogVideoX-5b",
+        #     precision="bf16",
+        #     quantization="disabled",
+        #     enable_sequential_cpu_offload=True,
+        #     attention_mode="sdpa",
+        #     load_device="main_device",
+        # )
+
+        # loadimage = NODE_CLASS_MAPPINGS["LoadImage"]()
+        loadimage_8 = loadimage.load_image(image=input_image)
+
+        # cliploader = NODE_CLASS_MAPPINGS["CLIPLoader"]()
+        # cliploader_20 = cliploader.load_clip(
+        #     clip_name="t5/google_t5-v1_1-xxl_encoderonly-fp8_e4m3fn.safetensors",
+        #     type="sd3",
+        #     device="default",
+        # )
+
+        # emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
+        emptylatentimage_161 = emptylatentimage.generate(
+            width=720, height=480, batch_size=1
         )
 
-        emptylatentimage = NODE_CLASS_MAPPINGS["EmptyLatentImage"]()
-        emptylatentimage_5 = emptylatentimage.generate(
-            width=512, height=512, batch_size=1
-        )
+        # cogvideotextencode = NODE_CLASS_MAPPINGS["CogVideoTextEncode"]()
+        # cogvideosampler = NODE_CLASS_MAPPINGS["CogVideoSampler"]()
+        # cogvideodecode = NODE_CLASS_MAPPINGS["CogVideoDecode"]()
+        # reactorfaceswap = NODE_CLASS_MAPPINGS["ReActorFaceSwap"]()
+        # cr_upscale_image = NODE_CLASS_MAPPINGS["CR Upscale Image"]()
+        # vhs_videocombine = NODE_CLASS_MAPPINGS["VHS_VideoCombine"]()
 
-        cliptextencode = NODE_CLASS_MAPPINGS["CLIPTextEncode"]()
-        cliptextencode_6 = cliptextencode.encode(
-            text=positive_prompt,
-            clip=get_value_at_index(checkpointloadersimple_4, 1),
-        )
-
-        cliptextencode_7 = cliptextencode.encode(
-            text=negative_prompt, clip=get_value_at_index(checkpointloadersimple_4, 1)
-        )
+        for q in range(1):
+            cogvideotextencode_30 = cogvideotextencode.process(
+                prompt=positive_prompt,
+                strength=1,
+                force_offload=True,
+                clip=get_value_at_index(cliploader_20, 0),
+            )
 
-        ksampler = NODE_CLASS_MAPPINGS["KSampler"]()
-        vaedecode = NODE_CLASS_MAPPINGS["VAEDecode"]()
-        saveimage = NODE_CLASS_MAPPINGS["SaveImage"]()
+            cogvideotextencode_31 = cogvideotextencode.process(
+                prompt='',
+                strength=1,
+                force_offload=True,
+                clip=get_value_at_index(cogvideotextencode_30, 1),
+            )
 
-        for q in range(1):
-            ksampler_3 = ksampler.sample(
+            cogvideosampler_155 = cogvideosampler.process(
+                num_frames=num_frames,
+                steps=50,
+                cfg=6,
                 seed=random.randint(1, 2**64),
-                steps=20,
-                cfg=8,
-                sampler_name="euler",
-                scheduler="normal",
-                denoise=1,
-                model=get_value_at_index(checkpointloadersimple_4, 0),
-                positive=get_value_at_index(cliptextencode_6, 0),
-                negative=get_value_at_index(cliptextencode_7, 0),
-                latent_image=get_value_at_index(emptylatentimage_5, 0),
+                scheduler="CogVideoXDDIM",
+                denoise_strength=1,
+                model=get_value_at_index(downloadandloadcogvideomodel_1, 0),
+                positive=get_value_at_index(cogvideotextencode_30, 0),
+                negative=get_value_at_index(cogvideotextencode_31, 0),
+                samples=get_value_at_index(emptylatentimage_161, 0),
             )
 
-            vaedecode_8 = vaedecode.decode(
-                samples=get_value_at_index(ksampler_3, 0),
-                vae=get_value_at_index(checkpointloadersimple_4, 2),
+            cogvideodecode_11 = cogvideodecode.decode(
+                enable_vae_tiling=False,
+                tile_sample_min_height=240,
+                tile_sample_min_width=360,
+                tile_overlap_factor_height=0.2,
+                tile_overlap_factor_width=0.2,
+                auto_tile_size=True,
+                vae=get_value_at_index(downloadandloadcogvideomodel_1, 1),
+                samples=get_value_at_index(cogvideosampler_155, 0),
             )
 
-            saveimage_9 = saveimage.save_images(
-                filename_prefix="ComfyUI", 
-                images=get_value_at_index(vaedecode_8, 0),
+            reactorfaceswap_3 = reactorfaceswap.execute(
+                enabled=True,
+                swap_model="inswapper_128.onnx",
+                facedetection="retinaface_resnet50",
+                face_restore_model="GFPGANv1.4.pth",
+                face_restore_visibility=1,
+                codeformer_weight=0.75,
+                detect_gender_input="no",
+                detect_gender_source="no",
+                input_faces_index="0",
+                source_faces_index="0",
+                console_log_level=1,
+                input_image=get_value_at_index(cogvideodecode_11, 0),
+                source_image=get_value_at_index(loadimage_8, 0),
             )
-            saved_path=f"output/{saveimage_9['ui']['images'][0]['filename']}"
+
+            cr_upscale_image_151 = cr_upscale_image.upscale(
+                upscale_model="4x_NMKD-Superscale-SP_178000_G.pth",
+                mode="rescale",
+                rescale_factor=4,
+                resize_width=720,
+                resampling_method="lanczos",
+                supersample="true",
+                rounding_modulus=16,
+                image=get_value_at_index(reactorfaceswap_3, 0),
+            )
+
+            vhs_videocombine_154 = vhs_videocombine.combine_video(
+                frame_rate=8,
+                loop_count=0,
+                filename_prefix="AnimateDiff",
+                format="video/h264-mp4",
+                pix_fmt="yuv420p",
+                crf=19,
+                save_metadata=True,
+                trim_to_audio=False,
+                pingpong=True,
+                save_output=True,
+                images=get_value_at_index(cr_upscale_image_151, 0),
+                unique_id=7214086815220268849,
+            )
+            saved_path = f"output/{vhs_videocombine_154['ui']['images'][0]['filename']}"
             return saved_path
-            
-# Start your Gradio app
-with gr.Blocks() as app:
-    # Add a title
-    gr.Markdown("# FLUX Style Shaping")
-
-    with gr.Row():
-        with gr.Column():
-            # Add an input
-            positive_prompt_input = gr.Textbox(label="Positive Prompt", placeholder="Enter your positive prompt here...")
-            negative_prompt_input = gr.Textbox(label="Negative Prompt", placeholder="Enter your negative prompt here...")
-            # The generate button
-            generate_btn = gr.Button("Generate")
+
+
+if __name__ == "__main__":
+
+    with gr.Blocks() as app:
+        with gr.Row():
+            positive_prompt = gr.Textbox(label="Positive Prompt", value="", lines=2)
+        with gr.Row():
+            num_frames = gr.Number(label="Number of Frames", value=60)
+        with gr.Row():
+            input_image = gr.Image(label="Input Image", type="pil")
+        submit = gr.Button("Submit")
+        output_video = gr.Video(label="Output Video")
         
-        with gr.Column():
-            # The output image
-            output_image = gr.Image(label="Generated Image")
-
-        # When clicking the button, it will trigger the `generate_image` function, with the respective inputs
-        # and the output an image
-        generate_btn.click(
-            fn=generate_image,
-            inputs=[positive_prompt_input, negative_prompt_input],
-            outputs=[output_image]
+        submit.click(
+            fn=generate_video, 
+            inputs=[positive_prompt, num_frames, input_image],
+            outputs=[output_video]
         )
 
-if __name__ == "__main__":
-        app.launch(share=True)
+    app.launch(share=True)