Spaces:
Sleeping
Sleeping
| import torch | |
| import gradio as gr | |
| from diffusers import AutoPipelineForText2Image | |
| import gc | |
| # Available models | |
| AVAILABLE_MODELS = [ | |
| "stabilityai/stable-diffusion-xl-base-1.0", | |
| "stabilityai/sd-turbo", | |
| "Lykon/dreamshaper-8", | |
| "runwayml/stable-diffusion-v1-5", | |
| ] | |
| # Global pipeline | |
| pipe = None | |
| def load_model(model_id): | |
| """Load a new model, clearing memory first""" | |
| global pipe | |
| try: | |
| # Clear existing model from memory | |
| if pipe is not None: | |
| del pipe | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| # Load new pipeline | |
| pipe = AutoPipelineForText2Image.from_pretrained( | |
| model_id, | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32 | |
| ) | |
| if torch.cuda.is_available(): | |
| pipe = pipe.to("cuda") | |
| return f"✓ Loaded: {model_id}" | |
| except Exception as e: | |
| pipe = None | |
| return f"✗ Error loading model: {str(e)}" | |
| # Load initial model | |
| load_model(AVAILABLE_MODELS[0]) | |
| def generate(model_id, prompt, negative_prompt, steps, guidance): | |
| global pipe | |
| status = "" | |
| if not prompt: | |
| return None, "Prompt required." | |
| # Load model if not loaded or changed | |
| if pipe is None or getattr(pipe, 'model_id', None) != model_id: | |
| status = load_model(model_id) | |
| # Attach model_id to pipe for tracking | |
| if pipe is not None: | |
| pipe.model_id = model_id | |
| else: | |
| status = f"✓ Loaded: {model_id}" | |
| if pipe is None: | |
| return None, status | |
| image = pipe( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt or None, | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(guidance), | |
| ).images[0] | |
| return image, status | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # ⚡ Text-to-Image Generator | |
| Select a model, type a prompt, tweak the sliders, and hit **Generate**. | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| model_dropdown = gr.Dropdown( | |
| choices=AVAILABLE_MODELS, | |
| value=AVAILABLE_MODELS[0], | |
| label="Select Model", | |
| interactive=True | |
| ) | |
| model_status = gr.Textbox( | |
| label="Model Status", | |
| value=f"✓ Loaded: {AVAILABLE_MODELS[0]}", | |
| interactive=False | |
| ) | |
| prompt = gr.Textbox( | |
| label="Prompt", | |
| lines=2, | |
| value="a cute robot teaching about Hugging Face Spaces, digital art, colorful" | |
| ) | |
| negative_prompt = gr.Textbox( | |
| label="Negative prompt (optional)", | |
| lines=1, | |
| placeholder="blurry, low quality, text" | |
| ) | |
| steps = gr.Slider( | |
| minimum=1, | |
| maximum=50, | |
| value=2, | |
| step=1, | |
| label="Inference steps" | |
| ) | |
| guidance = gr.Slider( | |
| minimum=0.0, | |
| maximum=20.0, | |
| value=1.5, | |
| step=0.1, | |
| label="Guidance scale (strength of text conditioning)" | |
| ) | |
| generate_btn = gr.Button("Generate 🚀") | |
| with gr.Column(scale=3): | |
| output = gr.Image(label="Generated image", height=512) | |
| generate_btn.click( | |
| fn=generate, | |
| inputs=[model_dropdown, prompt, negative_prompt, steps, guidance], | |
| outputs=[output, model_status] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |