app.py

import gradio as gr
import torch
import io
import smtplib
from email.message import EmailMessage
from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline

from diffusers.utils import load_image
from PIL import Image


# Load the base (text-to-image) model
pipe_t2i = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
).to("cuda")

# Use from_pipe to initialize img2img from the same base
pipe_i2i = StableDiffusionXLImg2ImgPipeline.from_pipe(pipe_t2i).to("cuda")

def generate_from_image(image: Image.Image, prompt: str,
                        strength: float = 0.75,
                        guidance_scale: float = 10.5,
                        steps: int = 40):
    """
    Generates a new SDXL image conditioned on both prompt and input image.
    """
    print(f"Prompt: {prompt}, Strength: {strength}, Guidance: {guidance_scale}")

    image = image.convert("RGB").resize((1024, 1024))
    with torch.inference_mode():
        output = pipe_i2i(
            prompt=prompt,
            image=image,
            strength=strength,
            guidance_scale=guidance_scale,
            num_inference_steps=steps
        ).images[0]
    return output

style_prompts = {
    "Single Photo With Dhaka Topi - Ghibli Style": "Create a Ghibli-style portrait using the uploaded photo as the person’s face to get a magical Ghibli-style portrait, wearing a traditional Nepali Dhaka topi. The dreamy background features Mount Everest, blooming laliguras flowers, Kathmandu’s Ghanta Ghar tower, elegant Newari window patterns, and colorful Nepali prayer flags waving in the wind. The art should blend Studio Ghibli’s enchanting style with authentic Nepali culture.",
    "Two Person Picture With Nepali Style Ghibli Art": "Create a Ghibli-style portrait using the uploaded photo as the person’s face. Background features Mount Everest, laliguras flowers, Kathmandu’s Ghanta Ghar tower, Newari window patterns, colorful Nepali prayer flags, the Nepali flag, Pashupatinath temple, and a flying Himalayan monal (Danfe). Make it vibrant, enchanting, and full of Nepali cultural details.",
    "Spirited Away Fantasy": "Studio Ghibli style portrait with fantasy elements, mysterious background, dreamy lighting, spirited away aesthetic",
    "Kiki's Cozy Town": "A warm, cozy background with bakery and rooftops, Studio Ghibli anime style, charming and colorful",
    "Princess Mononoke Wild": "Forest warrior aesthetic, powerful gaze, nature spirit background, strong composition, Ghibli anime style",
    "Ghibli Vintage Portrait": "Classic Ghibli portrait, soft facial features, minimal background, vintage anime color palette, watercolor ink"
}
sample_keys = list(style_prompts.keys())[:5]



def process(image, style, email):
    prompt = style_prompts.get(style, "")
    result = generate_from_image(image, prompt)
    email_status = ""
    if email:
        email_status = send_email(email, result)
    return result, "✅ Generated", email_status

def send_email(email, image):
    msg = EmailMessage()
    msg["Subject"] = "Your Ghibli-style Portrait"
    msg["From"] = "[email protected]"
    msg["To"] = email
    buf = io.BytesIO(); image.save(buf, format="PNG"); buf.seek(0)
    msg.add_attachment(buf.read(), maintype="image", subtype="png", filename="portrait.png")
    with smtplib.SMTP_SSL("smtp.example.com", 465) as s:
        s.login("[email protected]", "password")
        s.send_message(msg)
    return "✅ Sent to email!"

with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="emerald", secondary_hue="amber", neutral_hue="stone"),
    css="""
        body { font-family: 'Merriweather', serif; background: #f9f5f0; }
        h1 { color: #2a6f48; font-family: 'Karma', serif; }
        .sample-btn { margin: 4px; background: #a4c639; color: white; border-radius:8px; }
        .sample-btn:hover { background: #7fa32c; }
    """
) as demo:
    gr.Markdown("## 🌸 Nepali Roots Ghibli Art Booth")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Choose a Ghibli prompt style:")
            sample_buttons = []
            for key in sample_keys:
                btn = gr.Button(key, elem_classes="sample-btn")
                sample_buttons.append(btn)
            style_choice = gr.Dropdown(list(style_prompts.keys()), value=sample_keys[0], label="🎨 Style")
            upload = gr.Image(sources=["upload"], type="pil", label="📸 Take or Upload Image")
            email_input = gr.Textbox(label="📧 Email (optional)")
            generate_button = gr.Button("✨ Generate & Project")

        with gr.Column(scale=1):
            image_output = gr.Image(label="🌠 Result", show_fullscreen_button=True)
            status = gr.Textbox(label="Status")
            email_status = gr.Textbox(label="Email Status")
            projector = gr.HTML("", elem_id="projector")

    # sample button callback
    for btn, key in zip(sample_buttons, sample_keys):
        btn.click(lambda k=key: gr.update(value=k), inputs=None, outputs=style_choice)

    generate_button.click(process, [image_input, style_choice, email_input], [image_output, status, email_status])

    # projector link update
    def projector_link(img):
        if img is None: return ""
        return '<a href="#" onclick="window.open(window.images[0].src,\'_blank\')">🔍 Open full-screen on projector</a>'
    image_output.change(projector_link, inputs=[image_output], outputs=[projector])

demo.launch()