import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from transformers import pipeline
from gtts import gTTS
import os

def generate_character(description, emotion, physics, dialogue):
    # Load Stable Diffusion Model (Optimized for Speed)
    pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
    pipe.to("cuda" if torch.cuda.is_available() else "cpu")
    
    # Modify description based on emotion
    emotion_map = {"Happy": "smiling", "Angry": "furious", "Sad": "teary-eyed"}
    description = f"{description}, {emotion_map.get(emotion, '')}"
    
    # Generate Image
    image = pipe(description).images[0]
    image.save("output.png")
    
    # Generate Voiceover
    try:
        tts_model = pipeline("text-to-speech", model="facebook/mms-tts-eng")
        speech_output = tts_model(dialogue)
        speech_path = "output.wav"
        with open(speech_path, "wb") as f:
            f.write(speech_output["audio"])
    except Exception:
        tts = gTTS(dialogue)
        speech_path = "output.mp3"
        tts.save(speech_path)
    
    return image, speech_path

gui = gr.Interface(
    fn=generate_character,
    inputs=[
        gr.Textbox(label="Character Description"),
        gr.Radio(["Happy", "Angry", "Sad"], label="Emotion"),
        gr.Radio(["None", "Wind", "Falling"], label="Physics Effect"),
        gr.Textbox(label="Dialogue")
    ],
    outputs=[gr.Image(label="Generated Character"), gr.Audio(label="Voiceover")],
    allow_flagging="never"
)

gui.launch(debug=True)