Update app.py

app.py

@@ -1,55 +1,71 @@
 import torch
-import cv2
 import numpy as np
-import gradio as gr
-from diffusers import StableDiffusionPipeline, DDIMScheduler, IPAdapterFaceIDPipeline
-from insightface.app import FaceAnalysis
-from huggingface_hub import hf_hub_download
+from flask import Flask, request, jsonify
+from diffusers import DiffusionPipeline
+from PIL import Image
+import base64
+import io
+import gc
+import os
 
-device = "cpu"
-dtype = torch.float32
+app = Flask(__name__)
 
-# Initialize Face Detection
-face_app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
-face_app.prepare(ctx_id=0, det_size=(320, 320))
+# Device set to CPU and optimize threading
+device = "cpu"
+torch.set_num_threads(max(1, os.cpu_count() or 4))  # Use all available CPU threads
+torch.set_num_interop_threads(max(1, os.cpu_count() or 4))
 
-# Download IP-Adapter FaceID Weights
-ip_adapter_path = hf_hub_download(
-    repo_id="h94/IP-Adapter",
-    filename="ip-adapter_sd15_faceid.bin",
-    subfolder="models"
+# Load model on CPU with optimizations
+pipe = DiffusionPipeline.from_pretrained(
+    "stabilityai/sdxl-turbo",
+    use_safetensors=True,
+    torch_dtype=torch.float16  # FP16 for memory efficiency
 )
+pipe = pipe.to(device)
+pipe.eval()  # Evaluation mode to save memory
+pipe.unet.enable_model_cpu_offload()  # Offload U-Net to CPU memory if possible
 
-# Load Pipeline
-pipe = IPAdapterFaceIDPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=dtype)
-pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
-pipe.load_ip_adapter(ip_adapter_path)
-pipe.to(device)
-
-def generate(upload_image, prompt):
-    if upload_image is None:
-        return "No Image Uploaded", None
+def infer(prompt, steps=1, seed=0):  # Reduced to 1 step for speed
+    if not prompt or len(prompt.split()) > 77:  # Enforce 77 token limit
+        return "Prompt missing or exceeds 77 tokens!", 0
     
-    img = cv2.cvtColor(np.array(upload_image), cv2.COLOR_RGB2BGR)
-    faces = face_app.get(img)
+    # Set seed
+    generator = torch.Generator(device=device).manual_seed(seed) if seed != 0 else torch.Generator(device=device).manual_seed(np.random.randint(0, 2**32 - 1))
 
-    if len(faces) == 0:
-        return "No Face Detected", None
-    
-    face_emb = torch.tensor(faces[0].normed_embedding).unsqueeze(0).to(device)
-    result = pipe(prompt=prompt, face_embeds=face_emb, num_inference_steps=30).images[0]
-    return "Image Generated Successfully", result
-
-demo = gr.Interface(
-    generate,
-    inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Textbox(label="Prompt")
-    ],
-    outputs=[
-        gr.Textbox(label="Status"),
-        gr.Image(label="Generated Image")
-    ]
-)
+    # Generate image on CPU with minimal steps
+    with torch.no_grad():  # Disable gradient computation
+        image = pipe(
+            prompt=prompt,
+            num_inference_steps=steps,  # Reduced to 1 step
+            guidance_scale=0.0,
+            height=512,
+            width=512,
+            generator=generator,
+            output_type="pil",
+            num_images_per_prompt=1
+        ).images[0]
+
+    gc.collect()  # Free memory immediately
+    torch.cuda.empty_cache()  # No-op on CPU but included for consistency
+    return image, seed
+
+@app.route('/generate', methods=['POST'])
+def generate():
+    prompt = request.form.get('prompt')
+    steps = int(request.form.get('steps', 1))  # Default to 1 step
+    seed = int(request.form.get('seed', 0))    # Default to 0 or random
+
+    result, seed_used = infer(prompt, steps, seed)
+
+    if isinstance(result, str):  # Error message
+        return jsonify({'error': result, 'seed': seed_used}), 400
+
+    # Convert image to base64
+    buffered = io.BytesIO()
+    result.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    return jsonify({'image': img_str, 'seed': seed_used})
 
-demo.launch()
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=8000)