app.py

import base64
from openai import OpenAI
from dotenv import load_dotenv
import fal_client
import os
from datetime import datetime
import requests
import gradio as gr

load_dotenv()

client = OpenAI()

# Function to encode the image
def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")

def generate_description(image):
    """Generate description for uploaded image using OpenAI vision"""
    if image is None:
        return "Please upload an image first."
    
    try:
        # Convert PIL image to base64
        import io
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
        
        response = client.responses.create(
            model="gpt-4.1",
            input=[
                {
                    "role": "user",
                    "content": [
                        { "type": "input_text", "text": "Describe the image in rich, natural detail for txt2img model. Focus on the scene, the person, their casual clothes, accessories, makeup, the everyday environment, and what they're doing. Capture the emotions, expressions, mood, and candid vibe as if it were a real Instagram Reel, TikTok, or Pinterest-style snapshot. Mention the camera angle, framing, depth of field, shadows, colors, and textures in a way that feels spontaneous and unpolished, with natural light and subtle motion blur. Convey the overall feeling, atmosphere, and sense of being in that moment. Describe the angle of the camera and pose of the model precisely. The result will be used to generate a similar image with a txt2img model, emphasizing realism, casualness, and lived-in authenticity rather than studio-perfect or commercialized style. Highlight the photo taken with iPhone front camera in-the-moment." },
                        {
                            "type": "input_image",
                            "image_url": f"data:image/jpeg;base64,{base64_image}",
                        },
                    ],
                }
            ],
        )
        
        description = response.output_text + "\n Photo taken with iPhone front camera in-the-moment."
        return description
        
    except Exception as e:
        return f"Error generating description: {str(e)}"


model = "fal-ai/gemini-25-flash-image" # , "fal-ai/nano-banana",  

def on_queue_update(update):
    if isinstance(update, fal_client.InProgress):
        for log in update.logs:
           print(log["message"])

def generate_avatar(prompt, selected_model):
    """Generate avatar image with given prompt and selected model"""
    if not prompt or prompt.strip() == "":
        return None, "Please provide a prompt for avatar generation."
    
    try:
        result = fal_client.subscribe(
            selected_model,
            arguments={
                "prompt": prompt,
                "num_images": 1,
                "output_format": "png"
            },
            with_logs=True,
            on_queue_update=on_queue_update,
        )

        # Return the first generated image
        if 'images' in result and len(result['images']) > 0:
            image_url = result['images'][0]['url']
            
            # Download the image
            response = requests.get(image_url)
            if response.status_code == 200:
                # Convert to PIL Image for Gradio display
                from PIL import Image
                import io
                image = Image.open(io.BytesIO(response.content))
                return image, "Avatar generated successfully!"
            else:
                return None, f"Failed to download image: {response.status_code}"
        else:
            return None, "No images found in result"
            
    except Exception as e:
        return None, f"Error generating avatar: {str(e)}"

# Gradio Interface
def process_image_and_generate_avatar(image):
    """Process uploaded image and generate avatar"""
    if image is None:
        return None, "Please upload an image first.", "Upload an image to get started!"
    
    # Generate description
    description = generate_description(image)
    
    # Generate avatar
    avatar_image, message = generate_avatar(description)
    
    return avatar_image, description, message

# Available models
available_models = [
    "fal-ai/gemini-25-flash-image",
    "fal-ai/nano-banana", 
    "fal-ai/flux-pro",
    "fal-ai/stable-diffusion-v35-large",
    "fal-ai/qwen-image",
    "fal-ai/imagen4-preview",
    "fal-ai/imagen4-preview-fast",
    "fal-ai/imagen4-preview-ultra"
]

# Create Gradio interface
with gr.Blocks(title="Avatar Clone Generator") as demo:
    gr.Markdown("# Avatar Clone Generator")
    gr.Markdown("Upload a photo to generate a description and create an avatar!")
    
    with gr.Row():
        with gr.Column():
            # Image upload
            input_image = gr.Image(
                label="Upload Photo",
                type="pil",
                height=300
            )
            
            # Model selection
            model_dropdown = gr.Dropdown(
                choices=available_models,
                value="fal-ai/gemini-25-flash-image",
                label="Select Vision Model",
                info="Choose the AI model for generating your avatar"
            )
            
            # Generate description button
            generate_desc_btn = gr.Button("Generate Description", variant="primary")
            
            # Description text area (editable)
            description_text = gr.Textbox(
                label="Generated Description (You can edit this)",
                lines=8,
                placeholder="Description will appear here after uploading an image..."
            )
            
            # Generate avatar button
            generate_avatar_btn = gr.Button("Generate Avatar", variant="secondary")
            
            # Status message
            status_text = gr.Textbox(
                label="Status",
                lines=2,
                interactive=False
            )
        
        with gr.Column():
            # Generated avatar display
            output_image = gr.Image(
                label="Generated Avatar",
                height=400
            )
    
    # Event handlers
    def on_image_upload(image):
        if image is not None:
            description = generate_description(image)
            return description, "Image uploaded! Click 'Generate Description' to proceed."
        return "", "Please upload an image."
    
    def on_generate_description(image):
        if image is None:
            return "", "Please upload an image first."
        description = generate_description(image)
        return description, "Description generated! You can edit it if needed."
    
    def on_generate_avatar(description, selected_model):
        if not description or description.strip() == "":
            return None, "Please generate a description first."
        avatar_image, message = generate_avatar(description, selected_model)
        return avatar_image, message
    
    # Connect events
    input_image.change(
        fn=on_image_upload,
        inputs=[input_image],
        outputs=[description_text, status_text]
    )
    
    generate_desc_btn.click(
        fn=on_generate_description,
        inputs=[input_image],
        outputs=[description_text, status_text]
    )
    
    generate_avatar_btn.click(
        fn=on_generate_avatar,
        inputs=[description_text, model_dropdown],
        outputs=[output_image, status_text]
    )

if __name__ == "__main__":
    demo.launch(share=True)