import gradio as gr import spaces import torch from diffusers import FluxPipeline import os # Initialize model configuration dtype = torch.bfloat16 device = "cuda" if torch.cuda.is_available() else "cpu" # Global model instance model_pipe = None # 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=dtype ) pipe = pipe.to(device) if torch.cuda.is_available(): torch.cuda.empty_cache() return pipe except Exception as e: print(f"Error loading model: {e}") return None @spaces.GPU(duration=100) 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, "🚀 Loading AI model... Please wait a moment and try again!" try: # Enhance prompt for Picasso style (avoid duplication) picasso_style = "in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors" if picasso_style not in prompt: enhanced_prompt = f"{prompt}, {picasso_style}" else: enhanced_prompt = prompt # 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(device).manual_seed(42) ).images[0] # Clean up GPU memory after generation if torch.cuda.is_available(): torch.cuda.empty_cache() return image, f"Generated: {enhanced_prompt}" except Exception as e: error_msg = str(e).lower() if "gpu" in error_msg or "cuda" in error_msg or "memory" in error_msg: return None, "🔄 GPU temporarily unavailable. Please try again in a moment!" elif "timeout" in error_msg or "aborted" in error_msg: return None, "⏱️ Generation timed out. Please try again with fewer inference steps or try again later!" else: return None, f"⚠️ Something went wrong. Please try again! (Technical details: {str(e)})" # Create Gradio interface with gr.Blocks(title="🎨 FLUX.1-dev Picasso Edition") as demo: gr.Markdown("# 🎨 FLUX.1-dev Picasso Edition") gr.Markdown("Generate high-quality images in Pablo Picasso's distinctive cubist style using FLUX.1-dev") gr.Markdown(""" 12B param rectified flow transformer guidance-distilled from [FLUX.1 [pro]](https://bfl.ai/) [[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) """) 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(): guidance_slider = gr.Slider( minimum=1.0, maximum=10.0, value=3.5, step=0.1, label="Guidance Scale", info="Higher values = more faithful to prompt\nLower values = more creative interpretation" ) steps_slider = gr.Slider( minimum=20, maximum=50, value=28, step=1, label="Inference Steps", info="Higher values = better quality but slower\nLower values = faster but less refined" ) generate_btn = gr.Button("Generate Picasso-Style Image", variant="primary") # Example prompts as vertical buttons gr.Markdown("### 💡 Example Prompts") def set_example_1(): return "A portrait of a woman, in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors" def set_example_2(): return "A still life with fruits and bottles, in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors" def set_example_3(): return "A musician playing guitar, in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors" def set_example_4(): return "Two people having a conversation, in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors" def set_example_5(): return "A cityscape with buildings, in the style of Pablo Picasso, cubist painting, abstract art, geometric shapes, fragmented forms, bold colors" 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) with gr.Column(): output_image = gr.Image(label="Generated Image", type="pil", height=600, format="png") output_text = gr.Textbox(label="Status", lines=3) # Connect example buttons to only set prompt (no generation) 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] ) if __name__ == "__main__": demo.launch()