Update app.py

app.py

@@ -15,6 +15,8 @@ from tqdm import tqdm
 import cv2
 import numpy as np
 import torch
+import torch._dynamo
+from huggingface_hub import list_models
 from torch.nn import functional as F
 from PIL import Image
 
@@ -231,9 +233,30 @@ def interpolate_bits(frames_np, multiplier=2, scale=1.0):
 
 # WAN
 
-MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+ORG_NAME = "TestOrganizationPleaseIgnore"
+# MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
+MODEL_ID = os.getenv("REPO_ID") or random.choice(
+    list(list_models(author=ORG_NAME, filter='diffusers:WanImageToVideoPipeline'))
+).modelId
 CACHE_DIR = os.path.expanduser("~/.cache/huggingface/")
 
+LORA_MODELS = [
+    # {
+    #     "repo_id": "exampleuser/example_lora_1",
+    #     "high_tr": "example_lora_1_high.safetensors",
+    #     "low_tr": "example_lora_1_low.safetensors",
+    #     "high_scale": 0.5,
+    #     "low_scale": 0.5
+    # },
+    # {
+    #     "repo_id": "exampleuser/example_lora_2",
+    #     "high_tr": "subfolder/example_lora_2_high.safetensors",
+    #     "low_tr": "subfolder/example_lora_2_low.safetensors",
+    #     "high_scale": 0.4,
+    #     "low_scale": 0.4
+    # },
+]
+
 MAX_DIM = 832
 MIN_DIM = 480
 SQUARE_DIM = 640
@@ -258,11 +281,43 @@ SCHEDULER_MAP = {
 }
 
 pipe = WanImageToVideoPipeline.from_pretrained(
-    "TestOrganizationPleaseIgnore/WAMU_v1_WAN2.2_I2V_LIGHTNING",
+    MODEL_ID,
     torch_dtype=torch.bfloat16,
 ).to('cuda')
 original_scheduler = copy.deepcopy(pipe.scheduler)
 
+for i, lora in enumerate(LORA_MODELS):
+    name_high_tr = lora["high_tr"].split(".")[0].split("/")[-1] + "Hh"
+    name_low_tr = lora["low_tr"].split(".")[0].split("/")[-1] + "Ll"
+    
+    try: 
+        pipe.load_lora_weights(
+            lora["repo_id"],
+            weight_name=lora["high_tr"],
+            adapter_name=name_high_tr
+        )
+    
+        kwargs_lora = {"load_into_transformer_2": True}
+        pipe.load_lora_weights(
+            lora["repo_id"],
+            weight_name=lora["low_tr"],
+            adapter_name=name_low_tr,
+            **kwargs_lora
+        )
+    
+        pipe.set_adapters([name_high_tr, name_low_tr], adapter_weights=[1.0, 1.0])
+    
+        pipe.fuse_lora(adapter_names=[name_high_tr], lora_scale=lora["high_scale"], components=["transformer"])
+        pipe.fuse_lora(adapter_names=[name_low_tr], lora_scale=lora["low_scale"], components=["transformer_2"])
+    
+        pipe.unload_lora_weights()
+
+        print(f"Applied: {lora['high_tr']}, hs={lora['high_scale']}/ls={lora['low_scale']}, {i+1}/{len(LORA_MODELS)}") 
+    except Exception as e:
+        print("Error:", str(e))
+        print("Failed LoRA:", name_high_tr)
+        pipe.unload_lora_weights()
+
 if os.path.exists(CACHE_DIR):
     shutil.rmtree(CACHE_DIR)
     print("Deleted Hugging Face cache.")
@@ -270,8 +325,11 @@ else:
     print("No hub cache found.")
 
 quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
+torch._dynamo.reset()
 quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
+torch._dynamo.reset()
 quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())
+torch._dynamo.reset()
 
 aoti.aoti_blocks_load(pipe.transformer, 'zerogpu-aoti/Wan2', variant='fp8da')
 aoti.aoti_blocks_load(pipe.transformer_2, 'zerogpu-aoti/Wan2', variant='fp8da')
@@ -283,6 +341,12 @@ default_prompt_i2v = "make this image come alive, cinematic motion, smooth anima
 default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
 
 
+def model_title():
+    repo_name = MODEL_ID.split('/')[-1].replace("_", " ")
+    url = f"https://huggingface.co/{MODEL_ID}"
+    return f"## This space is currently running [{repo_name}]({url}) 🐢"
+
+
 def resize_image(image: Image.Image) -> Image.Image:
     width, height = image.size
     if width == height:
@@ -359,7 +423,7 @@ def get_inference_duration(
     factor = num_frames * width * height / BASE_FRAMES_HEIGHT_WIDTH
     step_duration = BASE_STEP_DURATION * factor ** 1.5
     gen_time = int(steps) * step_duration
-    print(gen_time)
+
     if guidance_scale > 1:
         gen_time = gen_time * 1.8
 
@@ -367,10 +431,8 @@ def get_inference_duration(
     if frame_factor > 1:
         total_out_frames = (num_frames * frame_factor) - num_frames
         inter_time = (total_out_frames * 0.02)
-        print(inter_time)
         gen_time += inter_time
 
-    print("Time GPU", gen_time + 10)
     return 10 + gen_time
 
 
@@ -562,7 +624,7 @@ CSS = """
 
 
 with gr.Blocks(delete_cache=(3600, 10800)) as demo:
-    gr.Markdown("## WAMU - Wan 2.2 I2V (14B) 🐢")
+    gr.Markdown(model_title())
     gr.Markdown("#### ℹ️ **A Note on Performance:** This version prioritizes a straightforward setup over maximum speed, so performance may vary.")
     gr.Markdown("Run Wan 2.2 in just 4-8 steps, fp8 quantization & AoT compilation - compatible with 🧨 diffusers and ZeroGPU")
 
@@ -594,8 +656,11 @@ with gr.Blocks(delete_cache=(3600, 10800)) as demo:
                 )
                 flow_shift_slider = gr.Slider(minimum=0.5, maximum=15.0, step=0.1, value=3.0, label="Flow Shift")
                 play_result_video = gr.Checkbox(label="Display result", value=True, interactive=True)
-                org_name = "TestOrganizationPleaseIgnore"
-                gr.Markdown(f"[ZeroGPU help, tips and troubleshooting](https://huggingface.co/datasets/{org_name}/help/blob/main/gpu_help.md)")
+                gr.Markdown(f"[ZeroGPU help, tips and troubleshooting](https://huggingface.co/datasets/{ORG_NAME}/help/blob/main/gpu_help.md)")
+                gr.Markdown(  # TestOrganizationPleaseIgnore/wamu-tools
+                    "To use a different model, **duplicate this Space** first, then change the `REPO_ID` environment variable. "
+                    "[See compatible models here](https://huggingface.co/models?other=diffusers:WanImageToVideoPipeline&sort=trending&search=WAN2.2_I2V_LIGHTNING)."
+                )
 
             generate_button = gr.Button("Generate Video", variant="primary")