replacebg / app.py
Munaf1987's picture
Update app.py
0f7ba4c verified
raw
history blame
2.36 kB
import gradio as gr
import json
from transformers import pipeline
from PIL import Image
import numpy as np
import os
import spaces
# Text Expansion Model (ZeroGPU compatible)
text_generator = pipeline("text-generation", model="tiiuae/falcon-7b-instruct")
# Text-to-Speech Model (Bark small)
tts_pipeline = pipeline("text-to-speech", model="suno/bark-small")
# Image Generation Model (LoRA-based)
image_generator = pipeline("text-to-image", model="stabilityai/sdxl-turbo")
# Main Processing Function
@spaces.GPU
def generate_cartoon(script_text):
# Step 1: Expand Script
prompt = f"""
You are a cartoon script writer. Convert the following story into a detailed cartoon scene plan.
For each scene, provide:
1. Scene description (setting + action)
2. Dialogue
3. Characters involved
4. Background description
Story:
"""
{script_text}
"""
Return result in JSON format.
"""
response = text_generator(prompt, max_new_tokens=1024)[0]['generated_text']
try:
scene_data = json.loads(response)
except:
return "Script expansion failed. Please refine input."
os.makedirs("generated_images", exist_ok=True)
os.makedirs("generated_audio", exist_ok=True)
scene_results = []
for idx, scene in enumerate(scene_data.get("scenes", [])):
# Generate Background Image
background_prompt = scene.get("background_description", "cartoon background")
background_image = image_generator(background_prompt).images[0]
bg_path = f"generated_images/scene_{idx+1}.png"
background_image.save(bg_path)
# Generate TTS Audio
dialogue = scene.get("dialogue", "")
audio_output = tts_pipeline(dialogue)
audio_path = f"generated_audio/scene_{idx+1}.wav"
audio_output['audio'].export(audio_path, format="wav")
scene_results.append((bg_path, audio_path))
return scene_results
# Gradio UI
demo = gr.Interface(
fn=generate_cartoon,
inputs=gr.Textbox(label="Enter Cartoon Script", lines=10),
outputs=gr.Gallery(label="Generated Scenes with Audio"),
title="Cartoon Scene Generator",
description="Enter a cartoon story script to generate scenes with background images and audio. Video merging should be done separately in Google Colab."
)
if __name__ == "__main__":
demo.launch()