Update app.py

app.py

@@ -10,10 +10,6 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from PIL import Image
 
-MORE = """ ## TRY Other Demos
-        ### Instant Image: 4k images in 5 Second -> https://huggingface.co/spaces/KingNish/Instant-Image
-        """
-
 # Constants
 bases = {
     "Cartoon": "frankjoshua/toonyou_beta6",
@@ -21,49 +17,63 @@ bases = {
     "3d": "Lykon/DreamShaper",
     "Anime": "Yntec/mistoonAnime2"
 }
-step_loaded = None
-base_loaded = "Realistic"
-motion_loaded = None
 
-# Ensure model and scheduler are initialized in GPU-enabled function
+motion_models = {
+    "Default": None,
+    "Zoom in": "guoyww/animatediff-motion-lora-zoom-in",
+    "Zoom out": "guoyww/animatediff-motion-lora-zoom-out",
+    "Tilt up": "guoyww/animatediff-motion-lora-tilt-up",
+    "Tilt down": "guoyww/animatediff-motion-lora-tilt-down",
+    "Pan left": "guoyww/animatediff-motion-lora-pan-left",
+    "Pan right": "guoyww/animatediff-motion-lora-pan-right",
+    "Roll left": "guoyww/animatediff-motion-lora-rolling-anticlockwise",
+    "Roll right": "guoyww/animatediff-motion-lora-rolling-clockwise",
+}
+
+
+# Preload models
 if not torch.cuda.is_available():
     raise NotImplementedError("No GPU detected!")
 
 device = "cuda"
 dtype = torch.float16
-pipe = AnimateDiffPipeline.from_pretrained(bases[base_loaded], torch_dtype=dtype).to(device)
-pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
 
-# Safety checkers
-from transformers import CLIPFeatureExtractor
+pipes = {}
+for base_name, base_path in bases.items():
+    pipe = AnimateDiffPipeline.from_pretrained(base_path, torch_dtype=dtype).to(device)
+    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", beta_schedule="linear")
+    pipes[base_name] = pipe
+
+# Load motion models
+for motion_name, motion_path in motion_models.items():
+    if motion_path:
+        motion_model = MotionAdapter.from_pretrained(motion_path, torch_dtype=dtype).to(device)
+        motion_models[motion_name] = motion_model
 
-feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
 
 # Function 
 @spaces.GPU(duration=60,queue=False)
-def generate_image(prompt, base="Realistic", motion="", step=8, progress=gr.Progress()):
-    global step_loaded
-    global base_loaded
-    global motion_loaded
-    print(prompt, base, step)
-
-    if step_loaded != step:
-        repo = "ByteDance/AnimateDiff-Lightning"
-        ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
-        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
-        step_loaded = step
+def generate_image(prompt, base="Realistic", motion="Default", step=8, progress=gr.Progress()):
+    global pipes
+    global motion_models
+
+    pipe = pipes[base]
 
-    if base_loaded != base:
-        pipe.unet.load_state_dict(torch.load(hf_hub_download(bases[base], "unet/diffusion_pytorch_model.bin"), map_location=device), strict=False)
-        base_loaded = base
+    if motion != "Default":
+        pipe.motion_adapter = motion_models[motion]
+    else:
+        pipe.motion_adapter = None
 
-    if motion_loaded != motion:
-        pipe.unload_lora_weights()
-        if motion != "":
-            pipe.load_lora_weights(motion, adapter_name="motion")
-            pipe.set_adapters(["motion"], [0.7])
-        motion_loaded = motion
 
+    # Load step model if not already loaded
+    repo = "ByteDance/AnimateDiff-Lightning"
+    ckpt = f"animatediff_lightning_{step}step_diffusers.safetensors"
+    try:
+        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt, local_files_only=True), device=device), strict=False)
+    except:
+        pipe.unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device=device), strict=False)
+
+    # Generate image
     output = pipe(prompt=f"{base} image of {prompt}", guidance_scale=1.2, num_inference_steps=step)
 
     name = str(uuid.uuid4()).replace("-", "")
@@ -72,6 +82,7 @@ def generate_image(prompt, base="Realistic", motion="", step=8, progress=gr.Prog
     return path
 
 
+
 # Gradio Interface
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(
@@ -162,7 +173,7 @@ with gr.Blocks(css="style.css") as demo:
         fn=generate_image,
         inputs=[prompt],
         outputs=[video],
-        cache_examples=True,
+        cache_examples="lazy",
 )
 
 demo.queue().launch()