Upload api.py

api.py

@@ -0,0 +1,720 @@
+import os
+import torch
+import boto3
+import random
+import string
+import numpy as np
+import logging
+import datetime
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, constr, conint
+from diffusers import (FluxPipeline, FluxControlNetPipeline, 
+                        FluxControlNetModel, FluxImg2ImgPipeline, 
+                        FluxInpaintPipeline, CogVideoXImageToVideoPipeline)
+from diffusers.utils import load_image
+from PIL import Image
+from collections import defaultdict
+import time
+
+# Setup logging
+logging.basicConfig(level=logging.INFO, 
+                    format='%(asctime)s - %(levelname)s - %(message)s',
+                    handlers=[
+                        logging.FileHandler("error.txt"),
+                        logging.StreamHandler()
+                    ])
+
+app = FastAPI()
+
+# Allow CORS for specific origins if needed
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Update with specific domains as necessary
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+MAX_SEED = np.iinfo(np.int32).max
+
+# AWS S3 Configuration
+AWS_ACCESS_KEY_ID = "your-access-key-id"
+AWS_SECRET_ACCESS_KEY = "your-secret-access-key"
+AWS_REGION = "your-region"
+S3_BUCKET_NAME = "your-bucket-name"
+
+# Initialize S3 client
+s3_client = boto3.client(
+    's3',
+    aws_access_key_id=AWS_ACCESS_KEY_ID,
+    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
+    region_name=AWS_REGION
+)
+
+# Asynchronously log requests
+async def log_requests(user_key: str, prompt: str):
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+    log_entry = f"{timestamp}, {user_key}, {prompt}\n"
+    async with aiofiles.open("key_requests.txt", "a") as log_file:
+        await log_file.write(log_entry)
+
+# Asynchronously upload image to S3
+async def upload_image_to_s3(image_path: str, s3_path: str):
+    try:
+        s3_client.upload_file(image_path, S3_BUCKET_NAME, s3_path)
+        return f"https://{S3_BUCKET_NAME}.s3.{AWS_REGION}.amazonaws.com/{s3_path}"
+    except Exception as e:
+        logging.error(f"Error uploading image to S3: {e}")
+        raise HTTPException(status_code=500, detail=f"Image upload failed: {str(e)}")
+
+# Generate a random sequence of 12 numbers and 11 words
+def generate_random_sequence():
+    random_numbers = ''.join(random.choices(string.digits, k=12))  # 12 random digits
+    random_words = ''.join(random.choices(string.ascii_lowercase, k=11))  # 11 random letters
+    return f"{random_numbers}_{random_words}"
+
+# Load the default pipeline once globally for efficiency
+flux_pipe = FluxPipeline.from_pretrained("pranavajay/flow", torch_dtype=torch.bfloat16)
+flux_pipe.enable_model_cpu_offload()
+logging.info("FluxPipeline loaded successfully.")
+
+img_pipe = FluxImg2ImgPipeline.from_pretrained("pranavajay/flow", torch_dtype=torch.bfloat16)
+img_pipe.enable_model_cpu_offload()
+logging.info("FluxImg2ImgPipeline loaded successfully.")
+
+inpainting_pipe = FluxInpaintPipeline.from_pretrained("pranavajay/flow", torch_dtype=torch.bfloat16)
+inpainting_pipe.enable_model_cpu_offload()
+logging.info("FluxInpaintPipeline loaded successfully.")
+
+video = CogVideoXImageToVideoPipeline.from_pretrained(
+    "THUDM/CogVideoX-5b-I2V", torch_dtype=torch.bfloat16
+)
+video.enable_sequential_cpu_offload()
+video.vae.enable_tiling()
+video.vae.enable_slicing()
+logging.info("CogVideoXImageToVideoPipeline loaded successfully.")
+
+flux_controlnet_pipe = None
+
+# Rate limiting variables
+request_timestamps = defaultdict(list)  # Store timestamps of requests per user key
+RATE_LIMIT = 30  # Maximum requests allowed
+TIME_WINDOW = 5  # Time window in seconds
+
+# Available LoRA styles and ControlNet adapters
+style_lora_mapping = {
+    "Uncensored": {"path": "enhanceaiteam/Flux-uncensored", "triggered_word": "nsfw"},
+    "Logo": {"path": "Shakker-Labs/FLUX.1-dev-LoRA-Logo-Design", "triggered_word": "logo"},
+    "Yarn": {"path": "Shakker-Labs/FLUX.1-dev-LoRA-MiaoKa-Yarn-World", "triggered_word": "mkym this is made of wool"},
+    "Anime": {"path": "prithivMLmods/Canopus-LoRA-Flux-Anime", "triggered_word": "anime"},
+    "Comic": {"path": "wkplhc/comic", "triggered_word": "comic"}
+}
+
+adapter_controlnet_mapping = {
+    "Canny": "InstantX/FLUX.1-dev-controlnet-canny",
+    "Depth": "Shakker-Labs/FLUX.1-dev-ControlNet-Depth",
+    "Pose": "Shakker-Labs/FLUX.1-dev-ControlNet-Pose",
+    "Upscale": "jasperai/Flux.1-dev-Controlnet-Upscaler"
+}
+
+# Request model for query parameters
+class GenerateImageRequest(BaseModel):
+    prompt: constr(min_length=1)  # Ensures prompt is not empty
+    guidance_scale: float = 7.5
+    seed: conint(ge=0, le=MAX_SEED) = 42
+    randomize_seed: bool = False
+    height: conint(gt=0) = 768
+    width: conint(gt=0) = 1360
+    control_image_url: str = "https://enhanceai.s3.amazonaws.com/792e2322-77fe-4070-aac4-7fa8d9e29c11_1.png"
+    controlnet_conditioning_scale: float = 0.6
+    num_inference_steps: conint(gt=0) = 50
+    num_images_per_prompt: conint(gt=0, le=5) = 1  # Limit to max 5 images per request
+    style: str = None  # Optional LoRA style
+    adapter: str = None  # Optional ControlNet adapter
+    user_key: str  # API user key
+
+# Apply LoRA style to the prompt
+async def apply_lora_style(pipe, style, prompt):
+    if style in style_lora_mapping:
+        lora_path = style_lora_mapping[style]["path"]
+        triggered_word = style_lora_mapping[style]["triggered_word"]
+        pipe.load_lora_weights(lora_path)
+        return f"{triggered_word} {prompt}"
+    return prompt
+
+# Set ControlNet adapter for the pipeline
+async def set_controlnet_adapter(adapter: str, is_inpainting: bool = False):
+    global flux_controlnet_pipe
+    if adapter not in adapter_controlnet_mapping:
+        raise ValueError(f"Invalid ControlNet adapter: {adapter}")
+
+    controlnet_model_path = adapter_controlnet_mapping[adapter]
+    controlnet = FluxControlNetModel.from_pretrained(controlnet_model_path, torch_dtype=torch.bfloat16)
+    pipeline_cls = FluxControlNetPipeline if not is_inpainting else FluxInpaintPipeline
+    flux_controlnet_pipe = pipeline_cls.from_pretrained(
+        "pranavajay/flow", controlnet=controlnet, torch_dtype=torch.bfloat16
+    )
+    flux_controlnet_pipe.to("cuda")
+    logging.info(f"ControlNet adapter '{adapter}' loaded successfully.")
+
+# Rate limit user requests
+async def rate_limit(user_key: str):
+    current_time = time.time()
+    request_timestamps[user_key] = [t for t in request_timestamps[user_key] if current_time - t < TIME_WINDOW]
+    if len(request_timestamps[user_key]) >= RATE_LIMIT:
+        logging.info(f"Rate limit exceeded for user_key: {user_key}")
+        return False
+    request_timestamps[user_key].append(current_time)
+    return True
+
+@app.post("/text_to_image/")
+async def generate_image(req: GenerateImageRequest):
+    seed = req.seed or random.randint(0, MAX_SEED)
+
+    # Rate limit check
+    if not await rate_limit(req.user_key):
+        await log_requests(req.user_key, req.prompt)
+        
+
+    retries = 3  # Number of retries for transient errors
+
+    for attempt in range(retries):
+        try:
+            # Check if prompt is None or empty
+            if not req.prompt or req.prompt.strip() == "":
+                raise ValueError("Prompt cannot be empty.")
+
+            original_prompt = req.prompt  # Save the original prompt
+
+            # Set ControlNet if adapter is provided
+            if req.adapter:
+                try:
+                    await set_controlnet_adapter(req.adapter)
+                except Exception as e:
+                    logging.error(f"Error setting ControlNet adapter: {e}")
+                    raise HTTPException(status_code=400, detail=f"Failed to load ControlNet adapter: {str(e)}")
+
+                await apply_lora_style(flux_controlnet_pipe, req.style, req.prompt)
+
+                # Load control image asynchronously
+                try:
+                    loop = asyncio.get_running_loop()
+                    control_image = await loop.run_in_executor(None, load_image, req.control_image_url)
+                except Exception as e:
+                    logging.error(f"Error loading control image from URL: {e}")
+                    raise HTTPException(status_code=400, detail="Invalid control image URL or image could not be loaded.")
+
+                # Image generation with ControlNet
+                try:
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    images = await loop.run_in_executor(None, flux_controlnet_pipe, {
+                        "prompt": req.prompt,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "control_image": control_image,
+                        "generator": generator,
+                        "controlnet_conditioning_scale": req.controlnet_conditioning_scale
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images with ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation with ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+            else:
+                # Image generation without ControlNet
+                try:
+                    await apply_lora_style(flux_pipe, req.style, req.prompt)
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    images = await loop.run_in_executor(None, flux_pipe, {
+                        "prompt": req.prompt,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "generator": generator
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images without ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation without ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+
+            # Saving images and uploading to S3 asynchronously
+            image_urls = []
+            for img in images:
+                image_path = f"generated_images/{generate_random_sequence()}.png"
+                await loop.run_in_executor(None, img.save, image_path)
+                image_url = await upload_image_to_s3(image_path, image_path)
+                image_urls.append(image_url)
+                os.remove(image_path)  # Clean up local files after upload
+            
+            return {
+                "status": "success",
+                "output": image_urls,
+                "prompt": original_prompt,
+                "height": req.height,
+                "width": req.width,
+                "scale": req.guidance_scale,
+                "steps": req.num_inference_steps,
+                "style": req.style,
+                "adapter": req.adapter
+            }
+
+        except Exception as e:
+            logging.error(f"Attempt {attempt + 1} failed: {e}")
+            if attempt == retries - 1:  # Last attempt
+                raise HTTPException(status_code=500, detail=f"Failed to generate image after multiple attempts: {str(e)}")
+            continue  # Retry on transient errors
+
+class GenerateImageToImageRequest(BaseModel):
+    prompt: str = None  # Prompt can be None
+    image: str = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+    strength: float = 0.7
+    guidance_scale: float = 7.5
+    seed: conint(ge=0, le=MAX_SEED) = 42
+    randomize_seed: bool = False
+    height: conint(gt=0) = 768
+    width: conint(gt=0) = 1360
+    control_image_url: str = None  # Optional ControlNet image
+    controlnet_conditioning_scale: float = 0.6
+    num_inference_steps: conint(gt=0) = 50
+    num_images_per_prompt: conint(gt=0, le=5) = 1
+    style: str = None  # Optional LoRA style
+    adapter: str = None  # Optional ControlNet adapter
+    user_key: str  # API user key
+
+@app.post("/image_to_image/")
+async def generate_image_to_image(req: GenerateImageToImageRequest):
+    seed = req.seed
+    original_prompt = req.prompt
+    modified_prompt = original_prompt
+    
+    # Check if user is exceeding rate limit
+    if not await rate_limit(req.user_key):
+        await log_requests(req.user_key, req.prompt if req.prompt else "No prompt")
+        raise HTTPException(status_code=429, detail="Rate limit exceeded")
+   
+    retries = 3  # Number of retries for transient errors
+    loop = asyncio.get_running_loop()
+
+    for attempt in range(retries):
+        try:
+            # Check if prompt is None or empty
+            if not req.prompt or req.prompt.strip() == "":
+                raise ValueError("Prompt cannot be empty.")
+            
+            original_prompt = req.prompt  # Save the original prompt
+            
+            # Set ControlNet if adapter is provided
+            if req.adapter:
+                try:
+                    await set_controlnet_adapter(req.adapter)
+                except Exception as e:
+                    logging.error(f"Error setting ControlNet adapter: {e}")
+                    raise HTTPException(status_code=400, detail=f"Failed to load ControlNet adapter: {str(e)}")
+                
+                await apply_lora_style(flux_controlnet_pipe, req.style, req.prompt)
+
+                # Load control image asynchronously
+                try:
+                    control_image = await loop.run_in_executor(None, load_image, req.control_image_url)
+                except Exception as e:
+                    logging.error(f"Error loading control image from URL: {e}")
+                    raise HTTPException(status_code=400, detail="Invalid control image URL or image could not be loaded.")
+                
+                # Image generation with ControlNet
+                try:
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    images = await loop.run_in_executor(None, flux_controlnet_pipe, {
+                        "prompt": modified_prompt,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "control_image": control_image,
+                        "generator": generator,
+                        "controlnet_conditioning_scale": req.controlnet_conditioning_scale
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images with ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation with ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+            else:
+                # Image generation without ControlNet
+                try:
+                    await apply_lora_style(img_pipe, req.style, req.prompt)
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    source = await loop.run_in_executor(None, load_image, req.image)
+
+                    images = await loop.run_in_executor(None, img_pipe, {
+                        "prompt": modified_prompt,
+                        "image": source,
+                        "strength": req.strength,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "generator": generator
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images without ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation without ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+
+            # Saving images and uploading to S3 asynchronously
+            image_urls = []
+            for img in images:
+                image_path = f"generated_images/{generate_random_sequence()}.png"
+                await loop.run_in_executor(None, img.save, image_path)
+                image_url = await upload_image_to_s3(image_path, image_path)
+                image_urls.append(image_url)
+                os.remove(image_path)  # Clean up local files after upload
+            
+            return {
+                "status": "success",
+                "output": image_urls,
+                "prompt": original_prompt,
+                "height": req.height,
+                "width": req.width,
+                "image": req.image,
+                "strength": req.strength,
+                "scale": req.guidance_scale,
+                "steps": req.num_inference_steps,
+                "style": req.style,
+                "adapter": req.adapter
+            }
+
+        except Exception as e:
+            logging.error(f"Attempt {attempt + 1} failed: {e}")
+            if attempt == retries - 1:  # Last attempt
+                raise HTTPException(status_code=500, detail=f"Failed to generate image after multiple attempts: {str(e)}")
+            continue  # Retry on transient errors
+  
+
+   
+class GenerateInpaintingRequest(BaseModel):
+    prompt: str = None  # Prompt can be None
+    image: str = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
+    mask_image: str = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"
+    guidance_scale: float = 7.5
+    seed: conint(ge=0, le=MAX_SEED) = 42
+    randomize_seed: bool = False
+    height: conint(gt=0) = 768
+    width: conint(gt=0) = 1360
+    control_image_url: str = None  # Optional ControlNet image
+    controlnet_conditioning_scale: float = 0.6
+    num_inference_steps: conint(gt=0) = 50
+    num_images_per_prompt: conint(gt=0, le=5) = 1
+    style: str = None  # Optional LoRA style
+    adapter: str = None  # Optional ControlNet adapter
+    user_key: str  # API user key
+    
+@app.post("/inpainting/")
+async def generate_inpainting(req: GenerateInpaintingRequest):
+    seed = req.seed
+    original_prompt = req.prompt
+    modified_prompt = original_prompt
+
+    # Check if user is exceeding rate limit
+    if not await rate_limit(req.user_key):
+        await log_requests(req.user_key, req.prompt if req.prompt else "No prompt")
+        raise HTTPException(status_code=429, detail="Rate limit exceeded")
+
+    retries = 3  # Number of retries for transient errors
+    loop = asyncio.get_running_loop()
+
+    for attempt in range(retries):
+        try:
+            # Check if prompt is None or empty
+            if not req.prompt or req.prompt.strip() == "":
+                raise ValueError("Prompt cannot be empty.")
+
+            # Set ControlNet if adapter is provided
+            if req.adapter:
+                try:
+                    await set_controlnet_adapter(req.adapter, is_inpainting=True)
+                except Exception as e:
+                    logging.error(f"Error setting ControlNet adapter: {e}")
+                    raise HTTPException(status_code=400, detail=f"Failed to load ControlNet adapter: {str(e)}")
+                
+                await apply_lora_style(flux_inpainting_controlnet_pipe, req.style, req.prompt)
+
+                # Load control image asynchronously
+                try:
+                    control_image = await loop.run_in_executor(None, load_image, req.control_image_url)
+                except Exception as e:
+                    logging.error(f"Error loading control image from URL: {e}")
+                    raise HTTPException(status_code=400, detail="Invalid control image URL or image could not be loaded.")
+
+                # Image generation with ControlNet
+                try:
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    source = await loop.run_in_executor(None, load_image, req.image)
+                    mask = await loop.run_in_executor(None, load_image, req.mask_image)
+
+                    images = await loop.run_in_executor(None, flux_controlnet_pipe, {
+                        "prompt": modified_prompt,
+                        "image": source,
+                        "mask_image": mask,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "control_image": control_image,
+                        "generator": generator,
+                        "controlnet_conditioning_scale": req.controlnet_conditioning_scale
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images with ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation with ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+            else:
+                # Image generation without ControlNet
+                try:
+                    await apply_lora_style(inpainting_pipe, req.style, req.prompt)
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    source = await loop.run_in_executor(None, load_image, req.image)
+                    mask = await loop.run_in_executor(None, load_image, req.mask_image)
+
+                    images = await loop.run_in_executor(None, inpainting_pipe, {
+                        "prompt": modified_prompt,
+                        "image": source,
+                        "mask_image": mask,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "generator": generator
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images without ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation without ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+
+            # Saving generated images
+            image_urls = []
+            for i, img in enumerate(images):
+                image_path = f"generated_images/inpainting_{generate_random_sequence()}.png"
+                img.save(image_path)
+
+                # Optionally, upload the image to S3
+                s3_path = f"inpainting/{original_prompt.replace(' ', '_')}_{generate_random_sequence()}_{i}.png"
+                s3_url = await upload_file_to_s3(image_path, s3_path)
+                image_urls.append(s3_url)
+
+                # Clean up temporary files
+                os.remove(image_path)
+
+            return {
+                "status": "success",
+                "output": image_urls,
+                "prompt": original_prompt,
+                "height": req.height,
+                "width": req.width,
+                "scale": req.guidance_scale,
+                "style": req.style,
+                "adapter": req.adapter
+            }
+
+        except Exception as e:
+            logging.error(f"Attempt {attempt + 1} failed: {e}")
+            if attempt == retries - 1:  # Last attempt
+                raise HTTPException(status_code=500, detail=f"Failed to generate inpainting after multiple attempts: {str(e)}")
+            continue  # Retry on transient errors
+            
+            
+class GenerateVideoRequest(BaseModel):
+    prompt: constr(min_length=1)  # Ensures prompt is not empty
+    guidance_scale: float = 7.5
+    seed: conint(ge=0, le=MAX_SEED) = 42
+    randomize_seed: bool = False
+    height: conint(gt=0) = 768
+    width: conint(gt=0) = 1360
+    control_image_url: str = "https://enhanceai.s3.amazonaws.com/792e2322-77fe-4070-aac4-7fa8d9e29c11_1.png"
+    controlnet_conditioning_scale: float = 0.6
+    num_inference_steps: conint(gt=0) = 50
+    num_images_per_prompt: conint(gt=0, le=5) = 1  # Limit to max 5 images per request
+    style: str = None  # Optional LoRA style
+    adapter: str = None  # Optional ControlNet adapter
+    user_key: str  # API user key
+    
+
+@app.post("/text_to_video/")
+async def generate_video(req: GenerateImageRequest):
+    seed = req.seed
+    if not rate_limit(req.user_key):
+        log_requests(req.user_key, req.prompt)  # Log the request when rate limit is exceeded
+
+    retries = 3  # Number of retries for transient errors
+    s3_urls = []  # List to store S3 URLs of generated videos
+    loop = asyncio.get_running_loop()  # Get the current event loop
+
+    for attempt in range(retries):
+        try:
+            # Check if prompt is None or empty
+            if not req.prompt or req.prompt.strip() == "":
+                raise ValueError("Prompt cannot be empty.")
+
+            original_prompt = req.prompt  # Save the original prompt
+            
+            # Set ControlNet if adapter is provided
+            if req.adapter:
+                try:
+                    await set_controlnet_adapter(req.adapter)
+                except Exception as e:
+                    logging.error(f"Error setting ControlNet adapter: {e}")
+                    raise HTTPException(status_code=400, detail=f"Failed to load ControlNet adapter: {str(e)}")
+
+                # Load control image asynchronously
+                try:
+                    control_image = await loop.run_in_executor(None, load_image, req.control_image_url)
+                except Exception as e:
+                    logging.error(f"Error loading control image from URL: {e}")
+                    raise HTTPException(status_code=400, detail="Invalid control image URL or image could not be loaded.")
+                
+                # Image generation with ControlNet
+                try:
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    images = await loop.run_in_executor(None, flux_controlnet_pipe, {
+                        "prompt": original_prompt,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "control_image": control_image,
+                        "generator": generator,
+                        "controlnet_conditioning_scale": req.controlnet_conditioning_scale
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images with ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation with ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+            else:
+                # Image generation without ControlNet
+                try:
+                    await apply_lora_style(flux_pipe, req.style, req.prompt)
+                    if req.randomize_seed:
+                        seed = random.randint(0, MAX_SEED)
+                    generator = torch.Generator().manual_seed(seed)
+
+                    images = await loop.run_in_executor(None, flux_pipe, {
+                        "prompt": original_prompt,
+                        "guidance_scale": req.guidance_scale,
+                        "height": req.height,
+                        "width": req.width,
+                        "num_inference_steps": req.num_inference_steps,
+                        "num_images_per_prompt": req.num_images_per_prompt,
+                        "generator": generator
+                    })
+                except torch.cuda.OutOfMemoryError:
+                    logging.error("GPU out of memory error while generating images without ControlNet.")
+                    raise HTTPException(status_code=500, detail="GPU overload occurred while generating images. Try reducing the resolution or number of steps.")
+                except Exception as e:
+                    logging.error(f"Error during image generation without ControlNet: {e}")
+                    raise HTTPException(status_code=500, detail=f"Error during image generation: {str(e)}")
+
+            # Saving images and uploading to S3
+            for i, img in enumerate(images):
+                image_path = f"generated_images/{generate_random_sequence()}.png"
+                
+                # Save image asynchronously
+                await loop.run_in_executor(None, img.save, image_path)
+
+                # Generate video from the image
+                if req.randomize_seed:
+                    seed = random.randint(0, MAX_SEED)
+                vido = await loop.run_in_executor(None, video, {
+                    "prompt": original_prompt,
+                    "image": image_path,
+                    "num_videos_per_prompt": 1,
+                    "num_inference_steps": req.num_inference_steps,
+                    "num_frames": req.num_frames,
+                    "guidance_scale": req.guidance_scale,
+                    "generator": torch.Generator(device="cuda").manual_seed(seed)
+                })
+                
+                # Export the video to a file asynchronously
+                video_path = f"generated_video_{i}_{generate_random_sequence()}.mp4"
+                await loop.run_in_executor(None, export_to_video, vido, video_path, 8)
+
+                # Upload the video to S3 asynchronously
+                s3_path = f"videos/{original_prompt.replace(' ', '_')}_{generate_random_sequence()}_{i}.mp4"
+                s3_url = await loop.run_in_executor(None, upload_file_to_s3, video_path, s3_path)
+                s3_urls.append(s3_url)
+
+                # Clean up temporary files
+                os.remove(image_path)
+                os.remove(video_path)
+
+            return {
+                "status": "success",
+                "output": s3_urls,
+                "prompt": original_prompt,
+                "height": req.height,
+                "width": req.width,
+                "num_frames": req.num_frames,
+                "scale": req.guidance_scale,
+                "style": req.style,
+                "adapter": req.adapter
+            }
+
+        except Exception as e:
+            logging.error(f"Attempt {attempt + 1} failed: {e}")
+            if attempt == retries - 1:  # Last attempt
+                raise HTTPException(status_code=500, detail=f"Failed to generate video after multiple attempts: {str(e)}")
+            continue  # Retry on transient errors
+
+@app.on_event("shutdown")
+def shutdown_event():
+    """ Perform any cleanup activities on shutdown. """
+    logging.info("Shutting down the application gracefully.")
+
+# Additional endpoints can be added as needed, such as image-to-image or inpainting.
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+