app.py

import gradio as gr
import numpy as np
import random
import spaces
import torch
import os
from diffusers import DiffusionPipeline
from transformers import pipeline
from huggingface_hub import login

# Login to Hugging Face Hub with token
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    login(token=hf_token)
else:
    print("Warning: HF_TOKEN environment variable not found. Authentication may fail.")

# Translation pipeline and hardware settings
device = "cuda" if torch.cuda.is_available() else "cpu"
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device=device)
dtype = torch.bfloat16

# Load the model with token authentication
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", 
    torch_dtype=dtype,
    use_auth_token=hf_token
).to(device)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

def enhance_logo_prompt(prompt):
    """Enhance prompt specifically for logo generation"""
    logo_keywords = [
        "minimalist logo design",
        "vector graphics style", 
        "simple geometric shapes",
        "flat design",
        "logo icon",
        "white background",
        "centered logo composition",
        "professional brand identity",
        "clean lines",
        "no photography",
        "no realistic details",
        "graphic design",
        "corporate logo style"
    ]
    
    # Add negative prompts to avoid photorealistic results
    negative_context = "not a photograph, not realistic, not 3d render, not complex scene"
    
    enhanced_prompt = f"{prompt}, {', '.join(logo_keywords)}, {negative_context}"
    
    return enhanced_prompt

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    
    # Always enhance prompt for logo generation
    prompt = enhance_logo_prompt(prompt)
    
    # Korean input detection and translation
    if any('\uAC00' <= char <= '\uD7A3' for char in prompt):
        print("Translating Korean prompt...")
        translated_prompt = translator(prompt, max_length=512)[0]['translation_text']
        print("Translated prompt:", translated_prompt)
        prompt = translated_prompt
        
    image = pipe(
            prompt = prompt,
            width = width,
            height = height,
            num_inference_steps = num_inference_steps,
            generator = generator,
            guidance_scale=0.0
    ).images[0]
    
    return image, seed

examples = [
    ["minimal tech company logo, circuit board pattern, blue and white, TEXT 'AI FOREVER'"],
    ["coffee shop logo, coffee bean icon, brown and cream colors, circular design, TEXT 'T. 123-1234'"],
    ["fitness gym logo, dumbbell symbol, red and black, dynamic angles, TEXT 'abc@email.com'"],
    ["eco friendly company logo, leaf design, green gradient, modern style, TEXT 'NAME CARD'"],
    ["fashion brand logo, elegant typography, gold and black, luxury style, TEXT 'abc.com'"],
    ["startup logo, rocket icon, purple and orange, innovative design, TEXT 'EVER AI'"]
]

css = """
/* Clean, minimal styling */
.container {
    max-width: 1200px;
    margin: auto;
    padding: 20px;
}

/* Simple title */
.title {
    text-align: center;
    font-size: 2.5em;
    font-weight: 600;
    margin-bottom: 20px;
    color: #333;
}

/* White background for input */
#prompt textarea {
    background-color: white !important;
    color: #333 !important;
    border: 2px solid #e0e0e0;
    border-radius: 8px;
    padding: 12px;
    font-size: 16px;
    transition: border-color 0.3s ease;
}

#prompt textarea:focus {
    border-color: #4CAF50;
    outline: none;
}

/* Clean button styling */
.gr-button-primary {
    background-color: #4CAF50 !important;
    border: none;
    padding: 12px 30px;
    font-size: 16px;
    font-weight: 500;
    border-radius: 6px;
    transition: background-color 0.3s ease;
}

.gr-button-primary:hover {
    background-color: #45a049 !important;
}

/* Result image styling */
#result {
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 10px;
    background-color: #f9f9f9;
}

/* Info box styling */
.info-box {
    background-color: #f0f0f0;
    border-left: 4px solid #4CAF50;
    padding: 15px;
    margin: 15px 0;
    border-radius: 4px;
}

/* Accordion styling */
.gr-accordion {
    background-color: #f9f9f9;
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    margin-top: 20px;
}

/* Examples section */
.gr-examples {
    margin-top: 30px;
    background-color: #f9f9f9;
    padding: 20px;
    border-radius: 8px;
    border: 1px solid #e0e0e0;
}

/* Hide footer */
footer { 
    visibility: hidden; 
}

/* Responsive design */
@media (max-width: 768px) {
    .container {
        padding: 15px;
    }
    .title {
        font-size: 2em;
    }
}
"""

