File size: 3,734 Bytes
882e052 c8f1f54 4fe456a 882e052 c8f1f54 882e052 c8f1f54 882e052 c8f1f54 882e052 4fe456a 882e052 c8f1f54 882e052 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import os
import requests
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import io
import time
# Configuration
HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0" # High-quality model
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
WATERMARK_TEXT = "SelamGPT"
MAX_RETRIES = 3
def add_watermark(image_bytes):
"""Add watermark to generated image"""
try:
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
draw = ImageDraw.Draw(image)
# Use default font (no need for external files)
try:
font = ImageFont.truetype("arialbd.ttf", 40)
except:
font = ImageFont.load_default(size=40) # Fallback
# Calculate text position (bottom-right corner)
bbox = draw.textbbox((0, 0), WATERMARK_TEXT, font=font)
text_width = bbox[2] - bbox[0]
text_height = bbox[3] - bbox[1]
margin = 20
position = (image.width - text_width - margin, image.height - text_height - margin)
# Draw watermark with semi-transparent white text and black outline
draw.text(
position,
WATERMARK_TEXT,
font=font,
fill=(255, 255, 255, 128),
stroke_width=2,
stroke_fill=(0, 0, 0, 128)
)
return image
except Exception as e:
print(f"Watermark error: {str(e)}")
return Image.open(io.BytesIO(image_bytes)) # Return original if watermark fails
def generate_image(prompt):
"""Generate image with retry logic"""
if not prompt.strip():
return "Error: Please enter a valid prompt"
for attempt in range(MAX_RETRIES):
try:
response = requests.post(
API_URL,
headers=headers,
json={"inputs": prompt, "options": {"wait_for_model": True}},
timeout=30
)
if response.status_code == 200:
return add_watermark(response.content)
elif response.status_code == 503: # Model loading
time.sleep(10 * (attempt + 1)) # Exponential backoff
continue
else:
return f"API Error: {response.text}"
except requests.Timeout:
return "Error: Request timed out (30s)"
except Exception as e:
return f"Unexpected error: {str(e)}"
return "Failed after multiple attempts. Try again later."
# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown("""
# 🎨 SelamGPT Image Generator
*Powered by Stable Diffusion XL with built-in watermark*
""")
with gr.Row():
with gr.Column():
prompt_input = gr.Textbox(
label="Describe your image",
placeholder="A futuristic city at sunset...",
lines=3
)
generate_btn = gr.Button("Generate", variant="primary")
examples = gr.Examples(
examples=["A cute robot reading a book", "Ethiopian landscape in oil painting style"],
inputs=prompt_input
)
with gr.Column():
output_image = gr.Image(label="Generated Image", height=512)
error_output = gr.Textbox(label="Status", visible=False)
generate_btn.click(
fn=generate_image,
inputs=prompt_input,
outputs=[output_image, error_output],
show_progress="full"
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860) |