snackshell's picture
Update app.py
4fe456a verified
raw
history blame
2.14 kB
import gradio as gr
from diffusers import StableDiffusionPipeline
from PIL import Image, ImageDraw, ImageFont
import torch
import random
# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
pipe = StableDiffusionPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch_dtype,
revision="fp16" if device == "cuda" else None
)
pipe = pipe.to(device)
MAX_SEED = 2**32 - 1
# Add "SelamGPT" watermark to image
def add_watermark(image):
draw = ImageDraw.Draw(image)
font = ImageFont.load_default()
text = "SelamGPT"
margin = 10
x = image.width - draw.textlength(text, font=font) - margin
y = image.height - 20
draw.text((x, y), text, font=font, fill=(255, 255, 255))
return image
# Main generation function
def generate(prompt, seed, randomize_seed):
if randomize_seed or seed == 0:
seed = random.randint(0, MAX_SEED)
generator = torch.Generator(device).manual_seed(seed)
image = pipe(prompt=prompt, generator=generator).images[0]
image = add_watermark(image)
return image, seed
examples = [
"αŠ α‹²αˆ΅ α‹˜αˆ˜αŠ“α‹Š αŠ¨α‰°αˆ› α‰ α‹°αˆ˜αŠ“ መካከል",
"A futuristic Ethiopian skyline at night",
"αŠ αŠ•α‹΅ α‹¨αα‰…αˆ­ α‹΅αˆ«αˆ› α‰ α‰°αˆ«αˆ« በታች",
]
with gr.Blocks() as demo:
gr.Markdown("# SelamGPT ፑ Text-to-Image Generator πŸ–ΌοΈ\nGenerate creative visuals from your imagination!")
prompt = gr.Textbox(label="Image Prompt (in Amharic or English)", placeholder="e.g. αŠ α‹²αˆ΅ αŠ¨α‰°αˆ› α‰ αŒ¨αˆˆαˆ› α‹αˆ΅αŒ₯")
run_button = gr.Button("Generate")
result = gr.Image(label="Generated Image")
with gr.Accordion("βš™οΈ Advanced Settings", open=False):
seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
randomize_seed = gr.Checkbox(label="🎲 Randomize seed", value=True)
gr.Examples(examples=examples, inputs=[prompt])
run_button.click(fn=generate, inputs=[prompt, seed, randomize_seed], outputs=[result, seed])
if __name__ == "__main__":
demo.launch()