app.py

import gradio as gr
import torch
from diffusers import DiffusionPipeline
import random

# --- Model Configuration ---
model_id = "NewBie-AI/NewBie-image-Exp0.1"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading model on: {device}...")

# Load pipeline with proper error handling
try:
    pipe = DiffusionPipeline.from_pretrained(
        model_id, 
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        use_safetensors=True
    )
    pipe.to(device)
    print("Model loaded successfully.")
except Exception as e:
    print(f"Error loading model: {e}")

# --- Generation Function ---
def generate_image(prompt, negative_prompt, steps, guidance_scale, width, height, seed):
    if seed == -1:
        seed = random.randint(0, 2147483647)
    
    # Configure generator for reproducibility - FIXED SYNTAX ERROR
    generator = torch.Generator(device).manual_seed(int(seed))
    
    print(f"Generating with seed: {seed}")
    
    try:
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=int(steps),
            guidance_scale=float(guidance_scale),
            width=int(width),
            height=int(height),
            generator=generator
        ).images[0]
        return image, seed
    except Exception as e:
        return None, f"Error: {str(e)}"

# --- Gradio Interface ---
css = """
#col-container {max-width: 800px; margin-left: auto; margin-right: auto;}
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"#  🎨 Image Generator: {model_id}")
        gr.Markdown("Write a prompt to generate an image using the NewBie AI model.")
        
        with gr.Group():
            prompt = gr.Textbox(
                label="Prompt (Positive description)", 
                placeholder="An astronaut riding a horse on Mars, 4k, realistic...", 
                lines=2
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt (What you DON'T want)", 
                placeholder="blurry, deformed, bad quality, text...", 
                value="bad quality, worst quality, low resolution, blurry, distorted"
            )
            
        with gr.Row():
            with gr.Column():
                width = gr.Slider(
                    label="Width", 
                    minimum=256, 
                    maximum=1024, 
                    step=64, 
                    value=512
                )
            with gr.Column():
                height = gr.Slider(
                    label="Height", 
                    minimum=256, 
                    maximum=1024, 
                    step=64, 
                    value=512
                )
            
        with gr.Row():
            with gr.Column():
                steps = gr.Slider(
                    label="Inference Steps", 
                    minimum=10, 
                    maximum=100, 
                    step=1, 
                    value=25
                )
            with gr.Column():
                guidance_scale = gr.Slider(
                    label="Guidance Scale (Prompt fidelity)", 
                    minimum=1.0, 
                    maximum=20.0, 
                    step=0.5, 
                    value=7.5
                )
            
        with gr.Row():
            seed = gr.Number(
                label="Seed (Use -1 for random)", 
                value=-1, 
                precision=0
            )
            
        with gr.Accordion("Advanced Configuration", open=False):
            gr.Markdown("Adjust these parameters for finer control over image generation.")
                
        run_button = gr.Button("Generate Image", variant="primary")
        result_image = gr.Image(label="Result", interactive=False)
        seed_output = gr.Number(label="Seed used", interactive=False)
        
        clear_button = gr.Button("Clear", variant="secondary")
        
        # Examples
        examples = gr.Examples(
            examples=[
                ["A beautiful sunset over mountains, digital art", "blurry, distorted", 25, 7.5, 512, 512, -1],
                ["A futuristic city with flying cars, cyberpunk style", "low quality, pixelated", 30, 5.0, 768, 768, -1],
                ["A magical forest with glowing mushrooms, fantasy", "text, watermark", 40, 3.0, 1024, 1024, -1]
            ],
            inputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed]
        )

    # Events
    run_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed],
        outputs=[result_image, seed_output],
        api_visibility="public"
    )

    def clear_all():
        return None, -1

    clear_button.click(
        fn=clear_all,
        inputs=None,
        outputs=[result_image, seed_output],
        api_visibility="private"
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="purple"),
        css=css,
        footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
    )