with gr.Blocks(theme="soft", css=css) as demo:
    

    gr.HTML(
        """
        <div class='container'>
            <h1 class='title'>Logo Generator AI</h1>
    
            <p style='text-align:center; color:#666; font-size:1.1em; margin-bottom:20px;'>
                Create simple, professional logos with AI
            </p>
            <div style='display:flex; justify-content:center; gap:12px; flex-wrap:wrap; margin-bottom:30px;'>
                <a href="https://discord.gg/openfreeai" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord badge">
                </a>
                
                <a href="https://huggingface.co/OpenFreeAI" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=Community&message=OpenFree_AI&color=%23800080&labelColor=%23000080&logo=HUGGINGFACE&logoColor=%23ffa500&style=for-the-badge" alt="badge">
                </a>

                <a href="https://huggingface.co/spaces/openfree/Best-AI" target="_blank">
                    <img src="https://img.shields.io/static/v1?label=OpenFree&message=BEST%20AI%20Services&color=%230000ff&labelColor=%23000080&logo=huggingface&logoColor=%23ffa500&style=for-the-badge" alt="OpenFree badge">
                </a>
            </div>
        </div>
        """
    )

    
    with gr.Column(elem_id="container"):
        # Main input section
        with gr.Group():
            gr.HTML("""
                <div class='info-box'>
                    <strong>Tip:</strong> Describe your logo using simple terms like: company type, icon/symbol, colors, style.
                    <br>Example: "tech startup logo, lightning bolt icon, blue and silver, minimalist"
                </div>
            """)
            
            prompt = gr.Textbox(
                label="Logo Description",
                placeholder="Describe your logo design...\nExample: tech company logo, abstract shape, blue and white, minimal design",
                lines=3,
                elem_id="prompt"
            )
            
            with gr.Row():
                run_button = gr.Button("Generate Logo", variant="primary", scale=2)
                clear_button = gr.Button("Clear", scale=1)
        
        # Result section
        with gr.Row():
            with gr.Column(scale=3):
                result = gr.Image(
                    label="Generated Logo", 
                    show_label=True,
                    elem_id="result",
                    interactive=False
                )
            
            with gr.Column(scale=1):
                seed_output = gr.Number(label="Seed Used", interactive=False)
        
        # Advanced settings
        with gr.Accordion("Advanced Settings", open=False):
            with gr.Row():
                seed = gr.Slider(
                    label="Seed", 
                    minimum=0, 
                    maximum=MAX_SEED, 
                    step=1, 
                    value=0,
                    info="Set to 0 for random seed"
                )
                randomize_seed = gr.Checkbox(
                    label="Random Seed", 
                    value=True,
                    info="Generate unique results each time"
                )
            
            with gr.Row():
                width = gr.Slider(
                    label="Width", 
                    minimum=256, 
                    maximum=MAX_IMAGE_SIZE, 
                    step=32, 
                    value=1024,
                    info="Logo width in pixels"
                )
                height = gr.Slider(
                    label="Height", 
                    minimum=256, 
                    maximum=MAX_IMAGE_SIZE, 
                    step=32, 
                    value=1024,
                    info="Logo height in pixels"
                )
            
            num_inference_steps = gr.Slider(
                label="Quality Steps", 
                minimum=1, 
                maximum=50, 
                step=1, 
                value=4,
                info="Higher = better quality but slower"
            )
        
        # Examples section
        gr.HTML("<div class='info-box'><strong>Example Prompts:</strong></div>")
        gr.Examples(
            examples=examples,
            fn=infer,
            inputs=[prompt],
            outputs=[result, seed_output],
            cache_examples="lazy"
        )
        
        # Event handlers
        gr.on(
            triggers=[run_button.click, prompt.submit],
            fn=infer,
            inputs=[prompt, seed, randomize_seed, width, height, num_inference_steps],
            outputs=[result, seed_output]
        )
        
        # Clear button functionality
        clear_button.click(
            fn=lambda: (None, None),
            outputs=[prompt, result]
        )
        
        # Tips section
        gr.HTML("""
            <div class='info-box' style='margin-top: 30px;'>
                <strong>Logo Design Tips:</strong>
                <ul style='margin: 10px 0; padding-left: 20px; color: #666;'>
                    <li>Use simple, clear descriptions</li>
                    <li>Specify icon type (abstract, letter, symbol)</li>
                    <li>Mention preferred colors</li>
                    <li>Include style (minimal, modern, classic)</li>
                    <li>Keep it simple - logos should be clean and scalable</li>
                </ul>
            </div>
        """)

demo.launch()