Update app.py

app.py

@@ -1,23 +1,20 @@
 import gradio as gr
-from diffusers import StableDiffusionPipeline
+from diffusers import AutoPipelineForText2Image
 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 = 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)
 
 MAX_SEED = 2**32 - 1
 
-# Add "SelamGPT" watermark to image
 def add_watermark(image):
     draw = ImageDraw.Draw(image)
     font = ImageFont.load_default()
@@ -28,37 +25,40 @@ def add_watermark(image):
     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)
+    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 = pipe(prompt=prompt, generator=generator).images[0]
     image = add_watermark(image)
     return image, seed
 
 examples = [
-    "አዲስ ዘመናዊ ከተማ በደመና መካከል", 
-    "A futuristic Ethiopian skyline at night", 
-    "አንድ የፍቅር ድራማ በተራራ በታች", 
+    "Futuristic Ethiopian city at sunset, detailed, cinematic",
+    "አንድ ጫካ ውስጥ የተሰወረ የሳይባር ከተማ በመስሎ የታሰረ ቀይ ብርሃን",
 ]
 
 with gr.Blocks() as demo:
-    gr.Markdown("# SelamGPT ፡ Text-to-Image Generator 🖼️\nGenerate creative visuals from your imagination!")
+    gr.Markdown("## 🎨 SelamGPT - Super Fast Text-to-Image Generator")
 
-    prompt = gr.Textbox(label="Image Prompt (in Amharic or English)", placeholder="e.g. አዲስ ከተማ በጨለማ ውስጥ")
-    run_button = gr.Button("Generate")
+    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):
+    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)
+        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])
+    run.click(fn=generate, inputs=[prompt, seed, randomize_seed], outputs=[result, seed])
 
 if __name__ == "__main__":
     demo.launch()