Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from diffusers import FluxPipeline | |
| import os | |
| import traceback | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| MODEL_IDS = { | |
| "FLUX.1 Schnell": "black-forest-labs/FLUX.1-schnell", | |
| "FLUX.1 DEV": "black-forest-labs/FLUX.1-dev", | |
| "FLUX.1 Kontext": "black-forest-labs/FLUX.1-kontext" | |
| } | |
| PIPELINES = {} | |
| def get_pipeline(model_name): | |
| try: | |
| if model_name not in PIPELINES: | |
| print(f"Loading pipeline for {model_name} ({MODEL_IDS[model_name]})...") | |
| pipe = FluxPipeline.from_pretrained( | |
| MODEL_IDS[model_name], | |
| token=HF_TOKEN, # Correct argument for Hugging Face tokens | |
| torch_dtype=torch.float32 # Use float32 for CPU compatibility | |
| ) | |
| pipe.enable_model_cpu_offload() # Offload to CPU if needed | |
| PIPELINES[model_name] = pipe | |
| print(f"Pipeline for {model_name} loaded successfully.") | |
| return PIPELINES[model_name] | |
| except Exception as e: | |
| print(f"Error loading pipeline for {model_name}: {e}") | |
| traceback.print_exc() | |
| raise RuntimeError(f"Failed to load {model_name}: {e}") | |
| def generate_image(model_name, prompt, height, width, guidance_scale, steps, seed): | |
| print(f"generate_image called with model: {model_name}, prompt: '{prompt}', height: {height}, width: {width}, guidance_scale: {guidance_scale}, steps: {steps}, seed: {seed}") | |
| if not prompt or not prompt.strip(): | |
| print("Prompt is empty.") | |
| return None | |
| try: | |
| pipe = get_pipeline(model_name) | |
| generator = torch.Generator("cpu").manual_seed(int(seed)) | |
| images = pipe( | |
| prompt, | |
| height=int(height), | |
| width=int(width), | |
| guidance_scale=float(guidance_scale), | |
| num_inference_steps=int(steps), | |
| max_sequence_length=512, | |
| generator=generator | |
| ).images | |
| print(f"Image generated successfully for prompt: '{prompt}'") | |
| return images[0] | |
| except Exception as e: | |
| print(f"Error during image generation: {e}") | |
| traceback.print_exc() | |
| return None # Gradio will show a blank image | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# FLUX Text-to-Image Generator\nSelect a model and enter your prompt.") | |
| with gr.Row(): | |
| model_selector = gr.Dropdown( | |
| choices=list(MODEL_IDS.keys()), | |
| value="FLUX.1 Schnell", | |
| label="Choose FLUX Model" | |
| ) | |
| prompt = gr.Textbox(label="Prompt", placeholder="Describe the image you want to generate") | |
| with gr.Row(): | |
| height = gr.Slider(256, 1024, value=1024, step=64, label="Height") | |
| width = gr.Slider(256, 1024, value=1024, step=64, label="Width") | |
| with gr.Row(): | |
| guidance_scale = gr.Slider(1, 10, value=3.5, step=0.1, label="Guidance Scale") | |
| steps = gr.Slider(10, 100, value=50, step=1, label="Steps") | |
| seed = gr.Slider(0, 10000, value=0, step=1, label="Seed") | |
| generate_btn = gr.Button("Generate Image") | |
| output_image = gr.Image(type="pil", label="Generated Image") | |
| generate_btn.click( | |
| fn=generate_image, | |
| inputs=[model_selector, prompt, height, width, guidance_scale, steps, seed], | |
| outputs=output_image | |
| ) | |
| if __name__ == "__main__": | |
| print("Starting FLUX Text-to-Image Gradio app...") | |
| demo.launch() | |