# app.py - Improved Custom Content Generator
import gradio as gr
from transformers import pipeline, set_seed, AutoTokenizer, AutoModelForCausalLM
from diffusers import StableDiffusionPipeline
import torch

# ---------- CONFIG ----------
TEXT_MODEL = "openai-community/gpt2"  # Text model
CODE_MODEL = "Salesforce/codegen-350M-multi"  # Better for code
IMAGE_MODEL = "stabilityai/stable-diffusion-2-1"  # Higher quality images
SEED = 42
# ----------------------------

# Load text generator
text_tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL)
text_model = AutoModelForCausalLM.from_pretrained(TEXT_MODEL)
text_pipe = pipeline("text-generation", model=text_model, tokenizer=text_tokenizer,
                     device_map="auto" if torch.cuda.is_available() else None)

# Load code generator
code_tokenizer = AutoTokenizer.from_pretrained(CODE_MODEL)
code_model = AutoModelForCausalLM.from_pretrained(CODE_MODEL)
code_pipe = pipeline("text-generation", model=code_model, tokenizer=code_tokenizer,
                     device_map="auto" if torch.cuda.is_available() else None)

# Load Stable Diffusion
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    sd_pipe = StableDiffusionPipeline.from_pretrained(
        IMAGE_MODEL,
        torch_dtype=torch.float16 if device == "cuda" else torch.float32
    ).to(device)
    has_sd = True
except Exception as e:
    sd_pipe = None
    has_sd = False
    print("Stable Diffusion not available:", e)

set_seed(SEED)

# ---------- Generation helpers ----------
def generate_text(prompt, max_len=150, temperature=0.7, top_k=50, top_p=0.95, num_return=1):
    out = text_pipe(prompt, max_length=max_len, do_sample=True,
                    temperature=float(temperature), top_k=int(top_k),
                    top_p=float(top_p), num_return_sequences=int(num_return))
    return "\n\n===\n\n".join(o['generated_text'] for o in out)

def generate_image(prompt, steps=35, guidance_scale=8.0, height=768, width=768):
    if not has_sd:
        return "Stable Diffusion model not loaded."
    return sd_pipe(prompt, num_inference_steps=int(steps),
                   guidance_scale=float(guidance_scale),
                   height=int(height), width=int(width)).images[0]

def generate_code(prompt, language="Python", max_len=256):
    full_prompt = f"# {language} code\n# {prompt}\n"
    out = code_pipe(full_prompt, max_length=max_len, do_sample=True,
                    temperature=0.2, top_k=50, top_p=0.95, num_return_sequences=1)
    return out[0]['generated_text']

# ---------- UI ----------
with gr.Blocks(title="Custom Content Generator") as demo:
    gr.Markdown("## Custom Content Generator — Text, Image, and Code.")

    with gr.Tabs():
        with gr.TabItem("Text"):
            inp = gr.Textbox(label="Prompt", lines=3)
            max_len = gr.Slider(50, 512, value=150, step=10, label="Max length")
            temp = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
            top_k = gr.Slider(0, 200, value=50, step=1, label="Top-k")
            top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="Top-p")
            out_text = gr.Textbox(label="Generated Text", lines=12)
            gr.Button("Generate Text").click(
                generate_text, [inp, max_len, temp, top_k, top_p, gr.Number(value=1, visible=False)], out_text
            )

        with gr.TabItem("Image"):
            img_prompt = gr.Textbox(label="Image Prompt", lines=2)
            steps = gr.Slider(10, 50, value=35, step=1, label="Steps")
            guidance = gr.Slider(1.0, 20.0, value=8.0, step=0.1, label="Guidance scale")
            size = gr.Dropdown(choices=["512x512", "768x768"], value="768x768", label="Image size")
            img_out = gr.Image(label="Generated Image")
            def img_gen(prompt, steps, guidance, size):
                w, h = map(int, size.split("x"))
                return generate_image(prompt, steps=steps, guidance_scale=guidance, height=h, width=w)
            gr.Button("Generate Image").click(img_gen, [img_prompt, steps, guidance, size], img_out)

        with gr.TabItem("Code"):
            code_prompt = gr.Textbox(label="Describe Code Task", lines=3)
            lang = gr.Dropdown(choices=["Python", "JavaScript", "TypeScript", "Java", "C#", "Go"], value="Python")
            code_out = gr.Textbox(label="Generated Code", lines=18)
            gr.Button("Generate Code").click(generate_code, [code_prompt, lang], code_out)

if __name__ == "__main__":
    demo.launch()