File size: 4,304 Bytes
882e052 c8f1f54 4fe456a 61b9893 df1e443 c8f1f54 61b9893 df1e443 882e052 df1e443 c8f1f54 61b9893 df1e443 61b9893 882e052 5d4b17c 61b9893 9aeab3c 61b9893 df1e443 c8f1f54 61b9893 882e052 087c578 df1e443 61b9893 df1e443 61b9893 faa166e e3617aa 61b9893 882e052 61b9893 df1e443 882e052 61b9893 882e052 faa166e 61b9893 882e052 61b9893 087c578 61b9893 087c578 5d4b17c 087c578 61b9893 087c578 dc45b2e 5d4b17c 61b9893 087c578 5d4b17c 882e052 61b9893 087c578 882e052 c8f1f54 61b9893 |
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
import os
import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import io
import torch
from diffusers import DiffusionPipeline
# ===== CONFIGURATION =====
MODEL_NAME = "HiDream-ai/HiDream-I1-Full"
WATERMARK_TEXT = "SelamGPT"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
TORCH_DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
# ===== MODEL LOADING =====
@gr.Cache() # Cache model between generations
def load_model():
pipe = DiffusionPipeline.from_pretrained(
MODEL_NAME,
torch_dtype=TORCH_DTYPE
).to(DEVICE)
# Optimizations
if DEVICE == "cuda":
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_attention_slicing()
return pipe
pipe = load_model()
# ===== WATERMARK FUNCTION =====
def add_watermark(image):
"""Add watermark with optimized PNG output"""
try:
draw = ImageDraw.Draw(image)
font_size = 24
try:
font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
except:
font = ImageFont.load_default(font_size)
text_width = draw.textlength(WATERMARK_TEXT, font=font)
x = image.width - text_width - 10
y = image.height - 34
draw.text((x+1, y+1), WATERMARK_TEXT, font=font, fill=(0, 0, 0, 128))
draw.text((x, y), WATERMARK_TEXT, font=font, fill=(255, 255, 255))
# Convert to optimized PNG
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG', optimize=True, quality=85)
img_byte_arr.seek(0)
return Image.open(img_byte_arr)
except Exception as e:
print(f"Watermark error: {str(e)}")
return image
# ===== IMAGE GENERATION =====
def generate_image(prompt):
if not prompt.strip():
return None, "⚠️ Please enter a prompt"
try:
# Generate image (1024x1024 by default)
image = pipe(
prompt,
num_inference_steps=30,
guidance_scale=7.5
).images[0]
# Add watermark
watermarked = add_watermark(image)
return watermarked, "✔️ Generation successful"
except torch.cuda.OutOfMemoryError:
return None, "⚠️ Out of memory! Try a simpler prompt"
except Exception as e:
return None, f"⚠️ Error: {str(e)[:200]}"
# ===== GRADIO THEME =====
theme = gr.themes.Default(
primary_hue="emerald",
secondary_hue="amber",
font=[gr.themes.GoogleFont("Poppins"), "Arial", "sans-serif"]
)
# ===== GRADIO INTERFACE =====
with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
gr.Markdown("""
# 🎨 SelamGPT Image Generator
*Powered by HiDream-I1-Full (1024x1024 PNG output)*
""")
with gr.Row():
with gr.Column(scale=3):
prompt_input = gr.Textbox(
label="Describe your image",
placeholder="A futuristic Ethiopian city with flying cars...",
lines=3,
max_lines=5
)
with gr.Row():
generate_btn = gr.Button("Generate Image", variant="primary")
clear_btn = gr.Button("Clear")
gr.Examples(
examples=[
["An ancient Aksumite warrior in cyberpunk armor, 4k detailed"],
["Traditional Ethiopian coffee ceremony in zero gravity"],
["Portrait of a Habesha queen with golden jewelry"]
],
inputs=prompt_input
)
with gr.Column(scale=2):
output_image = gr.Image(
label="Generated Image",
type="pil",
format="png",
height=512
)
status_output = gr.Textbox(
label="Status",
interactive=False
)
generate_btn.click(
fn=generate_image,
inputs=prompt_input,
outputs=[output_image, status_output],
queue=True
)
clear_btn.click(
fn=lambda: [None, ""],
outputs=[output_image, status_output]
)
if __name__ == "__main__":
demo.queue(max_size=2)
demo.launch(server_name="0.0.0.0", server_port=7860) |