Update app.py

app.py

@@ -6,64 +6,54 @@ from PIL import Image
 import os
 import time
 
-# ✅ Set base URL manually (since you're not using secrets)
+# Branding
+WATERMARK_PATH = "CodeHubb.png"
 BASE_URL = "https://image.pollinations.ai/prompt/"
 
-# ✅ Load watermark logo (must be named 'CodeHubb.png' and present in root)
-watermark_path = "CodeHubb.png"
-watermark = Image.open(watermark_path).convert("RGBA")
-
-# ✅ Reduce watermark opacity to 50%
-alpha = watermark.split()[3]
-alpha = alpha.point(lambda p: int(p * 0.5))
-watermark.putalpha(alpha)
-
-# ✅ Image generation function
 def generate_image(prompt):
     if not prompt:
         raise gr.Error("Please enter a prompt.")
 
     seed = random.randint(1, 999999)
-    model = "flux"  # Default model
-
+    model = "flux"
     url = f"{BASE_URL}{prompt}?width=1024&height=1024&seed={seed}&nologo=true&model={model}"
 
     try:
         response = requests.get(url, stream=True)
         response.raise_for_status()
-
         img = Image.open(BytesIO(response.content)).convert("RGBA")
 
-        # ✅ Resize watermark to fixed width
-        wm_width = 200
-        wm_ratio = wm_width / watermark.width
-        wm_size = (wm_width, int(watermark.height * wm_ratio))
-        resized_wm = watermark.resize(wm_size)
+        # Watermark
+        if os.path.exists(WATERMARK_PATH):
+            watermark = Image.open(WATERMARK_PATH).convert("RGBA")
+            watermark = watermark.resize((150, 50))
+            watermark.putalpha(128)  # 50% opacity
+            img.paste(watermark, (10, 10), watermark)
 
-        # ✅ Paste watermark at top-left
-        img.paste(resized_wm, (10, 10), resized_wm)
+        # Convert to JPG and Save
+        output_path = "generated_image.jpg"
+        rgb_img = img.convert("RGB")
+        rgb_img.save(output_path, "JPEG")
 
-        # ✅ Convert to RGB JPEG
-        final_img = img.convert("RGB")
-        buffer = BytesIO()
-        final_img.save(buffer, format="JPEG")
-        buffer.seek(0)
-
-        return Image.open(buffer)
+        return rgb_img, output_path
 
     except Exception as e:
         raise gr.Error(f"Unexpected error: {str(e)}")
 
-# ✅ Gradio Interface
-with gr.Blocks(css="custom_hide_footer.css") as demo:
-    gr.Markdown("## CodeHubb AI Image Generator 🖼️\nPowered by Flux Model\n*Design. Develop. Dominate.*")
-
-    prompt_input = gr.Textbox(label="Prompt", placeholder="e.g., A futuristic city at sunset", lines=2)
+# Gradio UI
+with gr.Blocks(css=".footer {display: none !important;}") as demo:
+    gr.Markdown("## CodeHubb AI Image Generator 🧠\n**Design. Develop. Dominate.**")
 
+    prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g. horse running in desert")
     generate_btn = gr.Button("Generate Image")
-    output_image = gr.Image(label="Generated Image", type="pil")
 
-    generate_btn.click(generate_image, inputs=prompt_input, outputs=output_image)
+    image_output = gr.Image(label="Generated Image")
+    download_output = gr.File(label="Download JPG")
+
+    generate_btn.click(
+        fn=generate_image,
+        inputs=prompt_input,
+        outputs=[image_output, download_output],
+    )
 
-demo.queue()
 demo.launch()