app.py

import os
import torch
import random
# Removed google.colab import for deployment
from diffusers import StableDiffusionPipeline
import gradio as gr

# ---------------------------------------------------------
# 1. SETUP & HARDWARE DETECTION (Auto-Switching)
# ---------------------------------------------------------
# This allows it to run on Free CPU (Slow) or Paid GPU (Fast) automatically.
if torch.cuda.is_available():
    device = "cuda"
    dtype = torch.float16
    print("🚀 Running on GPU (Fast Mode)")
else:
    device = "cpu"
    dtype = torch.float32
    print("🐢 Running on CPU (Slow Compatibility Mode)")

print(f"⏳ Loading DreamShaper V44 on {device}...")

pipe = StableDiffusionPipeline.from_pretrained(
    "Lykon/dreamshaper-8", 
    torch_dtype=dtype,
    use_safetensors=True
)
pipe.to(device)

if device == "cuda":
    pipe.enable_attention_slicing()
    
# Disable Safety Checker to save memory/speed
pipe.safety_checker = None
pipe.requires_safety_checker = False

# ---------------------------------------------------------
# 2. LOGICAL CONTENT ENGINE
# ---------------------------------------------------------
THEMES = {
    "Cyberpunk": {
        "subjects": ["A young cyberpunk hacker", "A futuristic android", "A rogue cyborg mercenary", "A holographic AI avatar"],
        "outfits": ["wearing a glowing neon jacket", "wearing a transparent plastic raincoat", "wearing high-tech tactical armor"],
        "locations": ["in a rain-soaked neon city", "inside a server room with green lasers", "in a grimy futuristic alleyway", "on a rooftop overlooking flying cars"]
    },
    "Fantasy": {
        "subjects": ["A high elf princess", "A rugged dwarf warrior", "A mysterious forest druid", "A dark necromancer"],
        "outfits": ["wearing intricate silver plate armor", "wearing a hooded velvet cloak", "wearing robes embroidered with gold"],
        "locations": ["deep in a bioluminescent forest", "inside a grand castle throne room", "standing on a cliff overlooking a dragon valley", "in a magical crystal cave"]
    },
    "Modern/Urban": {
        "subjects": ["A stylish university student", "A lo-fi hip hop girl", "A weary detective", "A graffiti artist"],
        "outfits": ["wearing an oversized streetwear hoodie", "wearing a vintage denim jacket", "wearing a cozy knitted sweater"],
        "locations": ["sitting in a cozy cafe with warm lighting", "standing on a subway platform", "walking through a rainy Tokyo street", "in a messy bedroom full of posters"]
    },
    "Historical": {
        "subjects": ["A traditional samurai", "A victorian gothic lady", "A viking shieldmaiden", "A 1920s jazz singer"],
        "outfits": ["wearing a floral kimono with a katana", "wearing a black lace victorian dress", "wearing fur-lined leather armor"],
        "locations": ["standing in a cherry blossom garden", "inside a foggy london street", "on the deck of a wooden ship", "in a workshop full of brass clocks"]
    }
}

ACTIONS = [
    "looking at the camera", "holding a weapon", "reading a book", 
    "drinking coffee", "adjusting hair", "standing confidently"
]

ATMOSPHERES = [
    "at sunset", "during a storm", "under a starry sky", "in the morning mist", 
    "with dramatic shadows", "bathed in moonlight", "during golden hour"
]

def get_logical_prompt(subject_mode):
    theme_name = random.choice(list(THEMES.keys()))
    data = THEMES[theme_name]
    l = random.choice(data["locations"])
    
    if subject_mode == "Scenery Only":
        clean_loc = l.replace("standing ", "").replace("sitting ", "").replace("walking ", "")
        atmos = random.choice(ATMOSPHERES)
        return f"{clean_loc}, {atmos}"
    else:
        s = random.choice(data["subjects"])
        o = random.choice(data["outfits"])
        a = random.choice(ACTIONS)
        return f"{s} {o}, {a}, {l}"

def on_subject_change(mode):
    new_text = ""
    if mode == "Scenery Only":
        new_shape = "Landscape"
    else:
        new_shape = "Portrait"
    return new_text, new_shape

def live_style_manager(user_text, selected_style, subject_mode):
    if not user_text: return ""
    cleaned = user_text

    if selected_style == "Anime (2D)":
        cleaned = cleaned.replace(", detailed iris", "").replace("detailed iris", "")
        cleaned = cleaned.replace(", realistic eyes", "").replace("realistic eyes", "")
        cleaned = cleaned.replace(", skin texture", "").replace("skin texture", "")
    elif selected_style == "Realistic":
        if subject_mode == "Character" and "detailed iris" not in cleaned:
            cleaned += ", detailed iris"
    return cleaned

def smart_enhance(user_text, selected_style, subject_mode):
    if not user_text: return "masterpiece"
    base = user_text + ", high quality, sharp lines, masterpiece, 8k resolution"
    if subject_mode == "Character":
        if "perfect eyes" not in base:
            base += ", perfect eyes, symmetrical eyes"
        if selected_style != "Anime (2D)" and "detailed iris" not in base:
            base += ", detailed iris"
    elif subject_mode == "Scenery Only":
        base += ", detailed background, atmospheric, cinematic composition, wide angle"
    return base

