Update app.py

app.py

@@ -1,26 +1,37 @@
 import os
-import requests
 import gradio as gr
 from PIL import Image, ImageDraw, ImageFont
 import io
-import time
-from concurrent.futures import ThreadPoolExecutor
+import torch
+from diffusers import DiffusionPipeline
 
 # ===== CONFIGURATION =====
-HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
-MODEL_NAME = "stabilityai/stable-diffusion-xl-base-1.0"
-API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
-headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
+MODEL_NAME = "HiDream-ai/HiDream-I1-Full"
 WATERMARK_TEXT = "SelamGPT"
-MAX_RETRIES = 3
-TIMEOUT = 60
-EXECUTOR = ThreadPoolExecutor(max_workers=2)
+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_bytes):
+def add_watermark(image):
     """Add watermark with optimized PNG output"""
     try:
-        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         draw = ImageDraw.Draw(image)
         
         font_size = 24
@@ -43,49 +54,29 @@ def add_watermark(image_bytes):
         return Image.open(img_byte_arr)
     except Exception as e:
         print(f"Watermark error: {str(e)}")
-        return Image.open(io.BytesIO(image_bytes))
+        return image
 
 # ===== IMAGE GENERATION =====
 def generate_image(prompt):
     if not prompt.strip():
         return None, "⚠️ Please enter a prompt"
     
-    def api_call():
-        return requests.post(
-            API_URL,
-            headers=headers,
-            json={
-                "inputs": prompt,
-                "parameters": {
-                    "height": 1024,
-                    "width": 1024,
-                    "num_inference_steps": 30
-                },
-                "options": {"wait_for_model": True}
-            },
-            timeout=TIMEOUT
-        )
-    
-    for attempt in range(MAX_RETRIES):
-        try:
-            future = EXECUTOR.submit(api_call)
-            response = future.result()
-            
-            if response.status_code == 200:
-                return add_watermark(response.content), "✔️ Generation successful"
-            elif response.status_code == 503:
-                wait_time = (attempt + 1) * 15
-                print(f"Model loading, waiting {wait_time}s...")
-                time.sleep(wait_time)
-                continue
-            else:
-                return None, f"⚠️ API Error: {response.text[:200]}"
-        except requests.Timeout:
-            return None, f"⚠️ Timeout: Model took >{TIMEOUT}s to respond"
-        except Exception as e:
-            return None, f"⚠️ Unexpected error: {str(e)[:200]}"
+    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"
     
-    return None, "⚠️ Failed after multiple attempts. Please try later."
+    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(
@@ -98,7 +89,7 @@ theme = gr.themes.Default(
 with gr.Blocks(theme=theme, title="SelamGPT Image Generator") as demo:
     gr.Markdown("""
     # 🎨 SelamGPT Image Generator
-    *Powered by Stable Diffusion XL (1024x1024 PNG output)*
+    *Powered by HiDream-I1-Full (1024x1024 PNG output)*
     """)
     
     with gr.Row():