update 11.22

app.py

@@ -0,0 +1,730 @@
+import spaces
+import gradio as gr
+import torch
+import cv2
+import numpy as np
+import mediapipe as mp
+import matplotlib.pyplot as plt
+from PIL import Image
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, StableDiffusionControlNetInpaintPipeline
+from transformers import AutoTokenizer
+import base64
+import requests
+import json
+from rembg import remove
+from scipy import ndimage
+from moviepy.editor import ImageSequenceClip
+from tqdm import tqdm
+import os
+import shutil
+import time
+from huggingface_hub import snapshot_download
+import subprocess
+import sys
+
+
+@spaces.GPU(duration=120)
+def download_liveportrait():
+    """
+    Clone the LivePortrait repository and prepare its dependencies.
+    """
+    liveportrait_path = "./LivePortrait"
+    try:
+        if not os.path.exists(liveportrait_path):
+            print("Cloning LivePortrait repository...")
+            os.system(f"git clone https://github.com/KwaiVGI/LivePortrait.git {liveportrait_path}")
+        
+        # 安装依赖
+        os.chdir(liveportrait_path)
+        print("Installing LivePortrait dependencies...")
+        os.system("pip install -r requirements.txt")
+        
+        # 构建 MultiScaleDeformableAttention 模块
+        dependency_path = "src/utils/dependencies/XPose/models/UniPose/ops"
+        os.chdir(dependency_path)
+        print("Building MultiScaleDeformableAttention...")
+        os.system("python setup.py build")
+        os.system("python setup.py install")
+        
+        # 确保模块路径可用
+        module_path = os.path.abspath(dependency_path)
+        if module_path not in sys.path:
+            sys.path.append(module_path)
+        
+        # 返回 LivePortrait 目录
+        os.chdir("../../../../../../../")
+        print("LivePortrait setup completed")
+    except Exception as e:
+        print("Failed to initialize LivePortrait:", e)
+        raise
+download_liveportrait()
+
+@spaces.GPU(duration=120)
+def download_huggingface_resources():
+    """
+    Download additional necessary resources from Hugging Face using the CLI.
+    """
+    try:
+        local_dir = "./pretrained_weights"
+        os.makedirs(local_dir, exist_ok=True)
+
+        # Use the Hugging Face CLI for downloading
+        cmd = [
+            "huggingface-cli", "download",
+            "KwaiVGI/LivePortrait",
+            "--local-dir", local_dir,
+            "--exclude", "*.git*", "README.md", "docs"
+        ]
+        print("Executing command:", " ".join(cmd))
+        subprocess.run(cmd, check=True)
+
+        print("Resources successfully downloaded to:", local_dir)
+    except subprocess.CalledProcessError as e:
+        print("Error during Hugging Face CLI download:", e)
+        raise
+    except Exception as e:
+        print("General error in downloading resources:", e)
+        raise
+
+download_huggingface_resources()
+
+
+@spaces.GPU(duration=120)
+def get_project_root():
+    """Get the root directory of the current project."""
+    return os.path.abspath(os.path.dirname(__file__))
+
+# Ensure working directory is project root
+os.chdir(get_project_root())
+
+# Initialize the necessary models and components
+mp_pose = mp.solutions.pose
+mp_drawing = mp.solutions.drawing_utils
+
+# Load ControlNet model
+controlnet = ControlNetModel.from_pretrained('lllyasviel/sd-controlnet-openpose', torch_dtype=torch.float16)
+
+# Load Stable Diffusion model with ControlNet
+pipe_controlnet = StableDiffusionControlNetPipeline.from_pretrained(
+    'runwayml/stable-diffusion-v1-5',
+    controlnet=controlnet,
+    torch_dtype=torch.float16
+)
+
+# Load Inpaint Controlnet
+pipe_inpaint_controlnet = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+    "runwayml/stable-diffusion-inpainting",
+    controlnet=controlnet,
+    torch_dtype=torch.float16
+)
+
+# Move to GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+pipe_controlnet.to(device)
+pipe_controlnet.enable_attention_slicing()
+pipe_inpaint_controlnet.to(device)
+pipe_inpaint_controlnet.enable_attention_slicing()
+
+
+@spaces.GPU(duration=120)
+def resize_to_multiple_of_64(width, height):
+    return (width // 64) * 64, (height // 64) * 64
+
+
+@spaces.GPU(duration=120)
+def expand_mask(mask, kernel_size):
+    mask_array = np.array(mask)
+    structuring_element = np.ones((kernel_size, kernel_size), dtype=np.uint8)
+    expanded_mask_array = ndimage.binary_dilation(
+        mask_array, structure=structuring_element
+    ).astype(np.uint8) * 255
+    return Image.fromarray(expanded_mask_array)
+
+
+@spaces.GPU(duration=120)
+def crop_face_to_square(image_rgb, padding_ratio=0.2):
+    """
+    Detects the face in the input image and crops an enlarged square region around it.
+    """
+    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    gray_image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY)
+    faces = face_cascade.detectMultiScale(gray_image, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
+
+    if len(faces) == 0:
+        print("No face detected.")
+        return None
+
+    x, y, w, h = faces[0]
+    center_x, center_y = x + w // 2, y + h // 2
+    side_length = max(w, h)
+    padded_side_length = int(side_length * (1 + padding_ratio))
+    half_side = padded_side_length // 2
+
+    top_left_x = max(center_x - half_side, 0)
+    top_left_y = max(center_y - half_side, 0)
+    bottom_right_x = min(center_x + half_side, image_rgb.shape[1])
+    bottom_right_y = min(center_y + half_side, image_rgb.shape[0])
+
+    cropped_image = image_rgb[top_left_y:bottom_right_y, top_left_x:bottom_right_x]
+    resized_image = cv2.resize(cropped_image, (768, 768), interpolation=cv2.INTER_AREA)
+
+    return resized_image
+
+
+@spaces.GPU(duration=120)
+def spirit_animal_baseline(image_path, num_images = 4):
+
+    image = cv2.imread(image_path)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    image_rgb = crop_face_to_square(image_rgb)
+
+    original_height, original_width, _ = image_rgb.shape
+    aspect_ratio = original_width / original_height
+
+    if aspect_ratio > 1:
+        gen_width = 768
+        gen_height = int(gen_width / aspect_ratio)
+    else:
+        gen_height = 768
+        gen_width = int(gen_height * aspect_ratio)
+
+    gen_width, gen_height = resize_to_multiple_of_64(gen_width, gen_height)
+
+    with mp_pose.Pose(static_image_mode=True) as pose:
+        results = pose.process(image_rgb)
+
+        if results.pose_landmarks:
+            annotated_image = image_rgb.copy()
+            mp_drawing.draw_landmarks(
+                annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS
+            )
+        else:
+            print("No pose detected.")
+            return "No pose detected.", []
+
+    pose_image = np.zeros_like(image_rgb)
+    for connection in mp_pose.POSE_CONNECTIONS:
+        start_idx, end_idx = connection
+        start, end = results.pose_landmarks.landmark[start_idx], results.pose_landmarks.landmark[end_idx]
+        if start.visibility > 0.5 and end.visibility > 0.5:
+            x1, y1 = int(start.x * pose_image.shape[1]), int(start.y * pose_image.shape[0])
+            x2, y2 = int(end.x * pose_image.shape[1]), int(end.y * pose_image.shape[0])
+            cv2.line(pose_image, (x1, y1), (x2, y2), (255, 255, 255), 2)
+
+    pose_pil = Image.fromarray(cv2.resize(pose_image, (gen_width, gen_height), interpolation=cv2.INTER_LANCZOS4))
+
+    base64_image = base64.b64encode(cv2.imencode('.jpg', image_rgb)[1]).decode()
+    api_key = "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"
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+    payload = {
+        "model": "gpt-4o-mini",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Based on the provided image, think of one spirit animal that is right for the person, and answer in the following format: An ultra-realistic, highly detailed photograph of a single {animal} with facial features characterized by {description}, standing upright in a human-like pose, looking directly at the camera, against a solid, neutral background. Generate one sentence without any other responses or numbering."},
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+                ]
+            }
+        ],
+        "max_tokens": 100
+    }
+
+    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    prompt = response.json()['choices'][0]['message']['content'] if 'choices' in response.json() else "A majestic animal"
+
+    num_images = num_images
+    generated_images = []
+    with torch.no_grad():
+        with torch.autocast(device_type=device.type):
+            for _ in range(num_images):
+                images = pipe_controlnet(
+                    prompt=prompt,
+                    negative_prompt="multiple heads, extra limbs, duplicate faces, mutated anatomy, disfigured, blurry",
+                    num_inference_steps=20,
+                    image=pose_pil,
+                    guidance_scale=5,
+                    width=gen_width,
+                    height=gen_height,
+                ).images
+                generated_images.append(images[0])
+
+    return prompt, generated_images
+
+
+@spaces.GPU(duration=120)
+def spirit_animal_with_background(image_path, num_images = 4):
+
+    image = cv2.imread(image_path)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    # image_rgb = crop_face_to_square(image_rgb)
+
+    original_height, original_width, _ = image_rgb.shape
+    aspect_ratio = original_width / original_height
+
+    if aspect_ratio > 1:
+        gen_width = 768
+        gen_height = int(gen_width / aspect_ratio)
+    else:
+        gen_height = 768
+        gen_width = int(gen_height * aspect_ratio)
+
+    gen_width, gen_height = resize_to_multiple_of_64(gen_width, gen_height)
+
+    with mp_pose.Pose(static_image_mode=True) as pose:
+        results = pose.process(image_rgb)
+
+        if results.pose_landmarks:
+            annotated_image = image_rgb.copy()
+            mp_drawing.draw_landmarks(
+                annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS
+            )
+        else:
+            print("No pose detected.")
+            return "No pose detected.", []
+
+    pose_image = np.zeros_like(image_rgb)
+    for connection in mp_pose.POSE_CONNECTIONS:
+        start_idx, end_idx = connection
+        start, end = results.pose_landmarks.landmark[start_idx], results.pose_landmarks.landmark[end_idx]
+        if start.visibility > 0.5 and end.visibility > 0.5:
+            x1, y1 = int(start.x * pose_image.shape[1]), int(start.y * pose_image.shape[0])
+            x2, y2 = int(end.x * pose_image.shape[1]), int(end.y * pose_image.shape[0])
+            cv2.line(pose_image, (x1, y1), (x2, y2), (255, 255, 255), 2)
+
+    pose_pil = Image.fromarray(cv2.resize(pose_image, (gen_width, gen_height), interpolation=cv2.INTER_LANCZOS4))
+
+    base64_image = base64.b64encode(cv2.imencode('.jpg', image_rgb)[1]).decode()
+    api_key = "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"
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+    payload = {
+        "model": "gpt-4o-mini",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Based on the provided image, think of one spirit animal that is right for the person, and answer in the following format: An ultra-realistic, highly detailed photograph of a single {animal} with facial features characterized by {description}, standing upright in a human-like pose, looking directly at the camera, against a solid, neutral background. Generate one sentence without any other responses or numbering."},
+                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
+                ]
+            }
+        ],
+        "max_tokens": 100
+    }
+
+    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    prompt = response.json()['choices'][0]['message']['content'] if 'choices' in response.json() else "A majestic animal"
+
+    mask_image = remove(Image.fromarray(image_rgb))
+    initial_mask = mask_image.split()[-1].convert('L')
+
+    kernel_size = min(gen_width, gen_height) // 15
+    expanded_mask = expand_mask(initial_mask, kernel_size)
+
+    num_images = num_images
+    generated_images = []
+    with torch.no_grad():
+        with torch.autocast(device_type=device.type):
+            for _ in range(num_images):
+                images = pipe_inpaint_controlnet(
+                    prompt=prompt,
+                    negative_prompt="multiple heads, extra limbs, duplicate faces, mutated anatomy, disfigured, blurry",
+                    num_inference_steps=20,
+                    image=Image.fromarray(image_rgb),
+                    mask_image=expanded_mask,
+                    control_image=pose_pil,
+                    width=gen_width,
+                    height=gen_height,
+                    guidance_scale=5,
+                ).images
+                generated_images.append(images[0])
+
+    return prompt, generated_images
+
+
+@spaces.GPU(duration=120)
+def generate_multiple_animals(image_path, keep_background=True, num_images = 4):
+
+    image = cv2.imread(image_path)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    image_rgb = crop_face_to_square(image_rgb)
+
+    original_image = Image.fromarray(image_rgb)
+    original_width, original_height = original_image.size
+
+    aspect_ratio = original_width / original_height
+    if aspect_ratio > 1:
+        gen_width = 768
+        gen_height = int(gen_width / aspect_ratio)
+    else:
+        gen_height = 768
+        gen_width = int(gen_height * aspect_ratio)
+
+    gen_width, gen_height = resize_to_multiple_of_64(gen_width, gen_height)
+
+    base64_image = base64.b64encode(cv2.imencode('.jpg', image_rgb)[1]).decode()
+    api_key = "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"
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+    payload = {
+        "model": "gpt-4o-mini",
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "Based on the provided image, think of " + str(num_images) + " different spirit animals that are right for the person, and answer in the following format for each: An ultra-realistic, highly detailed photograph of a {animal} with facial features characterized by {description}, standing upright in a human-like pose, looking directly at the camera, against a solid, neutral background. Generate these sentences without any other responses or numbering. For the animal choose between owl, bear, fox, koala, lion, dog"
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}
+                    }
+                ]
+            }
+        ],
+        "max_tokens": 500
+    }
+
+    response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+    response_json = response.json()
+
+    if 'choices' in response_json and len(response_json['choices']) > 0:
+      content = response_json['choices'][0]['message']['content']
+      prompts = [prompt.strip() for prompt in content.strip().split('.') if prompt.strip()]
+      negative_prompt = (
+          "multiple heads, extra limbs, duplicate faces, mutated anatomy, disfigured, "
+          "blurry, deformed, text, watermark, logo, low resolution"
+      )
+      formatted_prompts = "\n".join(f"{i+1}. {prompt}" for i, prompt in enumerate(prompts))
+
+    with mp_pose.Pose(static_image_mode=True) as pose:
+        results = pose.process(image_rgb)
+
+        if results.pose_landmarks:
+            annotated_image = image_rgb.copy()
+            mp_drawing.draw_landmarks(
+                annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS
+            )
+        else:
+            print("No pose detected.")
+            return "No pose detected.", []
+
+    pose_image = np.zeros_like(image_rgb)
+    for connection in mp_pose.POSE_CONNECTIONS:
+        start_idx, end_idx = connection
+        start, end = results.pose_landmarks.landmark[start_idx], results.pose_landmarks.landmark[end_idx]
+        if start.visibility > 0.5 and end.visibility > 0.5:
+            x1, y1 = int(start.x * pose_image.shape[1]), int(start.y * pose_image.shape[0])
+            x2, y2 = int(end.x * pose_image.shape[1]), int(end.y * pose_image.shape[0])
+            cv2.line(pose_image, (x1, y1), (x2, y2), (255, 255, 255), 2)
+
+    pose_pil = Image.fromarray(cv2.resize(pose_image, (gen_width, gen_height), interpolation=cv2.INTER_LANCZOS4))
+
+    if keep_background:
+        mask_image = remove(original_image)
+        initial_mask = mask_image.split()[-1].convert('L')
+        expanded_mask = expand_mask(initial_mask, kernel_size=min(gen_width, gen_height) // 15)
+    else:
+        expanded_mask = None
+
+    generated_images = []
+
+    if keep_background:
+        with torch.no_grad():
+            with torch.amp.autocast("cuda"):
+                for prompt in prompts:
+                    images = pipe_inpaint_controlnet(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_inference_steps=20,
+                        image=Image.fromarray(image_rgb),
+                        mask_image=expanded_mask,
+                        control_image=pose_pil,
+                        width=gen_width,
+                        height=gen_height,
+                        guidance_scale=5,
+                    ).images
+                    generated_images.append(images[0])
+    else:
+        with torch.no_grad():
+            with torch.amp.autocast("cuda"):
+                for prompt in prompts:
+                    images = pipe_controlnet(
+                        prompt=prompt,
+                        negative_prompt=negative_prompt,
+                        num_inference_steps=20,
+                        image=pose_pil,
+                        guidance_scale=5,
+                        width=gen_width,
+                        height=gen_height,
+                    ).images
+                    generated_images.append(images[0])
+
+    return formatted_prompts, generated_images
+
+
+@spaces.GPU(duration=120)
+def wait_for_file(file_path, timeout=500):
+    """
+    Wait for a file to be created, with a specified timeout.
+    Args:
+        file_path (str): The path of the file to wait for.
+        timeout (int): Maximum time to wait in seconds.
+    Returns:
+        bool: True if the file is created, False if timeout occurs.
+    """
+    start_time = time.time()
+    while not os.path.exists(file_path):
+        if time.time() - start_time > timeout:
+            return False
+        time.sleep(0.5)  # Check every 0.5 seconds
+    return True
+
+
+@spaces.GPU(duration=120)
+def generate_spirit_animal_video(driving_video_path):
+    os.chdir(".")
+    try:
+        # Step 1: Extract the first frame
+        cap = cv2.VideoCapture(driving_video_path)
+        if not cap.isOpened():
+            print("Error: Unable to open video.")
+            return None
+
+        ret, frame = cap.read()
+        cap.release()
+        if not ret:
+            print("Error: Unable to read the first frame.")
+            return None
+
+        # Save the first frame
+        first_frame_path = "./first_frame.jpg"
+        cv2.imwrite(first_frame_path, frame)
+        print(f"First frame saved to: {first_frame_path}")
+
+        # Generate spirit animal image
+        _, input_image = generate_multiple_animals(first_frame_path, True, 1)
+        if input_image is None or not input_image:
+            print("Error: Spirit animal generation failed.")
+            return None
+
+        spirit_animal_path = "./animal.jpeg"
+        cv2.imwrite(spirit_animal_path, cv2.cvtColor(np.array(input_image[0]), cv2.COLOR_RGB2BGR))
+        print(f"Spirit animal image saved to: {spirit_animal_path}")
+
+        # Step 3: Run inference
+        output_path = "./animations/animal--uploaded_video_compressed.mp4"
+        script_path = os.path.abspath("./LivePortrait/inference_animals.py")
+
+        if not os.path.exists(script_path):
+            print(f"Error: Inference script not found at {script_path}.")
+            return None
+
+        command = f"python {script_path} -s {spirit_animal_path} -d {driving_video_path} --driving_multiplier 1.75 --no_flag_stitching"
+        print(f"Running command: {command}")
+        result = os.system(command)
+
+        if result != 0:
+            print(f"Error: Command failed with exit code {result}.")
+            return None
+
+        # Verify output file exists
+        if not os.path.exists(output_path):
+            print(f"Error: Expected output video not found at {output_path}.")
+            return None
+
+        print(f"Output video generated at: {output_path}")
+        return output_path
+    except Exception as e:
+        print(f"Error occurred: {e}")
+        return None
+
+
+@spaces.GPU(duration=120)
+def generate_spirit_animal(image, animal_type, background):
+    if animal_type == "Single Animal":
+        if background == "Preserve Background":
+            prompt, generated_images = spirit_animal_with_background(image)
+        else:
+            prompt, generated_images = spirit_animal_baseline(image)
+    elif animal_type == "Multiple Animals":
+        if background == "Preserve Background":
+            prompt, generated_images = generate_multiple_animals(image, keep_background=True)
+        else:
+            prompt, generated_images = generate_multiple_animals(image, keep_background=False)
+    return prompt, generated_images
+
+
+@spaces.GPU(duration=120)
+def compress_video(input_path, output_path, target_size_mb):
+    target_size_bytes = target_size_mb * 1024 * 1024
+    temp_output = "./temp_compressed.mp4"
+
+    cap = cv2.VideoCapture(input_path)
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # 使用 mp4 编码
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    writer = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        writer.write(frame)
+
+    cap.release()
+    writer.release()
+
+    current_size = os.path.getsize(temp_output)
+    if current_size > target_size_bytes:
+        bitrate = int(target_size_bytes * 8 / (current_size / target_size_bytes))  # 按比例缩减比特率
+        os.system(f"ffmpeg -i {temp_output} -b:v {bitrate} -y {output_path}")
+        os.remove(temp_output)
+    else:
+        shutil.move(temp_output, output_path)
+
+
+@spaces.GPU(duration=120)
+def process_video(video_file):
+    
+    # # 初始化 LivePortrait
+    # try:
+    #     download_liveportrait()
+    # except Exception as e:
+    #     print("Failed to initialize LivePortrait:", e)
+    #     return gr.update(value=None, visible=False)
+    
+    #     # 下载 Hugging Face 资源
+    # try:
+    #     download_huggingface_resources()
+    # except Exception as e:
+    #     print("Failed to download Hugging Face resources:", e)
+    #     return gr.update(value=None, visible=False)
+    
+    compressed_path = "./uploaded_video_compressed.mp4"
+    compress_video(video_file, compressed_path, target_size_mb=1)
+    print(f"Compressed and moved video to: {compressed_path}")
+
+    output_video_path = "./animations/animal--uploaded_video_compressed.mp4"
+    
+    generate_spirit_animal_video(compressed_path)
+
+    # Wait until the output video is generated
+    timeout = 60000  # Timeout in seconds
+    if not wait_for_file(output_video_path, timeout=timeout):
+        print("Timeout occurred while waiting for video generation.")
+        return gr.update(value=None, visible=False)  # Hide output if failed
+
+    # Return the generated video path
+    print(f"Output video is ready: {output_video_path}")
+    return gr.update(value=output_video_path, visible=True)  # Show video
+
+
+# Custom CSS styling for the interface
+css = """
+#title-container {
+    font-family: 'Arial', sans-serif;
+    color: #4a4a4a;
+    text-align: center;
+    margin-bottom: 20px;
+}
+#title-container h1 {
+    font-size: 2.5em;
+    font-weight: bold;
+    color: #ff9900;
+}
+#title-container h2 {
+    font-size: 1.2em;
+    color: #6c757d;
+}
+#intro-text {
+    font-size: 1em;
+    color: #6c757d;
+    margin: 50px;
+    text-align: center;
+    font-style: italic;
+}
+#prompt-output {
+    font-family: 'Courier New', monospace;
+    color: #5a5a5a;
+    font-size: 1.1em;
+    padding: 10px;
+    background-color: #f9f9f9;
+    border: 1px solid #ddd;
+    border-radius: 5px;
+    margin-top: 10px;
+}
+"""
+
+# Title and description
+title_html = """
+<div id="title-container">
+    <h1>Spirit Animal Generator</h1>
+    <h2>Create your unique spirit animal with AI-assisted image generation.</h2>
+</div>
+"""
+
+description_text = """
+### Project Overview
+Welcome to the Spirit Animal Generator! This tool leverages advanced AI technologies to create unique visualizations of spirit animals from both videos and images.
+#### Key Features:
+1. **Video Transformation**: Upload a driving video to generate a creative spirit animal animation.
+2. **Image Creation**: Upload an image and customize the spirit animal type and background options.
+3. **AI-Powered Prompting**: OpenAI's GPT generates descriptive prompts for each input.
+4. **High-Quality Outputs**: Generated using Stable Diffusion and ControlNet for stunning visuals.
+---
+### How It Works:
+1. **Upload Your Media**:
+   - Videos: Ensure the file is in MP4 format.
+   - Images: Use clear, high-resolution photos for better results.
+2. **Customize Options**:
+   - For images, select the type of animal and background settings.
+3. **View Your Results**:
+   - Videos will be transformed into animations.
+   - Images will produce customized visual art along with a generated prompt.
+Discover your spirit animal and let your imagination run wild!
+---
+"""
+
+with gr.Blocks() as demo:
+    gr.HTML(title_html)
+    gr.Markdown(description_text)
+
+    with gr.Tabs():
+        with gr.Tab("Generate Spirit Animal Image"):
+            gr.Markdown("Upload an image to generate a spirit animal.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    image_input = gr.Image(type="filepath", label="Upload an image")
+                    animal_type = gr.Radio(choices=["Single Animal", "Multiple Animals"], label="Animal Type", value="Single Animal")
+                    background_option = gr.Radio(choices=["Preserve Background", "Don't Preserve Background"], label="Background Option", value="Preserve Background")
+                    generate_image_button = gr.Button("Generate Image")
+                with gr.Column(scale=1):
+                    generated_prompt = gr.Textbox(label="Generated Prompt")
+                    generated_gallery = gr.Gallery(label="Generated Images")
+
+            generate_image_button.click(
+                fn=generate_spirit_animal,
+                inputs=[image_input, animal_type, background_option],
+                outputs=[generated_prompt, generated_gallery],
+            )
+
+        with gr.Tab("Generate Spirit Animal Video"):
+            gr.Markdown("Upload a driving video to generate a spirit animal video.")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    video_input = gr.Video(label="Upload a driving video (MP4 format)")
+                    generate_video_button = gr.Button("Generate Video")
+                with gr.Column(scale=1):
+                    video_output = gr.Video(label="Generated Spirit Animal Video")
+
+            generate_video_button.click(
+                fn=process_video,
+                inputs=video_input,
+                outputs=video_output,
+            )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,20 @@
+torch == 2.1.2
+torchvision==0.16.2
+torchaudio==2.1.2
+moviepy==1.0.3
+imageio[ffmpeg]
+pillow==10.4.0
+tyro==0.8.5
+onnxruntime-gpu==1.18.1
+onnx==1.16.1
+gradio==4.37.1
+colorama
+ffmpeg-python==0.2.0
+mediapipe
+rembg
+huggingface_hub[cli]
+opencv-python
+matplotlib
+diffusers
+transformers
+accelerate