Update app.py

app.py

@@ -1,64 +1,110 @@
+import os
+import requests
 import gradio as gr
-from diffusers import AutoPipelineForText2Image
 from PIL import Image, ImageDraw, ImageFont
-import torch
-import random
+import io
+import time
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# 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
 
-pipe = AutoPipelineForText2Image.from_pretrained(
-    "stabilityai/sdxl-turbo",
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    variant="fp16" if device == "cuda" else None
-)
-pipe = pipe.to(device)
+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
 
-MAX_SEED = 2**32 - 1
+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."
 
-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
-
-def generate(prompt, seed, randomize_seed):
-    if randomize_seed or seed == 0:
-        seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(device=device).manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        num_inference_steps=2,
-        guidance_scale=0.0,
-        generator=generator,
-    ).images[0]
-
-    image = add_watermark(image)
-    return image, seed
-
-examples = [
-    "Futuristic Ethiopian city at sunset, detailed, cinematic",
-    "አንድ ጫካ ውስጥ የተሰወረ የሳይባር ከተማ በመስሎ የታሰረ ቀይ ብርሃን",
-]
-
-with gr.Blocks() as demo:
-    gr.Markdown("## 🎨 SelamGPT - Super Fast Text-to-Image Generator")
-
-    prompt = gr.Textbox(label="Prompt", placeholder="Type your idea in English or Amharic")
-    run = 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.click(fn=generate, inputs=[prompt, seed, randomize_seed], outputs=[result, seed])
+# 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()
+    demo.launch(server_name="0.0.0.0", server_port=7860)