Dyushik/Edit-image

1from flask import Flask, request, jsonify, send_file, render_template2from huggingface_hub import InferenceClient3from PIL import Image, ImageOps4import io5import os6import json7import time8 9app = Flask(__name__)10 11class MultiImageAI:12    def __init__(self):13        self.client = InferenceClient(14            token=os.environ.get("HF_TOKEN")15        )16        self.model_name = "Qwen/Qwen-Image-Edit"17    18    def apply_smart_crop(self, image, crop_data):19        """Умный кроп с оптимизацией размера для лучшего качества AI"""20        try:21            # Сначала исправляем ориентацию22            try:23                image = ImageOps.exif_transpose(image)24            except:25                pass26            27            if not crop_data:28                # Если нет кропа - просто оптимизируем размер29                print("🔄 No crop - optimizing size for AI")30                return self.optimize_for_ai(image, 768)31            32            # Парсим данные кропа33            crop_dict = json.loads(crop_data)34            crop_x = float(crop_dict.get('cropX', 0))35            crop_y = float(crop_dict.get('cropY', 0))36            crop_width = float(crop_dict.get('cropWidth', image.width))37            crop_height = float(crop_dict.get('cropHeight', image.height))38            original_width = float(crop_dict.get('originalWidth', image.width))39            original_height = float(crop_dict.get('originalHeight', image.height))40            41            print(f"🎯 Crop data: ({crop_x:.0f}, {crop_y:.0f}) {crop_width:.0f}x{crop_height:.0f}")42            43            # Масштабируем координаты если нужно44            scale_x = image.width / original_width45            scale_y = image.height / original_height46            47            # Вычисляем абсолютные координаты кропа48            abs_x = int(crop_x * scale_x)49            abs_y = int(crop_y * scale_y)50            abs_width = int(crop_width * scale_x)51            abs_height = int(crop_height * scale_y)52            53            # Корректируем границы54            abs_x = max(0, abs_x)55            abs_y = max(0, abs_y)56            abs_width = min(image.width - abs_x, abs_width)57            abs_height = min(image.height - abs_y, abs_height)58            59            # Гарантируем минимальный размер для качества60            min_size = 51261            if abs_width < min_size or abs_height < min_size:62                # Используем центр с минимальным размером63                crop_size = min(image.width, image.height, min_size)64                abs_x = (image.width - crop_size) // 265                abs_y = (image.height - crop_size) // 266                abs_width = crop_size67                abs_height = crop_size68                print(f"🔄 Using center crop: {crop_size}x{crop_size}")69            70            print(f"✂️ Applying crop: ({abs_x}, {abs_y}) {abs_width}x{abs_height}")71            72            # Применяем кроп73            cropped = image.crop((abs_x, abs_y, abs_x + abs_width, abs_y + abs_height))74            print(f"✅ After crop: {cropped.size}")75            76            # Оптимизируем размер для AI77            optimized = self.optimize_for_ai(cropped, 768)78            print(f"🚀 Final for AI: {optimized.size}")79            80            return optimized81            82        except Exception as e:83            print(f"❌ Smart crop failed: {str(e)}")84            # Возвращаем просто оптимизированное изображение85            return self.optimize_for_ai(image, 768)86    87    def optimize_for_ai(self, image, target_size=768):88        """Оптимизация размера для лучшего качества AI"""89        try:90            width, height = image.size91            92            # Если изображение уже меньше целевого размера - оставляем как есть93            if width <= target_size and height <= target_size:94                # Все равно конвертируем в RGB95                if image.mode != 'RGB':96                    image = image.convert('RGB')97                return image98            99            # Вычисляем новый размер с сохранением пропорций100            ratio = min(target_size / width, target_size / height)101            new_width = int(width * ratio)102            new_height = int(height * ratio)103            104            # Ресайз с хорошим качеством105            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)106            107            # Конвертируем в RGB108            if image.mode != 'RGB':109                image = image.convert('RGB')110            111            print(f"🔄 Optimized: {width}x{height} -> {new_width}x{new_height}")112            return image113            114        except Exception as e:115            print(f"❌ Optimization failed: {str(e)}")116            return image117    118    def create_ai_prompt(self, original_prompt):119        """Создает оптимальный промпт для AI"""120        return f"""121        {original_prompt}122        123        Use the original image as exact reference.124        Preserve all facial features and details.125        Maintain original colors and lighting.126        Do not add new elements or modify existing ones.127        """128    129    def edit_single_image(self, image, prompt, crop_data=None):130        try:131            start_time = time.time()132            print(f"⚡ Starting AI processing...")133            134            # Умный кроп с оптимизацией135            processed_image = self.apply_smart_crop(image, crop_data)136            137            # Создаем промпт138            enhanced_prompt = self.create_ai_prompt(prompt)139            140            # Подготавливаем изображение для API141            img_byte_arr = io.BytesIO()142            processed_image.save(img_byte_arr, format='JPEG', quality=90, optimize=True)143            image_bytes = img_byte_arr.getvalue()144            145            print(f"📤 Sending to AI: {len(image_bytes)} bytes, size: {processed_image.size}")146            147            # Вызов API с низкой креативностью148            result = self.client.image_to_image(149                image=image_bytes,150                prompt=enhanced_prompt,151                model=self.model_name,152                strength=0.5,  # Минимальные изменения153            )154            155            processing_time = time.time() - start_time156            print(f"✅ AI processing completed in {processing_time:.1f}s")157            158            return result159            160        except Exception as e:161            print(f"❌ AI processing error: {str(e)}")162            return None163 164ai_processor = MultiImageAI()165 166@app.route('/')167def index():168    return render_template('index.html')169 170@app.route('/api/edit', methods=['POST'])171def api_edit():172    try:173        start_time = time.time()174        175        if 'image' not in request.files:176            return jsonify({"error": "No image file"}), 400177        178        image_file = request.files['image']179        prompt = request.form.get('prompt', '').strip()180        crop_data = request.form.get('crop_data', '')181        182        if not prompt:183            return jsonify({"error": "Please enter edit instructions"}), 400184        185        # Загружаем изображение186        image = Image.open(image_file.stream)187        image.load()188        print(f"📥 Original image: {image.size}")189        190        # Обрабатываем191        result = ai_processor.edit_single_image(image, prompt, crop_data)192        193        if result is None:194            return jsonify({"error": "AI processing failed. Please try again with a different image or prompt."}), 500195        196        # Сохраняем результат197        if hasattr(result, 'save'):198            output_image = result199        else:200            output_image = Image.open(io.BytesIO(result))201        202        output = io.BytesIO()203        output_image.save(output, format='PNG', optimize=True)204        output.seek(0)205        206        total_time = time.time() - start_time207        print(f"🎉 Total request time: {total_time:.1f}s")208        209        return send_file(output, mimetype='image/png')210        211    except Exception as e:212        print(f"❌ Server error: {str(e)}")213        return jsonify({"error": str(e)}), 500214 215@app.route('/health', methods=['GET'])216def health():217    return jsonify({218        "status": "healthy",219        "message": "PictoMagic AI is running",220        "model": "Qwen/Qwen-Image-Edit"221    })222 223if __name__ == '__main__':224    app.run(host='0.0.0.0', port=7860, debug=False)