diff --git a/example_workflows/image_generation.json b/example_workflows/image_generation.json index b82390c..b397b25 100644 --- a/example_workflows/image_generation.json +++ b/example_workflows/image_generation.json @@ -1 +1,206 @@ -{"id":"4a7ae110-38ab-446e-97da-b826e950c065","revision":0,"last_node_id":85,"last_link_id":203,"nodes":[{"id":67,"type":"GetHeadPose","pos":[814.1528930664062,495.9793701171875],"size":[270,98],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":147},{"localized_name":"images","name":"images","type":"IMAGE","link":151},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":143},{"localized_name":"crop","name":"crop","type":"BOOLEAN","widget":{"name":"crop"},"link":null}],"outputs":[{"localized_name":"head_pose","name":"head_pose","type":"HEAD_POSE","links":[191]}],"properties":{"Node name for S&R":"GetHeadPose"},"widgets_values":[false]},{"id":69,"type":"GetHeadPose","pos":[819.799560546875,860.5927734375],"size":[270,98],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":149},{"localized_name":"images","name":"images","type":"IMAGE","link":188},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":145},{"localized_name":"crop","name":"crop","type":"BOOLEAN","widget":{"name":"crop"},"link":null}],"outputs":[{"localized_name":"head_pose","name":"head_pose","type":"HEAD_POSE","links":[193]}],"properties":{"Node name for S&R":"GetHeadPose"},"widgets_values":[true]},{"id":80,"type":"HMPipelineImage","pos":[1275.33935546875,652.6725463867188],"size":[270,306],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"hm_image_pipeline","name":"hm_image_pipeline","type":"HMIMAGEPIPELINE","link":190},{"localized_name":"ref_head_pose","name":"ref_head_pose","type":"HEAD_POSE","link":191},{"localized_name":"ref_expression","name":"ref_expression","type":"EXPRESSION","link":196},{"localized_name":"drive_head_pose","name":"drive_head_pose","type":"HEAD_POSE","link":193},{"localized_name":"drive_expression","name":"drive_expression","type":"EXPRESSION","link":197},{"localized_name":"trans_ratio","name":"trans_ratio","type":"FLOAT","widget":{"name":"trans_ratio"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"steps","name":"steps","type":"INT","widget":{"name":"steps"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"gpu_id","name":"gpu_id","type":"INT","widget":{"name":"gpu_id"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[195]},{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":null}],"properties":{"Node name for S&R":"HMPipelineImage"},"widgets_values":[0,"(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture","",25,1018802962613971,"randomize",2,0]},{"id":82,"type":"HMImagePipelineLoader","pos":[814.7327880859375,202.5532989501953],"size":[298.3472595214844,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"checkpoint","name":"checkpoint","shape":7,"type":"COMBO","widget":{"name":"checkpoint"},"link":null},{"localized_name":"lora","name":"lora","shape":7,"type":"COMBO","widget":{"name":"lora"},"link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"COMBO","widget":{"name":"vae"},"link":null},{"localized_name":"version","name":"version","shape":7,"type":"COMBO","widget":{"name":"version"},"link":null},{"localized_name":"stylize","name":"stylize","shape":7,"type":"COMBO","widget":{"name":"stylize"},"link":null},{"localized_name":"deployment","name":"deployment","shape":7,"type":"COMBO","widget":{"name":"deployment"},"link":null},{"localized_name":"lora_scale","name":"lora_scale","shape":7,"type":"FLOAT","widget":{"name":"lora_scale"},"link":null}],"outputs":[{"localized_name":"hm_image_pipeline","name":"hm_image_pipeline","type":"HMIMAGEPIPELINE","links":[190]}],"properties":{"Node name for S&R":"HMImagePipelineLoader"},"widgets_values":["[preset]DisneyPixarCartoonB","None","same as checkpoint","v5","x1","huggingface",1]},{"id":83,"type":"PreviewImage","pos":[1689.1593017578125,650.2534790039062],"size":[210,258],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":195}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":84,"type":"GetExpression2","pos":[824.4129638671875,694.6199951171875],"size":[191.63729858398438,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":201},{"localized_name":"images","name":"images","type":"IMAGE","link":202},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":198}],"outputs":[{"localized_name":"expression","name":"expression","type":"EXPRESSION","links":[196]}],"properties":{"Node name for S&R":"GetExpression2"},"widgets_values":[]},{"id":85,"type":"GetExpression2","pos":[826.0261840820312,1060.0404052734375],"size":[191.63729858398438,66],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":200},{"localized_name":"images","name":"images","type":"IMAGE","link":203},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":199}],"outputs":[{"localized_name":"expression","name":"expression","type":"EXPRESSION","links":[197]}],"properties":{"Node name for S&R":"GetExpression2"},"widgets_values":[]},{"id":60,"type":"GetFaceLandmarks","pos":[518.9132690429688,612.9461669921875],"size":[170.62362670898438,46],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":122},{"localized_name":"images","name":"images","type":"IMAGE","link":121}],"outputs":[{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","links":[143,198]}],"properties":{"Node name for S&R":"GetFaceLandmarks"},"widgets_values":[]},{"id":61,"type":"GetFaceLandmarks","pos":[519.7196655273438,938.032958984375],"size":[170.62362670898438,46],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":129},{"localized_name":"images","name":"images","type":"IMAGE","link":187}],"outputs":[{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","links":[145,199]}],"properties":{"Node name for S&R":"GetFaceLandmarks"},"widgets_values":[]},{"id":2,"type":"HMFaceToolkitsLoader","pos":[45.41996383666992,769.9866333007812],"size":[315,82],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"gpu_id","name":"gpu_id","type":"INT","widget":{"name":"gpu_id"},"link":null},{"localized_name":"deployment","name":"deployment","type":"COMBO","widget":{"name":"deployment"},"link":null}],"outputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","slot_index":0,"links":[122,129,147,149,200,201]}],"properties":{"Node name for S&R":"HMFaceToolkitsLoader"},"widgets_values":[0,"huggingface"]},{"id":58,"type":"LoadImage","pos":[37.4466552734375,260.9000244140625],"size":[315,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[121,151,202]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["chillout.jpg","image"]},{"id":79,"type":"LoadImage","pos":[52.432804107666016,1018.0931396484375],"size":[274.080078125,314.00006103515625],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[187,188,203]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["ysll.jpg","image"]}],"links":[[121,58,0,60,1,"IMAGE"],[122,2,0,60,0,"FACE_TOOLKITS"],[129,2,0,61,0,"FACE_TOOLKITS"],[143,60,0,67,2,"FACELANDMARKS222"],[145,61,0,69,2,"FACELANDMARKS222"],[147,2,0,67,0,"FACE_TOOLKITS"],[149,2,0,69,0,"FACE_TOOLKITS"],[151,58,0,67,1,"IMAGE"],[187,79,0,61,1,"IMAGE"],[188,79,0,69,1,"IMAGE"],[190,82,0,80,0,"HMIMAGEPIPELINE"],[191,67,0,80,1,"HEAD_POSE"],[193,69,0,80,3,"HEAD_POSE"],[195,80,0,83,0,"IMAGE"],[196,84,0,80,2,"EXPRESSION"],[197,85,0,80,4,"EXPRESSION"],[198,60,0,84,2,"FACELANDMARKS222"],[199,61,0,85,2,"FACELANDMARKS222"],[200,2,0,85,0,"FACE_TOOLKITS"],[201,2,0,84,0,"FACE_TOOLKITS"],[202,58,0,84,1,"IMAGE"],[203,79,0,85,1,"IMAGE"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.8264462809917358,"offset":[459.39712102477415,-14.59995707513657]},"frontendVersion":"1.19.2"},"version":0.4} \ No newline at end of file +{ + "2": { + "inputs": { + "gpu_id": 0, + "deployment": "huggingface" + }, + "class_type": "HMFaceToolkitsLoader", + "_meta": { + "title": "Load Face Toolkits" + } + }, + "58": { + "inputs": { + "image": "chillout.jpg" + }, + "class_type": "LoadImage", + "_meta": { + "title": "Load Image" + } + }, + "60": { + "inputs": { + "face_toolkits": [ + "2", + 0 + ], + "images": [ + "58", + 0 + ] + }, + "class_type": "GetFaceLandmarks", + "_meta": { + "title": "Get Face Landmarks" + } + }, + "61": { + "inputs": { + "face_toolkits": [ + "2", + 0 + ], + "images": [ + "79", + 0 + ] + }, + "class_type": "GetFaceLandmarks", + "_meta": { + "title": "Get Face Landmarks" + } + }, + "67": { + "inputs": { + "crop": true, + "face_toolkits": [ + "2", + 0 + ], + "images": [ + "58", + 0 + ], + "landmarks": [ + "60", + 0 + ] + }, + "class_type": "GetHeadPose", + "_meta": { + "title": "Get Head Pose" + } + }, + "69": { + "inputs": { + "crop": true, + "face_toolkits": [ + "2", + 0 + ], + "images": [ + "79", + 0 + ], + "landmarks": [ + "61", + 0 + ] + }, + "class_type": "GetHeadPose", + "_meta": { + "title": "Get Head Pose" + } + }, + "79": { + "inputs": { + "image": "ysll.jpg" + }, + "class_type": "LoadImage", + "_meta": { + "title": "Load Image" + } + }, + "80": { + "inputs": { + "trans_ratio": 0, + "prompt": "(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture", + "negative_prompt": "", + "steps": 25, + "seed": 452553676269313, + "guidance_scale": 2, + "gpu_id": 0, + "hm_image_pipeline": [ + "82", + 0 + ], + "ref_head_pose": [ + "67", + 0 + ], + "ref_expression": [ + "84", + 0 + ], + "drive_head_pose": [ + "69", + 0 + ], + "drive_expression": [ + "85", + 0 + ] + }, + "class_type": "HMPipelineImage", + "_meta": { + "title": "HelloMeme Image Pipeline" + } + }, + "82": { + "inputs": { + "checkpoint": "[preset]DisneyPixarCartoonB", + "lora": "None", + "vae": "same as checkpoint", + "version": "v5", + "stylize": "x1", + "deployment": "huggingface", + "lora_scale": 1 + }, + "class_type": "HMImagePipelineLoader", + "_meta": { + "title": "Load HelloMemeImage Pipeline" + } + }, + "83": { + "inputs": { + "images": [ + "80", + 0 + ] + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Preview Image" + } + }, + "84": { + "inputs": { + "face_toolkits": [ + "2", + 0 + ], + "images": [ + "58", + 0 + ], + "landmarks": [ + "60", + 0 + ] + }, + "class_type": "GetExpression2", + "_meta": { + "title": "Get Face Expression V2" + } + }, + "85": { + "inputs": { + "face_toolkits": [ + "2", + 0 + ], + "images": [ + "79", + 0 + ], + "landmarks": [ + "61", + 0 + ] + }, + "class_type": "GetExpression2", + "_meta": { + "title": "Get Face Expression V2" + } + } +} \ No newline at end of file diff --git a/example_workflows/video_generation.json b/example_workflows/video_generation.json index b097f09..4a90c4b 100644 --- a/example_workflows/video_generation.json +++ b/example_workflows/video_generation.json @@ -1 +1,819 @@ -{"id":"4c8b7397-a69c-46cc-b553-00e600557365","revision":0,"last_node_id":76,"last_link_id":178,"nodes":[{"id":67,"type":"GetHeadPose","pos":[814.1528930664062,495.9793701171875],"size":[270,98],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":147},{"localized_name":"images","name":"images","type":"IMAGE","link":151},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":143},{"localized_name":"crop","name":"crop","type":"BOOLEAN","widget":{"name":"crop"},"link":null}],"outputs":[{"localized_name":"head_pose","name":"head_pose","type":"HEAD_POSE","links":[166]}],"properties":{"Node name for S&R":"GetHeadPose"},"widgets_values":[false]},{"id":75,"type":"GetExpression2","pos":[814.1359252929688,689.6795654296875],"size":[191.63729858398438,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":173},{"localized_name":"images","name":"images","type":"IMAGE","link":175},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":178}],"outputs":[{"localized_name":"expression","name":"expression","type":"EXPRESSION","links":[171]}],"properties":{"Node name for S&R":"GetExpression2"},"widgets_values":[]},{"id":76,"type":"GetExpression2","pos":[833.6573486328125,1042.0396728515625],"size":[191.63729858398438,66],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":174},{"localized_name":"images","name":"images","type":"IMAGE","link":176},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":177}],"outputs":[{"localized_name":"expression","name":"expression","type":"EXPRESSION","links":[172]}],"properties":{"Node name for S&R":"GetExpression2"},"widgets_values":[]},{"id":2,"type":"HMFaceToolkitsLoader","pos":[45.41996383666992,769.9866333007812],"size":[315,82],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"gpu_id","name":"gpu_id","type":"INT","widget":{"name":"gpu_id"},"link":null},{"localized_name":"deployment","name":"deployment","type":"COMBO","widget":{"name":"deployment"},"link":null}],"outputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","slot_index":0,"links":[122,129,147,149,173,174]}],"properties":{"Node name for S&R":"HMFaceToolkitsLoader"},"widgets_values":[0,"huggingface"]},{"id":61,"type":"GetFaceLandmarks","pos":[519.7196655273438,938.032958984375],"size":[170.62362670898438,46],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":129},{"localized_name":"images","name":"images","type":"IMAGE","link":131}],"outputs":[{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","links":[145,177]}],"properties":{"Node name for S&R":"GetFaceLandmarks"},"widgets_values":[]},{"id":60,"type":"GetFaceLandmarks","pos":[518.9132690429688,612.9461669921875],"size":[170.62362670898438,46],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":122},{"localized_name":"images","name":"images","type":"IMAGE","link":121}],"outputs":[{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","links":[143,178]}],"properties":{"Node name for S&R":"GetFaceLandmarks"},"widgets_values":[]},{"id":59,"type":"VHS_VideoCombine","pos":[1658.52490234375,685.6950073242188],"size":[329.22479248046875,633.2247924804688],"flags":{},"order":11,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":170},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":118},{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"frame_rate","name":"frame_rate","type":"FLOAT","widget":{"name":"frame_rate"},"link":null},{"localized_name":"loop_count","name":"loop_count","type":"INT","widget":{"name":"loop_count"},"link":null},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"pingpong","name":"pingpong","type":"BOOLEAN","widget":{"name":"pingpong"},"link":null},{"localized_name":"save_output","name":"save_output","type":"BOOLEAN","widget":{"name":"save_output"},"link":null}],"outputs":[{"localized_name":"Filenames","name":"Filenames","type":"VHS_FILENAMES","links":null}],"properties":{"Node name for S&R":"VHS_VideoCombine"},"widgets_values":{"frame_rate":15,"loop_count":0,"filename_prefix":"AnimateDiff","format":"video/h264-mp4","pix_fmt":"yuv420p","crf":18,"save_metadata":true,"pingpong":false,"save_output":true,"videopreview":{"hidden":false,"paused":false,"params":{"filename":"AnimateDiff_00611-audio.mp4","subfolder":"","type":"output","format":"video/h264-mp4","frame_rate":15},"muted":false}}},{"id":74,"type":"HMPipelineVideo","pos":[1296.2967529296875,659.4104614257812],"size":[270,330],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"hm_video_pipeline","name":"hm_video_pipeline","type":"HMVIDEOPIPELINE","link":165},{"localized_name":"ref_head_pose","name":"ref_head_pose","type":"HEAD_POSE","link":166},{"localized_name":"ref_expression","name":"ref_expression","type":"EXPRESSION","link":171},{"localized_name":"drive_head_pose","name":"drive_head_pose","type":"HEAD_POSE","link":168},{"localized_name":"drive_expression","name":"drive_expression","type":"EXPRESSION","link":172},{"localized_name":"trans_ratio","name":"trans_ratio","type":"FLOAT","widget":{"name":"trans_ratio"},"link":null},{"localized_name":"patch_overlap","name":"patch_overlap","type":"INT","widget":{"name":"patch_overlap"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"steps","name":"steps","type":"INT","widget":{"name":"steps"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"gpu_id","name":"gpu_id","type":"INT","widget":{"name":"gpu_id"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[170]},{"localized_name":"LATENT","name":"LATENT","type":"LATENT","links":null}],"properties":{"Node name for S&R":"HMPipelineVideo"},"widgets_values":[0,4,"(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture","",25,1055114486951082,"randomize",3.5,0]},{"id":57,"type":"HMVideoPipelineLoader","pos":[798.4796752929688,193.7666015625],"size":[352.79998779296875,202],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"checkpoint","name":"checkpoint","shape":7,"type":"COMBO","widget":{"name":"checkpoint"},"link":null},{"localized_name":"lora","name":"lora","shape":7,"type":"COMBO","widget":{"name":"lora"},"link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"COMBO","widget":{"name":"vae"},"link":null},{"localized_name":"version","name":"version","shape":7,"type":"COMBO","widget":{"name":"version"},"link":null},{"localized_name":"stylize","name":"stylize","shape":7,"type":"COMBO","widget":{"name":"stylize"},"link":null},{"localized_name":"deployment","name":"deployment","shape":7,"type":"COMBO","widget":{"name":"deployment"},"link":null},{"localized_name":"lora_scale","name":"lora_scale","shape":7,"type":"FLOAT","widget":{"name":"lora_scale"},"link":null}],"outputs":[{"localized_name":"hm_video_pipeline","name":"hm_video_pipeline","type":"HMVIDEOPIPELINE","slot_index":0,"links":[165]}],"properties":{"Node name for S&R":"HMVideoPipelineLoader"},"widgets_values":["[preset]DisneyPixarCartoonB","None","same as checkpoint","v4","x1","huggingface",1]},{"id":58,"type":"LoadImage","pos":[37.4466552734375,260.9000244140625],"size":[315,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[121,151,175]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["chillout.jpg","image"]},{"id":69,"type":"GetHeadPose","pos":[819.799560546875,860.5927734375],"size":[270,98],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"face_toolkits","name":"face_toolkits","type":"FACE_TOOLKITS","link":149},{"localized_name":"images","name":"images","type":"IMAGE","link":153},{"localized_name":"landmarks","name":"landmarks","type":"FACELANDMARKS222","link":145},{"localized_name":"crop","name":"crop","type":"BOOLEAN","widget":{"name":"crop"},"link":null}],"outputs":[{"localized_name":"head_pose","name":"head_pose","type":"HEAD_POSE","links":[168]}],"properties":{"Node name for S&R":"GetHeadPose"},"widgets_values":[true]},{"id":48,"type":"VHS_LoadVideo","pos":[74.54000091552734,996.8733520507812],"size":[247.455078125,680.3645629882812],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"meta_batch","name":"meta_batch","shape":7,"type":"VHS_BatchManager","link":null},{"localized_name":"vae","name":"vae","shape":7,"type":"VAE","link":null},{"localized_name":"video","name":"video","type":"COMBO","widget":{"name":"video"},"link":null},{"localized_name":"force_rate","name":"force_rate","type":"INT","widget":{"name":"force_rate"},"link":null},{"localized_name":"force_size","name":"force_size","type":"COMBO","widget":{"name":"force_size"},"link":null},{"localized_name":"custom_width","name":"custom_width","type":"INT","widget":{"name":"custom_width"},"link":null},{"localized_name":"custom_height","name":"custom_height","type":"INT","widget":{"name":"custom_height"},"link":null},{"localized_name":"frame_load_cap","name":"frame_load_cap","type":"INT","widget":{"name":"frame_load_cap"},"link":null},{"localized_name":"skip_first_frames","name":"skip_first_frames","type":"INT","widget":{"name":"skip_first_frames"},"link":null},{"localized_name":"select_every_nth","name":"select_every_nth","type":"INT","widget":{"name":"select_every_nth"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","slot_index":0,"links":[131,153,176]},{"localized_name":"frame_count","name":"frame_count","type":"INT","links":null},{"localized_name":"audio","name":"audio","type":"AUDIO","slot_index":2,"links":[118]},{"localized_name":"video_info","name":"video_info","type":"VHS_VIDEOINFO","links":null}],"properties":{"Node name for S&R":"VHS_LoadVideo"},"widgets_values":{"video":"tiktok.mp4","force_rate":15,"force_size":"Disabled","custom_width":512,"custom_height":512,"frame_load_cap":0,"skip_first_frames":0,"select_every_nth":1,"choose video to upload":"image","videopreview":{"hidden":false,"paused":false,"params":{"force_rate":15,"frame_load_cap":0,"skip_first_frames":0,"select_every_nth":1,"filename":"tiktok.mp4","type":"input","format":"video/mp4"},"muted":false}}}],"links":[[118,48,2,59,1,"AUDIO"],[121,58,0,60,1,"IMAGE"],[122,2,0,60,0,"FACE_TOOLKITS"],[129,2,0,61,0,"FACE_TOOLKITS"],[131,48,0,61,1,"IMAGE"],[143,60,0,67,2,"FACELANDMARKS222"],[145,61,0,69,2,"FACELANDMARKS222"],[147,2,0,67,0,"FACE_TOOLKITS"],[149,2,0,69,0,"FACE_TOOLKITS"],[151,58,0,67,1,"IMAGE"],[153,48,0,69,1,"IMAGE"],[165,57,0,74,0,"HMVIDEOPIPELINE"],[166,67,0,74,1,"HEAD_POSE"],[168,69,0,74,3,"HEAD_POSE"],[170,74,0,59,0,"IMAGE"],[171,75,0,74,2,"EXPRESSION"],[172,76,0,74,4,"EXPRESSION"],[173,2,0,75,0,"FACE_TOOLKITS"],[174,2,0,76,0,"FACE_TOOLKITS"],[175,58,0,75,1,"IMAGE"],[176,48,0,76,1,"IMAGE"],[177,61,0,76,2,"FACELANDMARKS222"],[178,60,0,75,2,"FACELANDMARKS222"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.8264462809917358,"offset":[605.403886161493,-0.08003092767563647]},"frontendVersion":"1.19.2"},"version":0.4} \ No newline at end of file +{ + "id": "4c8b7397-a69c-46cc-b553-00e600557365", + "revision": 0, + "last_node_id": 76, + "last_link_id": 178, + "nodes": [ + { + "id": 75, + "type": "GetExpression2", + "pos": [ + 814.1359252929688, + 689.6795654296875 + ], + "size": [ + 191.63729858398438, + 66 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "link": 173 + }, + { + "name": "images", + "type": "IMAGE", + "link": 175 + }, + { + "name": "landmarks", + "type": "FACELANDMARKS222", + "link": 178 + } + ], + "outputs": [ + { + "name": "expression", + "type": "EXPRESSION", + "links": [ + 171 + ] + } + ], + "properties": { + "Node name for S&R": "GetExpression2" + }, + "widgets_values": [] + }, + { + "id": 76, + "type": "GetExpression2", + "pos": [ + 833.6573486328125, + 1042.0396728515625 + ], + "size": [ + 191.63729858398438, + 66 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "link": 174 + }, + { + "name": "images", + "type": "IMAGE", + "link": 176 + }, + { + "name": "landmarks", + "type": "FACELANDMARKS222", + "link": 177 + } + ], + "outputs": [ + { + "name": "expression", + "type": "EXPRESSION", + "links": [ + 172 + ] + } + ], + "properties": { + "Node name for S&R": "GetExpression2" + }, + "widgets_values": [] + }, + { + "id": 2, + "type": "HMFaceToolkitsLoader", + "pos": [ + 45.41996383666992, + 769.9866333007812 + ], + "size": [ + 315, + 82 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "slot_index": 0, + "links": [ + 122, + 129, + 147, + 149, + 173, + 174 + ] + } + ], + "properties": { + "Node name for S&R": "HMFaceToolkitsLoader" + }, + "widgets_values": [ + 0, + "huggingface" + ] + }, + { + "id": 61, + "type": "GetFaceLandmarks", + "pos": [ + 519.7196655273438, + 938.032958984375 + ], + "size": [ + 170.62362670898438, + 46 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "link": 129 + }, + { + "name": "images", + "type": "IMAGE", + "link": 131 + } + ], + "outputs": [ + { + "name": "landmarks", + "type": "FACELANDMARKS222", + "links": [ + 145, + 177 + ] + } + ], + "properties": { + "Node name for S&R": "GetFaceLandmarks" + }, + "widgets_values": [] + }, + { + "id": 60, + "type": "GetFaceLandmarks", + "pos": [ + 518.9132690429688, + 612.9461669921875 + ], + "size": [ + 170.62362670898438, + 46 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "link": 122 + }, + { + "name": "images", + "type": "IMAGE", + "link": 121 + } + ], + "outputs": [ + { + "name": "landmarks", + "type": "FACELANDMARKS222", + "links": [ + 143, + 178 + ] + } + ], + "properties": { + "Node name for S&R": "GetFaceLandmarks" + }, + "widgets_values": [] + }, + { + "id": 59, + "type": "VHS_VideoCombine", + "pos": [ + 1658.52490234375, + 685.6950073242188 + ], + "size": [ + 329.22479248046875, + 633.2247924804688 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 170 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": 118 + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 15, + "loop_count": 0, + "filename_prefix": "AnimateDiff", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 18, + "save_metadata": true, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "AnimateDiff_00703-audio.mp4", + "subfolder": "", + "type": "output", + "format": "video/h264-mp4", + "frame_rate": 15 + }, + "muted": false + } + } + }, + { + "id": 74, + "type": "HMPipelineVideo", + "pos": [ + 1296.2967529296875, + 659.4104614257812 + ], + "size": [ + 270, + 330 + ], + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "hm_video_pipeline", + "type": "HMVIDEOPIPELINE", + "link": 165 + }, + { + "name": "ref_head_pose", + "type": "HEAD_POSE", + "link": 166 + }, + { + "name": "ref_expression", + "type": "EXPRESSION", + "link": 171 + }, + { + "name": "drive_head_pose", + "type": "HEAD_POSE", + "link": 168 + }, + { + "name": "drive_expression", + "type": "EXPRESSION", + "link": 172 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 170 + ] + }, + { + "name": "LATENT", + "type": "LATENT", + "links": null + } + ], + "properties": { + "Node name for S&R": "HMPipelineVideo" + }, + "widgets_values": [ + 0, + 4, + "(best quality), highly detailed, ultra-detailed, headshot, person, well-placed five sense organs, looking at the viewer, centered composition, sharp focus, realistic skin texture", + "", + 25, + 1075819929118817, + "randomize", + 2.5, + 0 + ] + }, + { + "id": 69, + "type": "GetHeadPose", + "pos": [ + 819.799560546875, + 860.5927734375 + ], + "size": [ + 270, + 98 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "link": 149 + }, + { + "name": "images", + "type": "IMAGE", + "link": 153 + }, + { + "name": "landmarks", + "type": "FACELANDMARKS222", + "link": 145 + } + ], + "outputs": [ + { + "name": "head_pose", + "type": "HEAD_POSE", + "links": [ + 168 + ] + } + ], + "properties": { + "Node name for S&R": "GetHeadPose" + }, + "widgets_values": [ + true + ] + }, + { + "id": 57, + "type": "HMVideoPipelineLoader", + "pos": [ + 798.4796752929688, + 193.7666015625 + ], + "size": [ + 352.79998779296875, + 202 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "hm_video_pipeline", + "type": "HMVIDEOPIPELINE", + "slot_index": 0, + "links": [ + 165 + ] + } + ], + "properties": { + "Node name for S&R": "HMVideoPipelineLoader" + }, + "widgets_values": [ + "[preset]RealisticVisionV60B1", + "None", + "same as checkpoint", + "v5", + "x1", + "huggingface", + 1 + ] + }, + { + "id": 48, + "type": "VHS_LoadVideo", + "pos": [ + 74.54000091552734, + 996.8733520507812 + ], + "size": [ + 247.455078125, + 503.455078125 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 131, + 153, + 176 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "slot_index": 2, + "links": [ + 118 + ] + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "jgz.mp4", + "force_rate": 15, + "force_size": "Disabled", + "custom_width": 512, + "custom_height": 512, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "force_rate": 15, + "frame_load_cap": 0, + "skip_first_frames": 0, + "select_every_nth": 1, + "filename": "jgz.mp4", + "type": "input", + "format": "video/mp4" + }, + "muted": false + } + } + }, + { + "id": 58, + "type": "LoadImage", + "pos": [ + 37.4466552734375, + 260.9000244140625 + ], + "size": [ + 315, + 314 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "slot_index": 0, + "links": [ + 121, + 151, + 175 + ] + }, + { + "name": "MASK", + "type": "MASK", + "links": null + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "i5.jpg", + "image" + ] + }, + { + "id": 67, + "type": "GetHeadPose", + "pos": [ + 814.1528930664062, + 495.9793701171875 + ], + "size": [ + 270, + 98 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "face_toolkits", + "type": "FACE_TOOLKITS", + "link": 147 + }, + { + "name": "images", + "type": "IMAGE", + "link": 151 + }, + { + "name": "landmarks", + "type": "FACELANDMARKS222", + "link": 143 + } + ], + "outputs": [ + { + "name": "head_pose", + "type": "HEAD_POSE", + "links": [ + 166 + ] + } + ], + "properties": { + "Node name for S&R": "GetHeadPose" + }, + "widgets_values": [ + true + ] + } + ], + "links": [ + [ + 118, + 48, + 2, + 59, + 1, + "AUDIO" + ], + [ + 121, + 58, + 0, + 60, + 1, + "IMAGE" + ], + [ + 122, + 2, + 0, + 60, + 0, + "FACE_TOOLKITS" + ], + [ + 129, + 2, + 0, + 61, + 0, + "FACE_TOOLKITS" + ], + [ + 131, + 48, + 0, + 61, + 1, + "IMAGE" + ], + [ + 143, + 60, + 0, + 67, + 2, + "FACELANDMARKS222" + ], + [ + 145, + 61, + 0, + 69, + 2, + "FACELANDMARKS222" + ], + [ + 147, + 2, + 0, + 67, + 0, + "FACE_TOOLKITS" + ], + [ + 149, + 2, + 0, + 69, + 0, + "FACE_TOOLKITS" + ], + [ + 151, + 58, + 0, + 67, + 1, + "IMAGE" + ], + [ + 153, + 48, + 0, + 69, + 1, + "IMAGE" + ], + [ + 165, + 57, + 0, + 74, + 0, + "HMVIDEOPIPELINE" + ], + [ + 166, + 67, + 0, + 74, + 1, + "HEAD_POSE" + ], + [ + 168, + 69, + 0, + 74, + 3, + "HEAD_POSE" + ], + [ + 170, + 74, + 0, + 59, + 0, + "IMAGE" + ], + [ + 171, + 75, + 0, + 74, + 2, + "EXPRESSION" + ], + [ + 172, + 76, + 0, + 74, + 4, + "EXPRESSION" + ], + [ + 173, + 2, + 0, + 75, + 0, + "FACE_TOOLKITS" + ], + [ + 174, + 2, + 0, + 76, + 0, + "FACE_TOOLKITS" + ], + [ + 175, + 58, + 0, + 75, + 1, + "IMAGE" + ], + [ + 176, + 48, + 0, + 76, + 1, + "IMAGE" + ], + [ + 177, + 61, + 0, + 76, + 2, + "FACELANDMARKS222" + ], + [ + 178, + 60, + 0, + 75, + 2, + "FACELANDMARKS222" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.8264462809917358, + "offset": [ + 493.27712102477454, + -69.45337656976547 + ] + }, + "frontendVersion": "1.19.2" + }, + "version": 0.4 +} \ No newline at end of file diff --git a/examples/i5.jpg b/examples/i5.jpg new file mode 100644 index 0000000..c405110 Binary files /dev/null and b/examples/i5.jpg differ diff --git a/hellomeme/__init__.py b/hellomeme/__init__.py index 4fa767d..2aefa89 100644 --- a/hellomeme/__init__.py +++ b/hellomeme/__init__.py @@ -10,4 +10,4 @@ from .pipelines import (HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline, - HM5ImagePipeline) \ No newline at end of file + HM5ImagePipeline, HM5VideoPipeline) \ No newline at end of file diff --git a/hellomeme/models/__init__.py b/hellomeme/models/__init__.py index fe72eb7..205c5db 100644 --- a/hellomeme/models/__init__.py +++ b/hellomeme/models/__init__.py @@ -13,7 +13,7 @@ from .hm_control import (HMControlNet, HMControlNet2, HMV2ControlNet, HMV2Contro HMV3ControlNet, HMControlNetBase, HM5ControlNetBase, HM4SD15ControlProj, HM5SD15ControlProj) from .hm_adapters import (HMReferenceAdapter, HM3ReferenceAdapter, HM5ReferenceAdapter, - HM3MotionAdapter, HMPipeline) + HM3MotionAdapter, HM5MotionAdapter, HMPipeline) from .hm_denoising_3d import HMDenoising3D from .hm3_denoising_3d import HM3Denoising3D from .hm3_denoising_motion import HM3DenoisingMotion diff --git a/hellomeme/models/hm3_denoising_3d.py b/hellomeme/models/hm3_denoising_3d.py index e69050c..1e3faf3 100644 --- a/hellomeme/models/hm3_denoising_3d.py +++ b/hellomeme/models/hm3_denoising_3d.py @@ -193,7 +193,9 @@ class HM3Denoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter): if not control_hidden_states is None and f'down3_{idx}' in control_hidden_states: sample += rearrange(control_hidden_states[f'down3_{idx}'], "b c f h w -> (b f) c h w") if hasattr(self, 'motion_down') and use_motion: - sample = self.motion_down[idx](sample, motion_pad_hidden_states[f'down_{idx}'], emb, num_frames) + sample = self.motion_down[idx](sample, + None if motion_pad_hidden_states is None else motion_pad_hidden_states[f'down_{idx}'], + emb, num_frames) down_block_res_samples += res_samples @@ -251,7 +253,9 @@ class HM3Denoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter): if hasattr(self, "reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states: sample = self.reference_modules_up[i](sample, reference_hidden_states[f'up_{i}'], num_frames) if hasattr(self, 'motion_up') and use_motion: - sample = self.motion_up[i](sample, motion_pad_hidden_states[f'up_{i}'], emb, num_frames) + sample = self.motion_up[i](sample, + None if motion_pad_hidden_states is None else motion_pad_hidden_states[f'up_{i}'], + emb, num_frames) if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention: sample = upsample_block( diff --git a/hellomeme/models/hm_adapters.py b/hellomeme/models/hm_adapters.py index 064d3a6..af6fce2 100644 --- a/hellomeme/models/hm_adapters.py +++ b/hellomeme/models/hm_adapters.py @@ -42,7 +42,8 @@ from diffusers.models.lora import adjust_lora_scale_text_encoder from .hm_blocks import (SKReferenceAttentionV3, SKReferenceAttentionV5, SKReferenceAttention, - SKMotionModule) + SKMotionModule, + SKMotionModuleV5) logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -183,6 +184,32 @@ class HM3MotionAdapter(ModelMixin, ConfigMixin): ) +class HM5MotionAdapter(ModelMixin, ConfigMixin): + @register_to_config + def __init__(self, block_down_channels: Tuple[int] = (320, 640, 1280, 1280), + block_up_channels: Tuple[int] = (1280, 1280, 1280, 640), + num_attention_heads: int = 8): + super().__init__() + self.motion_down = nn.ModuleList([]) + + for i, in_channels in enumerate(block_down_channels): + self.motion_down.append( + SKMotionModuleV5( + in_channels=in_channels, + num_attention_heads=num_attention_heads, + ) + ) + + self.motion_up = nn.ModuleList([]) + for i, in_channels in enumerate(block_up_channels): + self.motion_up.append( + SKMotionModuleV5( + in_channels=in_channels, + num_attention_heads=num_attention_heads, + ) + ) + + class HMPipeline(StableDiffusionImg2ImgPipeline): @torch.no_grad() def load_lora_weights_sk(self, unet, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]], diff --git a/hellomeme/models/hm_blocks.py b/hellomeme/models/hm_blocks.py index bacb6e4..8c3cfd4 100644 --- a/hellomeme/models/hm_blocks.py +++ b/hellomeme/models/hm_blocks.py @@ -14,9 +14,12 @@ import torch.utils.checkpoint from typing import Tuple from einops import rearrange -from diffusers.models.attention import Attention, FeedForward, BasicTransformerBlock +from diffusers.models.attention import Attention, FeedForward +from diffusers.models.unets.unet_motion_model import AnimateDiffTransformer3D from torchvision.models.resnet import BasicBlock -from diffusers.models.embeddings import SinusoidalPositionalEmbedding, get_1d_sincos_pos_embed_from_grid, TimestepEmbedding +from diffusers.models.embeddings import (SinusoidalPositionalEmbedding, + get_1d_sincos_pos_embed_from_grid, + TimestepEmbedding) from diffusers.utils import logging @@ -725,6 +728,51 @@ class SKMotionModule(nn.Module): return hidden_states + self.proj(res_temp) +class SKMotionModuleV5(nn.Module): + def __init__( + self, + in_channels: int, + num_attention_heads: int = 8, + ): + super().__init__() + + self.attn1 = AnimateDiffTransformer3D( + in_channels=in_channels, + num_attention_heads=num_attention_heads, + attention_head_dim=in_channels // num_attention_heads, + dropout=0.0, + double_self_attention=True, + positional_embeddings='sinusoidal', + num_positional_embeddings=32, + + ) + + self.attn_pad = AnimateDiffTransformer3D( + in_channels=in_channels, + num_attention_heads=num_attention_heads, + attention_head_dim=in_channels // num_attention_heads, + dropout=0.0, + double_self_attention=True, + positional_embeddings='sinusoidal', + num_positional_embeddings=32, + ) + + + def forward(self, hidden_states, pad_states, temb, num_frames): + c, h, w = hidden_states.shape[-3:] + res_temp = hidden_states + + if not pad_states is None: + res_temp = rearrange(res_temp, "(b f) c h w -> b c f h w", f=num_frames) + res_temp = torch.cat([pad_states[:,:,:1], res_temp, pad_states[:,:,-1:]], dim=2).contiguous() + res_temp = rearrange(res_temp, "b c f h w -> (b f) c h w") + res_temp = self.attn_pad(res_temp, num_frames=num_frames+2) + res_temp = rearrange(res_temp, "(b f) c h w -> b c f h w", f=num_frames+2).contiguous() + res_temp = rearrange(res_temp[:,:,1:-1], "b c f h w -> (b f) c h w", h=h, w=w) + hidden_states + + res_temp = self.attn1(res_temp, num_frames=num_frames) + return res_temp + # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnet.py def zero_module(module): for p in module.parameters(): diff --git a/hellomeme/pipelines/__init__.py b/hellomeme/pipelines/__init__.py index ae89f78..d8bb482 100644 --- a/hellomeme/pipelines/__init__.py +++ b/hellomeme/pipelines/__init__.py @@ -12,4 +12,5 @@ from .pipline_hm_image import HMImagePipeline from .pipline_hm_video import HMVideoPipeline from .pipline_hm3_image import HM3ImagePipeline from .pipline_hm3_video import HM3VideoPipeline -from .pipline_hm5_image import HM5ImagePipeline \ No newline at end of file +from .pipline_hm5_image import HM5ImagePipeline +from .pipline_hm5_video import HM5VideoPipeline \ No newline at end of file diff --git a/hellomeme/pipelines/pipline_hm5_video.py b/hellomeme/pipelines/pipline_hm5_video.py new file mode 100644 index 0000000..58ddfda --- /dev/null +++ b/hellomeme/pipelines/pipline_hm5_video.py @@ -0,0 +1,510 @@ +# coding: utf-8 + +""" +@File : hm_pipline_image.py +@Author : Songkey +@Email : songkey@pku.edu.cn +@Date : 5/30/2025 +@Desc : +adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py +""" + +import copy +from typing import Any, Callable, Dict, List, Optional, Union + +import torch +from einops import rearrange + +from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback +from diffusers.image_processor import PipelineImageInput +from diffusers.utils import deprecate + +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents +from diffusers import DPMSolverMultistepScheduler +from diffusers.utils.torch_utils import randn_tensor + +from ..models import (HM3Denoising3D, HM5MotionAdapter, + HM5ReferenceAdapter, HMPipeline, HM5ControlNetBase, HM5SD15ControlProj) + +class HM5VideoPipeline(HMPipeline): + def caryomitosis(self, version='v5', modelscope=False, **kwargs): + if hasattr(self, "unet_ref"): + del self.unet_ref + self.unet_ref = HM3Denoising3D.from_unet2d(self.unet) + self.unet_ref.cpu() + + self.num_frames = 12 + + if not isinstance(self.unet, HM3Denoising3D): + unet = HM3Denoising3D.from_unet2d(unet=self.unet) + # todo: 不够优雅 + del self.unet + self.unet = unet + self.unet.cpu() + + self.vae.cpu() + self.vae_decode = copy.deepcopy(self.vae) + self.text_encoder.cpu() + self.text_encoder_ref = copy.deepcopy(self.text_encoder) + self.safety_checker.cpu() + + def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False): + self.version = version + + if modelscope: + from modelscope import snapshot_download + hm_reference_dir = snapshot_download('songkey/hm5_reference') + hm_control_dir = snapshot_download('songkey/hm5_control_base') + hm_control_proj_dir = snapshot_download('songkey/hm5_control_proj') + hm_motion_dir = snapshot_download('songkey/hm5_motion') + else: + hm_reference_dir = 'songkey/hm5_reference' + hm_control_dir = 'songkey/hm5_control_base' + hm_control_proj_dir = 'songkey/hm5_control_proj' + hm_motion_dir = 'songkey/hm5_motion' + + hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir) + motion_adapter = HM5MotionAdapter.from_pretrained(hm_motion_dir) + + if isinstance(self.unet, HM3Denoising3D): + self.unet.insert_reference_adapter(hm_adapter) + self.unet.insert_reference_adapter(motion_adapter) + self.unet.to(device='cpu', dtype=dtype).eval() + + if hasattr(self, "unet_ref"): + self.unet_ref.to(device='cpu', dtype=dtype).eval() + + if hasattr(self, "mp_control"): + del self.mp_control + if hasattr(self, "mp_control_proj"): + del self.mp_control_proj + + self.mp_control = HM5ControlNetBase.from_pretrained(hm_control_dir) + self.mp_control_proj = HM5SD15ControlProj.from_pretrained(hm_control_proj_dir) + self.mp_control_proj.to(device='cpu', dtype=dtype).eval() + self.mp_control.to(device='cpu', dtype=dtype).eval() + + self.vae.to(device='cpu', dtype=dtype).eval() + self.vae_decode.to(device='cpu', dtype=dtype).eval() + self.text_encoder.to(device='cpu', dtype=dtype).eval() + + @torch.no_grad() + def gen_ref_cache(self, latent, prompt_embeds, added_cond_kwargs, device, add_axis=False): + self.unet_ref.to(device=device) + latent = latent.to(device=device, dtype=prompt_embeds.dtype) + latent_model_input = torch.cat([torch.zeros_like(latent), latent]) if self.do_classifier_free_guidance else latent + cached_res = self.unet_ref( + latent_model_input.unsqueeze(2), + 0, + encoder_hidden_states=prompt_embeds, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[1] + self.unet_ref.cpu() + + if add_axis: + for k, v in cached_res.items(): + cached_res[k] = rearrange(v, 'b c h w -> b c 1 h w') + + return dicts_to_device([cached_res], device='cpu')[0] + + @torch.no_grad() + def gen_video(self, indexes, pad_dict, base_noise, depth, control_dict, step, more_params): + device = more_params['device'] + timesteps = more_params['timesteps'] + timestep_cond = more_params['timestep_cond'] + sigmas = more_params['sigmas'] + prompt_embeds = more_params['prompt_embeds'] + cached_res = more_params['cached_res'] + added_cond_kwargs = more_params['added_cond_kwargs'] + extra_step_kwargs = more_params['extra_step_kwargs'] + + num_indexes = len(indexes) + if num_indexes <= self.num_frames: + paded_indexes = indexes + + paded_indexes_noise = list(range(len(indexes))) + + paded_indexes = copy.deepcopy(paded_indexes) + + drive_idx_chunks = [paded_indexes] + drive_idx_chunks_noise = [paded_indexes_noise] + else: + stride = self.num_frames - 1 + paded_indexes = indexes + paded_indexes = copy.deepcopy(paded_indexes) + + num_paded_indexes = len(paded_indexes) + paded_indexes_noise = list(range(num_paded_indexes)) + + drive_idx_chunks = [paded_indexes[i:min(i + self.num_frames, num_paded_indexes)] for i in + range(0, num_paded_indexes, stride)] + drive_idx_chunks_noise = [paded_indexes_noise[i:min(i + self.num_frames, num_paded_indexes)] for i in + range(0, num_paded_indexes, stride)] + + if len(drive_idx_chunks[-1]) == 1: + drive_idx_chunks[-1] = drive_idx_chunks[:-1] + drive_idx_chunks_noise[-1] = drive_idx_chunks_noise[:-1] + + margin_indexes = set() + for chunk in drive_idx_chunks: + margin_indexes.add(chunk[0]) + margin_indexes.add(chunk[-1]) + margin_indexes = sorted(list(margin_indexes)) + + self.gen_video(margin_indexes, pad_dict, base_noise, depth+1, control_dict, step, more_params) + + ## generate latent + scheduler = DPMSolverMultistepScheduler( + num_train_timesteps=1000, + beta_start=0.00085, + beta_end=0.012, + beta_schedule="scaled_linear", + algorithm_type="sde-dpmsolver++", + ) + + + timesteps, _ = retrieve_timesteps(scheduler, step, device, timesteps, sigmas) + + N, C, _, H, W = base_noise.shape + + latents = base_noise[:,:,:len(paded_indexes)] * scheduler.init_noise_sigma + + self.unet.to(device=device) + with self.progress_bar(total=len(timesteps)) as progress_bar: + for i, t in enumerate(timesteps): + progress_bar.set_description(f"D{depth} GEN Frames") + progress_bar.update() + + noise_pred = torch.zeros_like(latents) + noise_pred = torch.cat([noise_pred, noise_pred], dim=0) if \ + self.do_classifier_free_guidance else noise_pred + noise_pred_counter = torch.zeros([1, 1, noise_pred.size(2), 1, 1], dtype=prompt_embeds.dtype).cpu() + + latent_model_input_all = torch.cat([latents] * 2, dim=0) if \ + self.do_classifier_free_guidance else latents + + latent_model_input_all = scheduler.scale_model_input(latent_model_input_all, t) + for chunk, chunk_noise in zip(drive_idx_chunks, drive_idx_chunks_noise): + if self.interrupt: + continue + + control_latent = cat_dicts([control_dict[ix] for ix in chunk], dim=2) + + if chunk[0] in pad_dict and chunk[-1] in pad_dict: + pad_latent = cat_dicts([pad_dict[chunk[0]], pad_dict[chunk[-1]]], dim=2) + + else: + pad_latent = None + + latent_model_input = latent_model_input_all[:, :, chunk_noise].to(device=device) + + noise_pred_chunk = self.unet( + latent_model_input, + t, + encoder_hidden_states=prompt_embeds, + reference_hidden_states=dicts_to_device([cached_res], device=device)[0], + control_hidden_states=dicts_to_device([control_latent], device=device)[0], + motion_pad_hidden_states=None if pad_latent is None else dicts_to_device([pad_latent], device=device)[0], + use_motion=True, + timestep_cond=timestep_cond, + cross_attention_kwargs=self.cross_attention_kwargs, + added_cond_kwargs=added_cond_kwargs, + return_dict=False, + )[0].cpu() + + noise_pred[:, :, chunk_noise] = noise_pred[:, :, chunk_noise] + noise_pred_chunk + noise_pred_counter[:, :, chunk_noise] = noise_pred_counter[:, :, chunk_noise] + 1 + + noise_pred = noise_pred / noise_pred_counter + if self.do_classifier_free_guidance: + noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) + noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) + latents = scheduler.step(noise_pred.cpu(), t, latents, + **extra_step_kwargs, return_dict=False)[0] + self.unet.cpu() + + if depth == 0: + return latents[:,:,:num_indexes] + else: + for chunk, chunk_noise in zip(drive_idx_chunks, drive_idx_chunks_noise): + for idx, idx_noise in zip(chunk, chunk_noise): + # if not idx in pad_dict: + pad_dict[idx] = self.gen_ref_cache(latents[:,:,idx_noise], prompt_embeds, + added_cond_kwargs, device, add_axis=True) + + @torch.no_grad() + def __call__( + self, + prompt: Union[str, List[str]] = None, + image: PipelineImageInput = None, + chunk_overlap: int = 4, + drive_params: Dict[str, Any] = None, + strength: float = 0.8, + num_inference_steps: Optional[int] = 50, + timesteps: List[int] = None, + sigmas: List[float] = None, + guidance_scale: Optional[float] = 7.5, + negative_prompt: Optional[Union[str, List[str]]] = None, + eta: Optional[float] = 0.0, + generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, + prompt_embeds: Optional[torch.Tensor] = None, + negative_prompt_embeds: Optional[torch.Tensor] = None, + ip_adapter_image: Optional[PipelineImageInput] = None, + ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None, + output_type: Optional[str] = "pil", + device: Optional[str] = "cpu", + return_dict: bool = True, + cross_attention_kwargs: Optional[Dict[str, Any]] = None, + clip_skip: int = None, + callback_on_step_end: Optional[ + Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks] + ] = None, + callback_on_step_end_tensor_inputs: List[str] = ["latents"], + **kwargs, + ): + callback = kwargs.pop("callback", None) + callback_steps = kwargs.pop("callback_steps", None) + num_images_per_prompt = 1 + + if callback is not None: + deprecate( + "callback", + "1.0.0", + "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + if callback_steps is not None: + deprecate( + "callback_steps", + "1.0.0", + "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`", + ) + + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. Raise error if not correct + self.check_inputs( + prompt, + strength, + callback_steps, + negative_prompt, + prompt_embeds, + negative_prompt_embeds, + ip_adapter_image, + ip_adapter_image_embeds, + callback_on_step_end_tensor_inputs, + ) + + self._guidance_scale = guidance_scale + self._clip_skip = clip_skip + self._cross_attention_kwargs = cross_attention_kwargs + self._interrupt = False + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + # device = self.device + + # 3. Encode input prompt + text_encoder_lora_scale = ( + self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None + ) + self.text_encoder_ref.to(device=device) + prompt_embeds_ref, negative_prompt_embeds_ref = self.encode_prompt_sk( + self.text_encoder_ref, + prompt, + device, + num_images_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + clip_skip=self.clip_skip, + ) + self.text_encoder_ref.cpu() + + self.text_encoder.to(device=device) + prompt_embeds, negative_prompt_embeds = self.encode_prompt_sk( + self.text_encoder, + prompt, + device, + num_images_per_prompt, + self.do_classifier_free_guidance, + negative_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + lora_scale=text_encoder_lora_scale, + clip_skip=self.clip_skip, + ) + self.text_encoder.cpu() + + # For classifier free guidance, we need to do two forward passes. + # Here we concatenate the unconditional and text embeddings into a single batch + # to avoid doing two forward passes + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) + prompt_embeds_ref = torch.cat([negative_prompt_embeds_ref, prompt_embeds_ref]) + + if ip_adapter_image is not None or ip_adapter_image_embeds is not None: + image_embeds = self.prepare_ip_adapter_image_embeds( + ip_adapter_image, + ip_adapter_image_embeds, + device, + batch_size * num_images_per_prompt, + self.do_classifier_free_guidance, + ) + + # 4. Preprocess image + image = self.image_processor.preprocess(image).to(device=device, dtype=prompt_embeds.dtype) + + raw_video_len = drive_params['condition'].size(2) + + # 6. Prepare latent variables + self.vae.to(device=device) + ref_latents = [ + retrieve_latents(self.vae.encode(image[i: i + 1].to(device=device)), generator=generator) + for i in range(batch_size) + ] + self.vae.cpu() + + ref_latents = torch.cat(ref_latents, dim=0) + ref_latents = self.vae.config.scaling_factor * ref_latents + c, h, w = ref_latents.shape[1:] + + # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline + extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) + + # 7.1 Add image embeds for IP-Adapter + added_cond_kwargs = ( + {"image_embeds": image_embeds} + if ip_adapter_image is not None or ip_adapter_image_embeds is not None + else None + ) + + # 7.2 Optionally get Guidance Scale Embedding + timestep_cond = None + if self.unet.config.time_cond_proj_dim is not None: + guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt) + timestep_cond = self.get_guidance_scale_embedding( + guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim + ).to(device=device, dtype=prompt_embeds.dtype) + + input_noise = randn_tensor([batch_size, c, raw_video_len + self.num_frames + 1, h, w], + generator=generator, dtype=prompt_embeds.dtype) + + cached_res = self.gen_ref_cache(ref_latents, prompt_embeds_ref, added_cond_kwargs, device) + + control_latents = [] + control_latent1_neg, control_latent2_neg = None, None + self.mp_control.to(device=device) + self.mp_control_proj.to(device=device) + with self.progress_bar(total=raw_video_len) as progress_bar: + for idx in range(raw_video_len): + progress_bar.set_description(f"GEN stage1") + progress_bar.update() + + condition = drive_params['condition'][:, :, idx:idx+1].clone().to(device=device) + + control_latent_input = {} + if 'drive_coeff' in drive_params: + drive_coeff = drive_params['drive_coeff'][:, idx:idx+1].clone().to(device=device) + face_parts = drive_params['face_parts'][:, idx:idx+1].clone().to(device=device) + + if control_latent1_neg is None and self.do_classifier_free_guidance: + control_latent1_neg = self.mp_control(condition=torch.ones_like(condition)*-1, + drive_coeff=torch.zeros_like(drive_coeff), + face_parts=torch.zeros_like(face_parts)) + control_latent1_neg = self.mp_control_proj(control_latent1_neg) + + control_latent1 = self.mp_control(condition=condition, + drive_coeff=drive_coeff, + face_parts=face_parts) + control_latent1 = self.mp_control_proj(control_latent1) + if self.do_classifier_free_guidance: + control_latent1 = cat_dicts([control_latent1_neg, control_latent1], dim=0) + + control_latent_input.update(control_latent1) + else: + pd_fpg = drive_params['pd_fpg'][:, idx:idx+1].clone().to(device=device) + + if control_latent2_neg is None and self.do_classifier_free_guidance: + control_latent2_neg = self.mp_control(condition=torch.ones_like(condition)*-1, + emo_embedding=torch.zeros_like(pd_fpg)) + control_latent2_neg = self.mp_control_proj(control_latent2_neg) + + control_latent2 = self.mp_control(condition=condition, emo_embedding=pd_fpg) + control_latent2 = self.mp_control_proj(control_latent2) + if self.do_classifier_free_guidance: + control_latent2 = cat_dicts([control_latent2_neg, control_latent2], dim=0) + + control_latent_input.update(control_latent2) + control_latents.append(dicts_to_device([control_latent_input], device='cpu')[0]) + + self.mp_control.cpu() + self.mp_control_proj.cpu() + + more_params = dict( + timesteps=timesteps, + timestep_cond=timestep_cond, + sigmas=sigmas, + prompt_embeds=prompt_embeds, + cached_res=cached_res, + device=device, + added_cond_kwargs=added_cond_kwargs, + extra_step_kwargs=extra_step_kwargs, + ) + + indexes = list(range(raw_video_len)) + latents = self.gen_video(indexes, dict(), input_noise, 0, control_latents, num_inference_steps, more_params) + latents_res = latents / self.vae_decode.config.scaling_factor + + self.vae_decode.to(device=device) + res_frames = [] + with self.progress_bar(total=raw_video_len) as progress_bar: + for i in range(raw_video_len): + ret_tensor = self.vae_decode.decode(latents_res[:, :, i, :, :].to(device=device), + return_dict=False, generator=generator)[0] + ret_image = self.image_processor.postprocess(ret_tensor, output_type=output_type) + res_frames.append(ret_image) + progress_bar.set_description(f"VAE DECODE") + progress_bar.update() + self.vae_decode.cpu() + return res_frames, latents_res + + +def merge_dicts(dictl, dictr, wl=0.5): + res = {} + for k in dictl.keys(): + res[k] = dictl[k] * wl + dictr[k] * (1-wl) + return res + + +def cat_dicts(dicts, dim=0): + res = {} + for k in dicts[0].keys(): + res[k] = torch.cat([d[k].clone() for d in dicts], dim=dim) + return res + + +def cat_dicts_ref(dicts): + res = {} + for k in dicts[0].keys(): + res[k] = rearrange(torch.cat([d[k].clone().unsqueeze(1) for d in dicts], dim=1), "b f c h w -> (b f) c h w ") + return res + + +def dicts_to_device(dicts, device): + ret = [] + for d in dicts: + tmpd = {} + for k in d.keys(): + tmpd[k] = d[k].clone().to(device) + ret.append(tmpd) + return ret \ No newline at end of file diff --git a/hellomeme/tools/sr.py b/hellomeme/tools/sr.py new file mode 100644 index 0000000..702483e --- /dev/null +++ b/hellomeme/tools/sr.py @@ -0,0 +1,394 @@ +# coding: utf-8 + +""" +@File : sr.py +@Author : Songkey +@Email : songkey@pku.edu.cn +@Date : 5/30/2025 +@Desc : adapted from: https://github.com/xinntao/Real-ESRGAN +""" + +import torch +from torch import nn as nn +from torch.nn import functional as F +import cv2 +import numpy as np +import math +import os.path as osp + +def pixel_unshuffle(x, scale): + """ Pixel unshuffle. + + Args: + x (Tensor): Input feature with shape (b, c, hh, hw). + scale (int): Downsample ratio. + + Returns: + Tensor: the pixel unshuffled feature. + """ + b, c, hh, hw = x.size() + out_channel = c * (scale**2) + assert hh % scale == 0 and hw % scale == 0 + h = hh // scale + w = hw // scale + x_view = x.view(b, c, h, scale, w, scale) + return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w) + +def make_layer(basic_block, num_basic_block, **kwarg): + """Make layers by stacking the same blocks. + + Args: + basic_block (nn.module): nn.module class for basic block. + num_basic_block (int): number of blocks. + + Returns: + nn.Sequential: Stacked blocks in nn.Sequential. + """ + layers = [] + for _ in range(num_basic_block): + layers.append(basic_block(**kwarg)) + return nn.Sequential(*layers) + +class ResidualDenseBlock(nn.Module): + """Residual Dense Block. + + Used in RRDB block in ESRGAN. + + Args: + num_feat (int): Channel number of intermediate features. + num_grow_ch (int): Channels for each growth. + """ + + def __init__(self, num_feat=64, num_grow_ch=32): + super(ResidualDenseBlock, self).__init__() + self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) + self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1) + + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + + # initialization + # default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1) + + def forward(self, x): + x1 = self.lrelu(self.conv1(x)) + x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1))) + x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))) + x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1))) + x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) + # Empirically, we use 0.2 to scale the residual for better performance + return x5 * 0.2 + x + + +class RRDB(nn.Module): + """Residual in Residual Dense Block. + + Used in RRDB-Net in ESRGAN. + + Args: + num_feat (int): Channel number of intermediate features. + num_grow_ch (int): Channels for each growth. + """ + + def __init__(self, num_feat, num_grow_ch=32): + super(RRDB, self).__init__() + self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch) + self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch) + self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch) + + def forward(self, x): + out = self.rdb1(x) + out = self.rdb2(out) + out = self.rdb3(out) + # Empirically, we use 0.2 to scale the residual for better performance + return out * 0.2 + x + +class RRDBNet(nn.Module): + """Networks consisting of Residual in Residual Dense Block, which is used + in ESRGAN. + + ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks. + + We extend ESRGAN for scale x2 and scale x1. + Note: This is one option for scale 1, scale 2 in RRDBNet. + We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size + and enlarge the channel size before feeding inputs into the main ESRGAN architecture. + + Args: + num_in_ch (int): Channel number of inputs. + num_out_ch (int): Channel number of outputs. + num_feat (int): Channel number of intermediate features. + Default: 64 + num_block (int): Block number in the trunk network. Defaults: 23 + num_grow_ch (int): Channels for each growth. Default: 32. + """ + + def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32): + super(RRDBNet, self).__init__() + self.scale = scale + if scale == 2: + num_in_ch = num_in_ch * 4 + elif scale == 1: + num_in_ch = num_in_ch * 16 + self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) + self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch) + self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + # upsample + self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + + def forward(self, x): + if self.scale == 2: + feat = pixel_unshuffle(x, scale=2) + elif self.scale == 1: + feat = pixel_unshuffle(x, scale=4) + else: + feat = x + feat = self.conv_first(feat) + body_feat = self.conv_body(self.body(feat)) + feat = feat + body_feat + # upsample + feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest'))) + feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest'))) + out = self.conv_last(self.lrelu(self.conv_hr(feat))) + return out + +class RealESRGANer(): + def __init__(self, + scale, + tile=0, + tile_pad=10, + pre_pad=10, + half=True, + device=None, + gpu_id=None, + modelscope=False): + self.scale = scale + self.tile_size = tile + self.tile_pad = tile_pad + self.pre_pad = pre_pad + self.mod_scale = None + self.half = half + + # initialize model + if gpu_id: + self.device = torch.device( + f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device + else: + self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device + + if modelscope: + from modelscope import snapshot_download + model_path = osp.join(snapshot_download('songkey/ESRGAN'), 'RealESRGAN_x2plus.pth') + else: + from huggingface_hub import hf_hub_download + model_path = hf_hub_download('songkey/ESRGAN', filename='RealESRGAN_x2plus.pth') + + loadnet = torch.load(model_path, map_location=torch.device('cpu')) + + # prefer to use params_ema + if 'params_ema' in loadnet: + keyname = 'params_ema' + else: + keyname = 'params' + self.model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) + self.model.load_state_dict(loadnet[keyname], strict=True) + + self.model.eval() + self.model = self.model.to(self.device) + if self.half: + self.model = self.model.half() + + def dni(self, net_a, net_b, dni_weight, key='params', loc='cpu'): + """Deep network interpolation. + + ``Paper: Deep Network Interpolation for Continuous Imagery Effect Transition`` + """ + net_a = torch.load(net_a, map_location=torch.device(loc)) + net_b = torch.load(net_b, map_location=torch.device(loc)) + for k, v_a in net_a[key].items(): + net_a[key][k] = dni_weight[0] * v_a + dni_weight[1] * net_b[key][k] + return net_a + + def pre_process(self, img): + """Pre-process, such as pre-pad and mod pad, so that the images can be divisible + """ + img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float() + self.img = img.unsqueeze(0).to(self.device) + if self.half: + self.img = self.img.half() + + # pre_pad + if self.pre_pad != 0: + self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect') + # mod pad for divisible borders + if self.scale == 2: + self.mod_scale = 2 + elif self.scale == 1: + self.mod_scale = 4 + if self.mod_scale is not None: + self.mod_pad_h, self.mod_pad_w = 0, 0 + _, _, h, w = self.img.size() + if (h % self.mod_scale != 0): + self.mod_pad_h = (self.mod_scale - h % self.mod_scale) + if (w % self.mod_scale != 0): + self.mod_pad_w = (self.mod_scale - w % self.mod_scale) + self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect') + + def process(self): + # model inference + self.output = self.model(self.img) + + def tile_process(self): + """It will first crop input images to tiles, and then process each tile. + Finally, all the processed tiles are merged into one images. + + Modified from: https://github.com/ata4/esrgan-launcher + """ + batch, channel, height, width = self.img.shape + output_height = height * self.scale + output_width = width * self.scale + output_shape = (batch, channel, output_height, output_width) + + # start with black image + self.output = self.img.new_zeros(output_shape) + tiles_x = math.ceil(width / self.tile_size) + tiles_y = math.ceil(height / self.tile_size) + + # loop over all tiles + for y in range(tiles_y): + for x in range(tiles_x): + # extract tile from input image + ofs_x = x * self.tile_size + ofs_y = y * self.tile_size + # input tile area on total image + input_start_x = ofs_x + input_end_x = min(ofs_x + self.tile_size, width) + input_start_y = ofs_y + input_end_y = min(ofs_y + self.tile_size, height) + + # input tile area on total image with padding + input_start_x_pad = max(input_start_x - self.tile_pad, 0) + input_end_x_pad = min(input_end_x + self.tile_pad, width) + input_start_y_pad = max(input_start_y - self.tile_pad, 0) + input_end_y_pad = min(input_end_y + self.tile_pad, height) + + # input tile dimensions + input_tile_width = input_end_x - input_start_x + input_tile_height = input_end_y - input_start_y + tile_idx = y * tiles_x + x + 1 + input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad] + + # upscale tile + try: + with torch.no_grad(): + output_tile = self.model(input_tile) + except RuntimeError as error: + print('Error', error) + print(f'\tTile {tile_idx}/{tiles_x * tiles_y}') + + # output tile area on total image + output_start_x = input_start_x * self.scale + output_end_x = input_end_x * self.scale + output_start_y = input_start_y * self.scale + output_end_y = input_end_y * self.scale + + # output tile area without padding + output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale + output_end_x_tile = output_start_x_tile + input_tile_width * self.scale + output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale + output_end_y_tile = output_start_y_tile + input_tile_height * self.scale + + # put tile into output image + self.output[:, :, output_start_y:output_end_y, + output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile, + output_start_x_tile:output_end_x_tile] + + def post_process(self): + # remove extra pad + if self.mod_scale is not None: + _, _, h, w = self.output.size() + self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale] + # remove prepad + if self.pre_pad != 0: + _, _, h, w = self.output.size() + self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale] + return self.output + + @torch.no_grad() + def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'): + h_input, w_input = img.shape[0:2] + # img: numpy + img = img.astype(np.float32) + if np.max(img) > 256: # 16-bit image + max_range = 65535 + print('\tInput is a 16-bit image') + else: + max_range = 255 + img = img / max_range + if len(img.shape) == 2: # gray image + img_mode = 'L' + img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) + elif img.shape[2] == 4: # RGBA image with alpha channel + img_mode = 'RGBA' + alpha = img[:, :, 3] + img = img[:, :, 0:3] + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + if alpha_upsampler == 'realesrgan': + alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB) + else: + img_mode = 'RGB' + img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) + + # ------------------- process image (without the alpha channel) ------------------- # + self.pre_process(img) + if self.tile_size > 0: + self.tile_process() + else: + self.process() + output_img = self.post_process() + output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy() + output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0)) + if img_mode == 'L': + output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY) + + # ------------------- process the alpha channel if necessary ------------------- # + if img_mode == 'RGBA': + if alpha_upsampler == 'realesrgan': + self.pre_process(alpha) + if self.tile_size > 0: + self.tile_process() + else: + self.process() + output_alpha = self.post_process() + output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy() + output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0)) + output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY) + else: # use the cv2 resize for alpha channel + h, w = alpha.shape[0:2] + output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR) + + # merge the alpha channel + output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA) + output_img[:, :, 3] = output_alpha + + # ------------------------------ return ------------------------------ # + if max_range == 65535: # 16-bit image + output = (output_img * 65535.0).round().astype(np.uint16) + else: + output = (output_img * 255.0).round().astype(np.uint8) + + if outscale is not None and outscale != float(self.scale): + output = cv2.resize( + output, ( + int(w_input * outscale), + int(h_input * outscale), + ), interpolation=cv2.INTER_LANCZOS4) + + return output, img_mode \ No newline at end of file diff --git a/meme.py b/meme.py index 41892e8..81335fc 100644 --- a/meme.py +++ b/meme.py @@ -1,9 +1,7 @@ import json import os import os.path as osp -import random import sys -import subprocess import importlib.metadata cur_dir = osp.dirname(osp.abspath(__file__)) @@ -11,7 +9,7 @@ cur_dir = osp.dirname(osp.abspath(__file__)) installed_packages = [package.name for package in importlib.metadata.distributions()] REQUIRED = { - 'diffusers':'0.31.0', 'transformers':'4.46.3', 'einops':'0.8.0', 'opencv-python':'4.10.0.84', 'tqdm':'4.67.0', + 'diffusers':'0.33.1', 'transformers':'4.46.3', 'einops':'0.8.0', 'opencv-python':'4.10.0.84', 'tqdm':'4.67.0', 'pillow':'10.2.0', 'onnxruntime-gpu':'1.18.1', 'onnx':'1.17.0', 'safetensors':'0.4.5', 'accelerate':'1.1.1', 'peft':'0.13.2', 'modelscope':'', 'huggingface-hub':'' } @@ -54,7 +52,7 @@ from .hellomeme.utils import (get_drive_expression, ) from .hellomeme import (HMImagePipeline, HMVideoPipeline, HM3ImagePipeline, HM3VideoPipeline, - HM5ImagePipeline) + HM5ImagePipeline, HM5VideoPipeline) config_path = osp.join(cur_dir, 'hellomeme', 'model_config.json') with open(config_path, 'r') as f: @@ -169,7 +167,7 @@ class HMVideoPipelineLoader: "checkpoint": (checkpoint_files, ), "lora": (lora_files, ), "vae": (vae_files, ), - "version": (['v1', 'v2', 'v3', 'v4'], ), + "version": (['v1', 'v2', 'v3', 'v4', 'v5'], ), "stylize": (['x1', 'x2'], ), "deployment": (['huggingface', 'modelscope'], ), "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), @@ -192,6 +190,8 @@ class HMVideoPipelineLoader: if version == 'v3' or version == 'v4': pipeline = HM3VideoPipeline.from_pretrained(sd1_5_dir) + elif version == 'v5': + pipeline = HM5VideoPipeline.from_pretrained(sd1_5_dir) else: pipeline = HMVideoPipeline.from_pretrained(sd1_5_dir) pipeline.to(dtype=dtype) diff --git a/pyproject.toml b/pyproject.toml index aee7bab..e3d39a1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui_hellomeme" description = "This repository is the official implementation of the [a/HelloMeme](https://arxiv.org/pdf/2410.22901) ComfyUI interface" -version = "1.3.3" +version = "1.3.4" license = {file = "LICENSE"} [project.urls]