Update app.py

app.py

@@ -4,29 +4,35 @@ import random
 from PIL import Image, ImageDraw, ImageFont
 import torch
 from diffusers import DiffusionPipeline
+import io
 
 # ===== CONFIG =====
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if device == "cuda" else torch.float32
 model_repo_id = "stabilityai/sdxl-turbo"
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype, variant="fp16" if device == "cuda" else None)
+pipe = DiffusionPipeline.from_pretrained(
+    model_repo_id,
+    torch_dtype=torch_dtype,
+    variant="fp16" if device == "cuda" else None
+)
 pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
-IMAGE_SIZE = 1024
+IMAGE_WIDTH = 768
+IMAGE_HEIGHT = 768
 WATERMARK_TEXT = "SelamGPT"
 
 # ===== WATERMARK FUNCTION =====
 def add_watermark(image):
     draw = ImageDraw.Draw(image)
-    font_size = 24
+    font_size = int(image.width * 0.03)
     try:
         font = ImageFont.truetype("Roboto-Bold.ttf", font_size)
     except:
         font = ImageFont.load_default()
     text_width = draw.textlength(WATERMARK_TEXT, font=font)
-    x = image.width - text_width - 10
-    y = image.height - 34
+    x = image.width - text_width - 12
+    y = image.height - font_size - 10
     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))
     return image
@@ -51,15 +57,18 @@ def generate(
     result = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
-        width=IMAGE_SIZE,
-        height=IMAGE_SIZE,
+        width=IMAGE_WIDTH,
+        height=IMAGE_HEIGHT,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
         generator=generator,
     ).images[0]
 
-    image = add_watermark(result)
-    return image, seed
+    watermarked = add_watermark(result)
+    buffer = io.BytesIO()
+    watermarked.convert("RGB").save(buffer, format="JPEG", quality=70)
+    buffer.seek(0)
+    return Image.open(buffer), seed
 
 # ===== EXAMPLES =====
 examples = [
@@ -85,7 +94,7 @@ with gr.Blocks(css=css, title="SelamGPT Turbo Generator") as demo:
             )
             generate_btn = gr.Button("Generate", variant="primary")
 
-        output_image = gr.Image(label="Generated Image", type="pil", format="png")
+        output_image = gr.Image(label="Generated Image", type="pil", format="jpeg")
         seed_display = gr.Textbox(label="Seed Used", interactive=False)
 
         with gr.Accordion("⚙️ Advanced Settings", open=False):
@@ -94,7 +103,7 @@ with gr.Blocks(css=css, title="SelamGPT Turbo Generator") as demo:
             seed = gr.Slider(0, MAX_SEED, step=1, label="Seed", value=0)
 
             guidance_scale = gr.Slider(0.0, 10.0, step=0.1, label="Guidance Scale", value=0.0)
-            num_inference_steps = gr.Slider(1, 50, step=1, label="Inference Steps", value=2)
+            num_inference_steps = gr.Slider(1, 10, step=1, label="Inference Steps", value=2)
 
         gr.Examples(examples=examples, inputs=[prompt])
 
@@ -112,4 +121,4 @@ with gr.Blocks(css=css, title="SelamGPT Turbo Generator") as demo:
         )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()