def apply_style(prompt, preset):
    if preset == "Anime (2D)":
        clean_prompt = prompt.replace("detailed iris", "").replace("realistic eyes", "")
        suffix = ", anime style, studio ghibli, cel shading, flat color, vector art, 2d, vibrant colors, simple shading"
        neg = "3d, realistic, photorealistic, oil painting, render, octane, unreal engine, clay, blender, shading, shadows, detailed iris"
        return clean_prompt + suffix, neg
    else:
        suffix = ", photograph, 8k, raw photo, realistic texture, ray tracing, dslr, cinematic lighting, film grain"
        neg = "anime, cartoon, drawing, painting, illustration, 2d, sketch, flat color"
        return prompt + suffix, neg

MASTER_NEGATIVE = "blurry, low quality, bad anatomy, ugly, distorted, watermark, text, signature"
FACE_NEGATIVE = ", strabismus, crossed eyes, cross-eyed, uneven eyes, lazy eye, asymmetric eyes, bad eyes, open mouth, missing teeth, disfigured, extra fingers, missing limbs"
SCENERY_NEGATIVE = ", woman, man, person, people, human, face, portrait, character, girl, boy, body, skin, eyes, limbs"

def detect_face_request(prompt):
    triggers = ["face", "portrait", "eyes", "mouth", "closeup", "close up", "woman", "man", "girl", "boy", "character", "selfie"]
    return any(trigger in prompt.lower() for trigger in triggers)

def generate(prompt, aspect_ratio, style_preset, subject_mode, mode, cust_steps, cust_guidance, cust_seed):
    
    final_prompt, style_negative = apply_style(prompt, style_preset)
    
    if aspect_ratio == "Portrait": w, h = 512, 768
    elif aspect_ratio == "Landscape": w, h = 768, 512
    else: w, h = 512, 768

    if mode == "Custom (Expert)":
        steps = int(cust_steps)
        guidance = float(cust_guidance)
        seed = random.randint(0, 2147483647) if cust_seed == -1 else int(cust_seed)
        final_negative = MASTER_NEGATIVE + style_negative
        if subject_mode == "Scenery Only":
             final_negative += SCENERY_NEGATIVE
        info_prefix = "🔧 Custom"
    else: 
        base_neg = MASTER_NEGATIVE + style_negative
        if subject_mode == "Scenery Only":
            steps = 28
            guidance = 7.5
            final_negative = base_neg + SCENERY_NEGATIVE
            info_prefix = "🌍 Scenery Mode"
        else: 
            is_portrait = detect_face_request(prompt)
            if is_portrait:
                steps = 35 
                guidance = 6.5 
                final_negative = base_neg + FACE_NEGATIVE
                info_prefix = "🤖 Auto-Face"
            else:
                steps = 28
                guidance = 7.5
                final_negative = base_neg
                info_prefix = "✈️ Auto-Creative"
        
        if style_preset == "Anime (2D)":
            guidance = 7.0
            info_prefix += " (Anime)"

        seed = random.randint(0, 2147483647)

    # Use 'device' variable determined at startup
    generator = torch.Generator(device).manual_seed(seed)
    image = pipe(
        prompt=final_prompt, 
        negative_prompt=final_negative, 
        num_inference_steps=steps, 
        guidance_scale=guidance, 
        width=w, 
        height=h, 
        generator=generator
    ).images[0]
    
    return image, f"{info_prefix} | {w}x{h} | Steps: {steps} | G: {guidance} | Seed: {seed}"

with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🚀 DreamShaper V44 (Hosted)")
    
    with gr.Row():
        with gr.Column():
            txt = gr.Textbox(label="Prompt", placeholder="Type here or click 'Logical Idea'...", lines=4)
            subject_mode = gr.Radio(["Character", "Scenery Only"], value="Character", label="1. Subject Type")
            style_preset = gr.Radio(["Realistic", "Anime (2D)"], value="Realistic", label="2. Style Preset")

            with gr.Row():
                random_btn = gr.Button("🧠 Logical Idea", variant="primary")
                refine_btn = gr.Button("✨ Enhance Text", variant="secondary") 
            
            shape = gr.Radio(["Portrait", "Landscape"], value="Portrait", label="Shape")
            mode = gr.Radio(["Random (Auto-Pilot)", "Custom (Expert)"], value="Random (Auto-Pilot)", label="Mode")
            
            with gr.Group(visible=False) as custom_box:
                cust_steps = gr.Slider(10, 60, value=30, label="Steps")
                cust_guidance = gr.Slider(1.0, 15.0, value=7.0, label="Guidance")
                cust_seed = gr.Number(value=-1, label="Seed")

            btn = gr.Button("Generate 🎨", variant="primary")
        
        with gr.Column():
            out_img = gr.Image(label="Result")
            out_info = gr.Textbox(label="Settings Used", interactive=False)

    random_btn.click(get_logical_prompt, [subject_mode], [txt])
    refine_btn.click(smart_enhance, [txt, style_preset, subject_mode], [txt])
    style_preset.change(live_style_manager, [txt, style_preset, subject_mode], [txt])
    subject_mode.change(on_subject_change, inputs=[subject_mode], outputs=[txt, shape])
    
    def toggle(val): return gr.update(visible=(val == "Custom (Expert)"))
    mode.change(toggle, mode, custom_box)

    btn.click(generate, [txt, shape, style_preset, subject_mode, mode, cust_steps, cust_guidance, cust_seed], [out_img, out_info])

app.launch()