Update app.py

app.py

@@ -1,13 +1,11 @@
 from flask import Flask, request, jsonify
 import torch
 from transformers import (
-    UNet2DConditionModel,
-    AutoTokenizer,
-    CLIPTextModel,
-    CLIPTextModelWithProjection,
-    CLIPVisionModelWithProjection,
-    AutoencoderKL,
-    DDPMScheduler
+    UNet2DConditionModel, 
+    AutoTokenizer, 
+    CLIPTextModel, 
+    CLIPTextModelWithProjection, 
+    CLIPVisionModelWithProjection
 )
 from PIL import Image
 import base64
@@ -28,71 +26,70 @@ UNet_Encoder = None
 
 # Load models once at startup
 def load_models():
-    global unet, tokenizer_one, tokenizer_two, noise_scheduler, text_encoder_one, text_encoder_two, image_encoder, vae, UNet_Encoder
-    base_path = "your_base_path_here"
-    unet = UNet2DConditionModel.from_pretrained(base_path, subfolder="unet", torch_dtype=torch.float16, force_download=False)
-    tokenizer_one = AutoTokenizer.from_pretrained(base_path, subfolder="tokenizer", use_fast=False, force_download=False)
-    tokenizer_two = AutoTokenizer.from_pretrained(base_path, subfolder="tokenizer_2", use_fast=False, force_download=False)
-    noise_scheduler = DDPMScheduler.from_pretrained(base_path, subfolder="scheduler")
-    text_encoder_one = CLIPTextModel.from_pretrained(base_path, subfolder="text_encoder", torch_dtype=torch.float16, force_download=False)
-    text_encoder_two = CLIPTextModelWithProjection.from_pretrained(base_path, subfolder="text_encoder_2", torch_dtype=torch.float16, force_download=False)
-    image_encoder = CLIPVisionModelWithProjection.from_pretrained(base_path, subfolder="image_encoder", torch_dtype=torch.float16, force_download=False)
-    vae = AutoencoderKL.from_pretrained(base_path, subfolder="vae", torch_dtype=torch.float16, force_download=False)
-    UNet_Encoder = UNet2DConditionModel_ref.from_pretrained(base_path, subfolder="unet_encoder", torch_dtype=torch.float16, force_download=False)
-
-# Call the function to load models at startup
-load_models()
-
-# Helper function to free up GPU memory after processing
-def clear_gpu_memory():
-    torch.cuda.empty_cache()
-    torch.cuda.synchronize()
-
-# Helper function to convert base64 to image
-def base64_to_image(base64_str):
-    image_data = base64.b64decode(base64_str)
+    global unet, tokenizer_one, tokenizer_two, noise_scheduler
+    global text_encoder_one, text_encoder_two, image_encoder, vae, UNet_Encoder
+
+    if unet is None:
+        # Load models only when required to reduce memory usage
+        unet = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-v1-4")
+    
+    if tokenizer_one is None:
+        tokenizer_one = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14")
+    
+    if tokenizer_two is None:
+        tokenizer_two = AutoTokenizer.from_pretrained("openai/clip-vit-large-patch14-336")
+    
+    if noise_scheduler is None:
+        noise_scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4")
+
+    if text_encoder_one is None:
+        text_encoder_one = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14")
+
+    if text_encoder_two is None:
+        text_encoder_two = CLIPTextModelWithProjection.from_pretrained("openai/clip-vit-large-patch14-336")
+    
+    if image_encoder is None:
+        image_encoder = CLIPVisionModelWithProjection.from_pretrained("openai/clip-vit-large-patch14")
+
+    if vae is None:
+        vae = AutoencoderKL.from_pretrained("stabilityai/stable-diffusion-v1-4")
+
+    if UNet_Encoder is None:
+        UNet_Encoder = UNet2DConditionModel.from_pretrained("stabilityai/stable-diffusion-v1-4")
+
+# Helper function to process base64 image
+def decode_image(image_base64):
+    image_data = base64.b64decode(image_base64)
     image = Image.open(BytesIO(image_data)).convert("RGB")
     return image
 
-# Helper function to resize images for faster processing
-def resize_image(image, size=(512, 768)):
-    return image.resize(size)
+# Helper function to encode image to base64
+def encode_image(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-# Example try-on function
-@app.route('/start_tryon', methods=['POST'])
-def start_tryon():
-    data = request.get_json()
-    garm_img_base64 = data['garm_img']
-    human_img_base64 = data['human_img']
-    
-    # Decode and resize images
-    garm_img = resize_image(base64_to_image(garm_img_base64))
-    human_img = resize_image(base64_to_image(human_img_base64))
+# Route for image processing
+@app.route('/process_image', methods=['POST'])
+def process_image():
+    data = request.json
 
-    # Convert images to tensors and move to GPU
-    garm_img_tensor = torch.tensor(garm_img, dtype=torch.float16).unsqueeze(0).to('cuda')
-    human_img_tensor = torch.tensor(human_img, dtype=torch.float16).unsqueeze(0).to('cuda')
-    
-    try:
-        # Processing steps (dummy example, replace with your logic)
-        with torch.inference_mode():
-            # Run the inference for both images
-            result_tensor = unet(garm_img_tensor, human_img_tensor)  # Replace with your actual logic
-            
-        # Free GPU memory after inference
-        clear_gpu_memory()
-
-        # Convert result back to base64 for return
-        result_img = Image.fromarray(result_tensor.squeeze(0).cpu().numpy())
-        buffered = BytesIO()
-        result_img.save(buffered, format="JPEG")
-        result_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-        return jsonify({"result": result_base64})
-    
-    except Exception as e:
-        clear_gpu_memory()
-        return jsonify({"error": str(e)}), 500
+    # Load the models (this will only happen once)
+    load_models()
+
+    # Extract the image from the request
+    image_base64 = data.get('image_base64')
+    if not image_base64:
+        return jsonify({"error": "No image provided"}), 400
+
+    image = decode_image(image_base64)
+
+    # Perform inference with the models (example, modify as needed)
+    processed_image = image  # Placeholder for actual image processing
+
+    # Return the processed image as base64
+    processed_image_base64 = encode_image(processed_image)
+    return jsonify({"processed_image": processed_image_base64})
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=7860)