Update app.py

app.py

@@ -10,6 +10,10 @@ 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",
@@ -17,29 +21,41 @@ bases = {
     "3d": "Lykon/DreamShaper",
     "Anime": "Yntec/mistoonAnime2"
 }
+step_loaded = None
+base_loaded = "Realistic"
+motion_loaded = None
 
-# Preload models
+# Ensure model and scheduler are initialized in GPU-enabled function
 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")
 
-motion_loaded = None
+# 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
+feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")
 
 # Function 
 @spaces.GPU(duration=60,queue=False)
-def generate_image(prompt, base="Realistic", motion="Default", step=8, progress=gr.Progress()):
-    global pipes
+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
 
-    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_loaded != motion:
         pipe.unload_lora_weights()
@@ -48,16 +64,6 @@ def generate_image(prompt, base="Realistic", motion="Default", step=8, progress=
             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("-", "")
@@ -66,7 +72,6 @@ def generate_image(prompt, base="Realistic", motion="Default", step=8, progress=
     return path
 
 
-
 # Gradio Interface
 with gr.Blocks(css="style.css") as demo:
     gr.HTML(
@@ -90,7 +95,7 @@ with gr.Blocks(css="style.css") as demo:
                     "3d",
                     "Anime",
                 ],
-                value="Realistic",
+                value=base_loaded,
                 interactive=True
             )
             select_motion = gr.Dropdown(
@@ -157,7 +162,9 @@ with gr.Blocks(css="style.css") as demo:
         fn=generate_image,
         inputs=[prompt],
         outputs=[video],
-        cache_examples="lazy",
+        cache_examples=True,
 )
 
-demo.queue().launch()
+demo.queue().launch()
+
+Translate