| import gradio as gr |
| import torch |
| from diffusers import AutoPipelineForText2Image |
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| dtype = torch.float16 if device == "cuda" else torch.float32 |
| model_id = "stabilityai/sd-turbo" |
|
|
| print("[IMAGE] Loading", model_id, "on", device) |
| pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=dtype) |
| pipe = pipe.to(device) |
| print("[IMAGE] Pipeline ready.") |
|
|
|
|
| def generate_image(prompt): |
| full_prompt = prompt + ", high quality, realistic photo" |
| return pipe(prompt=full_prompt, num_inference_steps=2, guidance_scale=0.0).images[0] |
|
|
|
|
| with gr.Blocks(title="Tachelhit Image Generator") as demo: |
| prompt_input = gr.Textbox(label="Prompt") |
| generate_btn = gr.Button("Generate") |
| image_output = gr.Image(label="Result", type="filepath") |
| generate_btn.click(generate_image, prompt_input, image_output, api_name="/generate") |
|
|
| demo.launch(server_name="0.0.0.0", server_port=7860) |
