app.py

import gradio as gr
import spaces
import torch
from diffusers import FluxPipeline
import os

# Initialize the model pipeline
@spaces.GPU
def load_model():
    """Load the FLUX.1-dev model for text-to-image generation"""
    try:
        # Use FLUX.1-dev model for text-to-image
        pipe = FluxPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-dev",
            torch_dtype=torch.bfloat16
        )
        pipe = pipe.to("cuda")
        return pipe
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

# Global model instance
model_pipe = None

@spaces.GPU
def generate_image(prompt, num_inference_steps=28, guidance_scale=3.5):
    """Generate image using FLUX.1-Kontext-dev with Picasso-style prompt enhancement"""
    global model_pipe
    
    if model_pipe is None:
        model_pipe = load_model()
        if model_pipe is None:
            return None, "Failed to load model"
    
    try:
        # Enhance prompt for Picasso style
        enhanced_prompt = f"{prompt}, in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors"
        
        # Generate image with FLUX parameters
        image = model_pipe(
            enhanced_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            width=1024,
            height=1024,
            generator=torch.Generator("cuda").manual_seed(42)
        ).images[0]
        
        return image, f"Generated: {enhanced_prompt}"
    except Exception as e:
        return None, f"Error generating image: {str(e)}"

# Create Gradio interface
with gr.Blocks(title="🎨 FLUX.1-dev-Picasso") as demo:
    gr.Markdown("# 🎨 FLUX.1-dev-Picasso")
    gr.Markdown("Generate high-quality images in Pablo Picasso's distinctive cubist style using FLUX.1-dev")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Describe what you want to see in Picasso's style...",
                lines=3
            )
            
            with gr.Row():
                steps_slider = gr.Slider(
                    minimum=20, maximum=50, value=28, step=1,
                    label="Inference Steps",
                    info="Higher values = better quality but slower generation (20-50)"
                )
                guidance_slider = gr.Slider(
                    minimum=1.0, maximum=10.0, value=3.5, step=0.1,
                    label="Guidance Scale", 
                    info="How closely to follow the prompt (1.0=loose, 10.0=strict)"
                )
            
            generate_btn = gr.Button("Generate Picasso-Style Image", variant="primary")
        
        with gr.Column():
            output_image = gr.Image(label="Generated Image", type="pil")
            output_text = gr.Textbox(label="Status", lines=2)
    
    # Example prompts as vertical buttons
    gr.Markdown("### 💡 Example Prompts")
    
    def set_example_1():
        return "A portrait of a woman"
    def set_example_2():
        return "A still life with fruits and bottles"
    def set_example_3():
        return "A musician playing guitar"
    def set_example_4():
        return "Two people having a conversation"
    def set_example_5():
        return "A cityscape with buildings"
    
    with gr.Column():
        example_btn_1 = gr.Button("A portrait of a woman", size="sm", scale=0, min_width=0)
        example_btn_2 = gr.Button("A still life with fruits and bottles", size="sm", scale=0, min_width=0)
        example_btn_3 = gr.Button("A musician playing guitar", size="sm", scale=0, min_width=0)
        example_btn_4 = gr.Button("Two people having a conversation", size="sm", scale=0, min_width=0)
        example_btn_5 = gr.Button("A cityscape with buildings", size="sm", scale=0, min_width=0)
    
    # Connect example buttons to prompt input
    example_btn_1.click(fn=set_example_1, outputs=prompt_input)
    example_btn_2.click(fn=set_example_2, outputs=prompt_input)
    example_btn_3.click(fn=set_example_3, outputs=prompt_input)
    example_btn_4.click(fn=set_example_4, outputs=prompt_input)
    example_btn_5.click(fn=set_example_5, outputs=prompt_input)
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, steps_slider, guidance_slider],
        outputs=[output_image, output_text]
    )
    
    # Attribution section
    gr.Markdown("""
    ---
    **Model Attribution:**  
    12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://huggingface.co/black-forest-labs/FLUX.1-pro)  
    [[non-commercial license]](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md) [[blog]](https://blackforestlabs.ai/announcing-black-forest-labs/) [[model]](https://huggingface.co/black-forest-labs/FLUX.1-dev)
    """)

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