update to use tooncrafter

app.py

@@ -1,93 +1,135 @@
 import gradio as gr
-import replicate
+import spaces
+import torch
+from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from gradio_client import Client, handle_file
 import os
 
-def process_animation(sketch_1_path, sketch_2_path, prompt):
-    # 1. Check for API Token
-    if not os.environ.get("REPLICATE_API_TOKEN"):
-        raise gr.Error("REPLICATE_API_TOKEN is missing! Add it to your Space Secrets.")
-        
-    try:
-        # ==========================================
-        # STEP 1: Stylize the First Sketch
-        # ==========================================
-        print("🎨 Step 1: Stylizing Start Sketch...")
-        img1_output = replicate.run(
-            "jagilley/controlnet-scribble", # Replicate's standard Sketch-to-Image model
-            input={
-                "image": open(sketch_1_path, "rb"),
-                "prompt": prompt,
-                "a_prompt": "best quality, highly detailed cartoon, colorful, professional animation",
-                "n_prompt": "photorealistic, ugly, bad anatomy, grayscale, messy lines"
-            }
+# ==========================================
+# 1. LAZY LOAD LOCAL CARTOON MODELS
+# ==========================================
+image_pipe = None
+
+def load_cartoon_models():
+    global image_pipe
+    if image_pipe is None:
+        print("🟢 Loading Animagine XL 3.1 & Scribble ControlNet...")
+        dtype = torch.float16
+        controlnet = ControlNetModel.from_pretrained(
+            "xinsir/controlnet-scribble-sdxl-1.0", 
+            torch_dtype=dtype
         )
-        # ControlNet returns a list: [control_map, final_image]. We grab the final image.
-        url_1 = img1_output[-1] 
-        
-        # ==========================================
-        # STEP 2: Stylize the Second Sketch
-        # ==========================================
-        print("🎨 Step 2: Stylizing End Sketch...")
-        img2_output = replicate.run(
-            "jagilley/controlnet-scribble",
-            input={
-                "image": open(sketch_2_path, "rb"),
-                "prompt": prompt,
-                "a_prompt": "best quality, highly detailed cartoon, colorful, professional animation",
-                "n_prompt": "photorealistic, ugly, bad anatomy, grayscale, messy lines"
-            }
+        # Animagine XL 3.1 is the premier model for 2D animation/cartoon styles
+        image_pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
+            "cagliostrolab/animagine-xl-3.1",
+            controlnet=controlnet,
+            torch_dtype=dtype
         )
-        url_2 = img2_output[-1]
-        
-        # ==========================================
-        # STEP 3: Generate Video Interpolation
-        # ==========================================
-        print("🎬 Step 3: Generating Video with ToonCrafter...")
-        video_output = replicate.run(
-            "fofr/tooncrafter",
-            input={
-                # We pass the newly generated Image URLs directly to ToonCrafter!
-                "image_1": url_1, 
-                "image_2": url_2,
-                "prompt": prompt,
-                "loop": False
-            }
+    return True
+
+# ==========================================
+# 2. LOCAL ZEROGPU GENERATION (Sketch -> Cartoon)
+# ==========================================
+@spaces.GPU(duration=180)
+def generate_cartoons(sketch_1, sketch_2, user_prompt, ctrl_scale):
+    load_cartoon_models()
+    image_pipe.to("cuda")
+    
+    # Animagine responds best to these specific quality tags
+    master_prompt = f"{user_prompt}, masterpiece, best quality, highly detailed, professional 2d animation, flat colors, anime style"
+    neg_prompt = "nsfw, photorealistic, 3d render, ugly, messy lines, bad anatomy, bad hands, missing fingers, lowres, worst quality"
+    
+    print("🎨 Stylizing Start Sketch...")
+    img_1 = image_pipe(
+        prompt=master_prompt, 
+        negative_prompt=neg_prompt, 
+        image=sketch_1, 
+        num_inference_steps=25, 
+        guidance_scale=7.0, 
+        controlnet_conditioning_scale=ctrl_scale
+    ).images[0]
+    
+    print("🎨 Stylizing End Sketch...")
+    img_2 = image_pipe(
+        prompt=master_prompt, 
+        negative_prompt=neg_prompt, 
+        image=sketch_2, 
+        num_inference_steps=25, 
+        guidance_scale=7.0, 
+        controlnet_conditioning_scale=ctrl_scale
+    ).images[0]
+    
+    image_pipe.to("cpu")
+    
+    # Save temporarily to pass to the ToonCrafter API
+    img_1_path, img_2_path = "frame1.png", "frame2.png"
+    img_1.save(img_1_path)
+    img_2.save(img_2_path)
+    
+    return img_1_path, img_2_path
+
+# ==========================================
+# 3. REMOTE API CALL (Cartoon -> Video)
+# ==========================================
+def run_tooncrafter(img_1_path, img_2_path, prompt):
+    print("🎬 Sending frames to official ToonCrafter Space...")
+    try:
+        # Connect silently to the free ToonCrafter HF Space
+        client = Client("Doubiiu/tooncrafter")
+        result = client.predict(
+            image0=handle_file(img_1_path),
+            image1=handle_file(img_2_path),
+            prompt=prompt,
+            seed=123,
+            eta=1.0,
+            cfg_scale=7.5,
+            steps=15,
+            frame_num=16,
+            api_name="/predict"
         )
-        print("✅ Generation Complete!")
-        
-        # Return all three outputs to the UI
-        return url_1, url_2, video_output[0] 
-        
+        print("✅ Animation Complete!")
+        # The result returns a tuple, the video path is the second item [1]
+        return result[1] 
     except Exception as e:
-        raise gr.Error(f"API Error: {str(e)}")
+        raise gr.Error(f"ToonCrafter API Error: {str(e)}")
+
+# ==========================================
+# 4. MASTER PIPELINE CONTROLLER
+# ==========================================
+def process_full_animation(sketch_1, sketch_2, prompt, ctrl_scale):
+    # Step 1: Run local ZeroGPU
+    img_1_path, img_2_path = generate_cartoons(sketch_1, sketch_2, prompt, ctrl_scale)
+    # Step 2: Pass to ToonCrafter
+    video_path = run_tooncrafter(img_1_path, img_2_path, prompt)
+    return img_1_path, img_2_path, video_path
 
 # ==========================================
-# GRADIO INTERFACE
+# 5. GRADIO INTERFACE
 # ==========================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 Sketch-to-Motion Studio")
-    gr.Markdown("1. Upload two rough sketches. 2. We turn them into cartoons. 3. We animate the transition!")
+    gr.Markdown("# ✏️ Sketch-to-ToonCrafter Animation Studio")
+    gr.Markdown("Upload two sketches. We use **Animagine XL 3.1** locally to turn them into professional cartoons, and pass them to **ToonCrafter** to animate the motion!")
     
     with gr.Row():
         with gr.Column():
-            sketch_1 = gr.Image(type="filepath", label="Start Sketch (Black & White lines)")
-            sketch_2 = gr.Image(type="filepath", label="End Sketch (Black & White lines)")
+            sketch_1 = gr.Image(type="pil", label="Start Sketch (Black & White lines)", image_mode="RGB")
+            sketch_2 = gr.Image(type="pil", label="End Sketch (Black & White lines)", image_mode="RGB")
             prompt = gr.Textbox(
-                label="Cartoon Description", 
-                placeholder="e.g., A cute golden retriever dog running in a grassy field"
+                label="Character & Motion Description", 
+                placeholder="e.g., A boy in a red shirt and blue jeans jumping"
             )
+            ctrl_scale = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.05, label="Sketch Adherence Strength")
             generate_btn = gr.Button("Create Animation", variant="primary")
             
         with gr.Column():
             with gr.Row():
-                out_img_1 = gr.Image(label="AI Stylized Start")
-                out_img_2 = gr.Image(label="AI Stylized End")
-            out_video = gr.Video(label="Final Animated Video")
+                out_img_1 = gr.Image(label="Animagine Start Frame")
+                out_img_2 = gr.Image(label="Animagine End Frame")
+            out_video = gr.Video(label="ToonCrafter Animated Video")
             
-    # Link the UI to the function
     generate_btn.click(
-        fn=process_animation,
-        inputs=[sketch_1, sketch_2, prompt],
+        fn=process_full_animation,
+        inputs=[sketch_1, sketch_2, prompt, ctrl_scale],
         outputs=[out_img_1, out_img_2, out_video]
     )
 

requirements.txt

@@ -1,2 +1,7 @@
 gradio
-replicate
+spaces
+torch
+diffusers>=0.33.0
+transformers
+accelerate
+gradio_client