Spaces:
Sleeping
Sleeping
| 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 | |
| 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 | |
| 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() |