make image style match and fix sizes

app.py

@@ -1,16 +1,21 @@
 import gradio as gr
 import spaces
 import torch
-from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel
+from diffusers import (
+    StableDiffusionXLControlNetPipeline, 
+    StableDiffusionXLControlNetImg2ImgPipeline, 
+    ControlNetModel
+)
 from gradio_client import Client, handle_file
 
 # ==========================================
 # 1. LAZY LOAD LOCAL CARTOON MODELS
 # ==========================================
 image_pipe = None
+img2img_pipe = None
 
 def load_cartoon_models():
-    global image_pipe
+    global image_pipe, img2img_pipe
     if image_pipe is None:
         print("🟢 Loading Animagine XL 3.1 & Scribble ControlNet...")
         dtype = torch.float16
@@ -18,25 +23,32 @@ def load_cartoon_models():
             "xinsir/controlnet-scribble-sdxl-1.0", 
             torch_dtype=dtype
         )
+        # Pipeline 1: Text + Sketch -> Cartoon
         image_pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
             "cagliostrolab/animagine-xl-3.1",
             controlnet=controlnet,
             torch_dtype=dtype
         )
+        
+        # Pipeline 2: Previous Cartoon + New Sketch -> Consistent Cartoon
+        # We share the components so it uses 0 extra VRAM!
+        img2img_pipe = StableDiffusionXLControlNetImg2ImgPipeline(**image_pipe.components)
+        
     return True
 
 # ==========================================
 # 2. LOCAL ZEROGPU GENERATION (Sketch -> Cartoon)
 # ==========================================
 @spaces.GPU(duration=180)
-def generate_cartoons(sketch_1, sketch_2, user_prompt, ctrl_scale):
+def generate_cartoons(sketch_1, sketch_2, user_prompt, ctrl_scale, consistency_strength):
     load_cartoon_models()
     image_pipe.to("cuda")
+    img2img_pipe.to("cuda") # Moving both to GPU
     
     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...")
+    print("🎨 Stylizing Start Sketch (From Scratch)...")
     img_1 = image_pipe(
         prompt=master_prompt, 
         negative_prompt=neg_prompt, 
@@ -46,16 +58,23 @@ def generate_cartoons(sketch_1, sketch_2, user_prompt, ctrl_scale):
         controlnet_conditioning_scale=ctrl_scale
     ).images[0]
     
-    print("🎨 Stylizing End Sketch...")
-    img_2 = image_pipe(
+    print("🎨 Stylizing End Sketch (Inheriting colors from Start Sketch)...")
+    img_2 = img2img_pipe(
         prompt=master_prompt, 
         negative_prompt=neg_prompt, 
-        image=sketch_2, 
+        image=img_1,                 
+        control_image=sketch_2,      
+        strength=consistency_strength, 
         num_inference_steps=25, 
         guidance_scale=7.0, 
         controlnet_conditioning_scale=ctrl_scale
     ).images[0]
     
+    # 🚨 THE FIX: Resize to 512x320 so ToonCrafter doesn't crash!
+    print("📏 Resizing images for ToonCrafter compatibility...")
+    img_1 = img_1.resize((512, 320))
+    img_2 = img_2.resize((512, 320))
+    
     img_1_path, img_2_path = "frame1.png", "frame2.png"
     img_1.save(img_1_path)
     img_2.save(img_2_path)
@@ -71,25 +90,21 @@ def run_tooncrafter(img_1_path, img_2_path, prompt):
     
     try:
         print("🔄 Submitting to ToonCrafter API (/get_image)...")
-        # We now use the EXACT keyword arguments and order required by the API
         result = client.predict(
             image=handle_file(img_1_path),
             prompt=prompt,
-            steps=25,          # Lowered from default 50 to 25 so you don't wait in queue forever
+            steps=25,
             cfg_scale=7.5,
             eta=1.0,
-            fs=10,             # FPS
+            fs=10,
             seed=123,
             image2=handle_file(img_2_path),
             api_name="/get_image"
         )
         print("✅ ToonCrafter Generation Complete!")
         
-        # The API documentation says it returns a Dict: {video: filepath, subtitles: None}
         if isinstance(result, dict) and 'video' in result:
             return result['video']
-        
-        # Fallback just in case they return the raw string
         return result
         
     except Exception as e:
@@ -98,21 +113,17 @@ def run_tooncrafter(img_1_path, img_2_path, prompt):
 # ==========================================
 # 4. MASTER PIPELINE CONTROLLER
 # ==========================================
-def process_full_animation(sketch_1, sketch_2, prompt, ctrl_scale):
-    # Step 1: Run local ZeroGPU (Cartoons)
-    img_1_path, img_2_path = generate_cartoons(sketch_1, sketch_2, prompt, ctrl_scale)
-    
-    # Step 2: Pass to remote API (Video)
+def process_full_animation(sketch_1, sketch_2, prompt, ctrl_scale, consistency_strength):
+    img_1_path, img_2_path = generate_cartoons(sketch_1, sketch_2, prompt, ctrl_scale, consistency_strength)
     video_path = run_tooncrafter(img_1_path, img_2_path, prompt)
-    
     return img_1_path, img_2_path, video_path
 
 # ==========================================
 # 5. GRADIO INTERFACE
 # ==========================================
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# ✏️ Sketch-to-ToonCrafter 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!")
+    gr.Markdown("# ✏️ Sketch-to-ToonCrafter Studio (Pro Workflow)")
+    gr.Markdown("Upload two sketches. We use **Img2Img ControlNet** to ensure the characters maintain their exact clothing and colors, then pass them to **ToonCrafter** to animate!")
     
     with gr.Row():
         with gr.Column():
@@ -122,18 +133,23 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 label="Character & Motion Description", 
                 placeholder="e.g., A boy in a red shirt jumping"
             )
-            ctrl_scale = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.05, label="Sketch Adherence Strength")
+            
+            with gr.Accordion("Advanced AI Settings", open=False):
+                ctrl_scale = gr.Slider(minimum=0.0, maximum=2.0, value=1.0, step=0.05, label="Sketch Strictness (How closely to follow your lines)")
+                # New slider for Image-to-Image strength
+                consistency_strength = gr.Slider(minimum=0.5, maximum=1.0, value=0.85, step=0.05, label="Color & Style Consistency (Lower = More like Frame 1, Higher = More creative)")
+                
             generate_btn = gr.Button("Create Animation", variant="primary")
             
         with gr.Column():
             with gr.Row():
                 out_img_1 = gr.Image(label="Animagine Start Frame")
-                out_img_2 = gr.Image(label="Animagine End Frame")
+                out_img_2 = gr.Image(label="Animagine End Frame (Color Matched)")
             out_video = gr.Video(label="ToonCrafter Animated Video")
             
     generate_btn.click(
         fn=process_full_animation,
-        inputs=[sketch_1, sketch_2, prompt, ctrl_scale],
+        inputs=[sketch_1, sketch_2, prompt, ctrl_scale, consistency_strength],
         outputs=[out_img_1, out_img_2, out_video]
     )