app.py

import gradio as gr
from gradio_client import Client
from datasets import load_dataset
from huggingface_hub import InferenceClient
import random
from PIL import Image
import os
from random import randint
from datetime import datetime
import pathlib
import numpy as np
from prompts import __all__ as system_prompts
from all_models import __all__ as all_models
# Initialize API clients and other setups
token = os.getenv('HF_TOKEN')
sdxl_client = Client("K00B404/Stable-Diffusion-XL-Serverless", hf_token=token)
seed = randint(0, 12345678)
hf_client = InferenceClient(
    #"georgesung/llama2_7b_chat_uncensored",
    "Xennon-BD/Qwen-uncensored-v2",
    #"meta-llama/Llama-3.2-1B-Instruct", 
    token=token
)
hf_img_client = InferenceClient(all_models[0], token=token)

qwen_client = Client("K00B404/chatQwenne", hf_token=token)
personas_dataset = load_dataset("MohamedRashad/FinePersonas-Lite", split="train")
personas_list = personas_dataset["persona"][:100]
OUTPUT_DIR = "generated_characters"
pathlib.Path(OUTPUT_DIR).mkdir(exist_ok=True)

def get_timestamp_filename():
    return datetime.now().strftime("%Y%m%d_%H%M%S")

def save_outputs(image, description, base_filename):
    try:
        image_path = os.path.join(OUTPUT_DIR, f"{base_filename}.png")
        text_path = os.path.join(OUTPUT_DIR, f"{base_filename}.txt")
        image.save(image_path, format='PNG')
        with open(text_path, 'w', encoding='utf-8') as f:
            f.write(description)
        return image_path, text_path
    except Exception:
        return None, None



# Define the respond function using the HuggingFace InferenceClient
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    stream=False,
):
    messages = [{"role": "system", "content": system_message}]

    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})

    messages.append({"role": "user", "content": message+"|"})

    response = ""
    if stream:
        for message in hf_client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=stream,
            temperature=temperature,
            top_p=top_p,
        ):
            token = message.choices[0].delta.content
            response += token
            yield response
    else:
        response_data = hf_client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=stream,
            temperature=temperature,
            top_p=top_p,
        )
        yield response_data.choices[0].message.content


# Character Creator Functions
def generate_character_description(character_prompt, max_length):
    pre_force="a full body image , from  head to toe image of "
    pro_force=" 8k photorealistic "
    
    try:
        response_generator = respond(
            message=pre_force+character_prompt+pro_force,
            history=[],
            system_message=system_prompts[0],
            max_tokens=max_length // 4,
            temperature=0.9,
            top_p=0.95,
            stream=False,
        )
        result = "".join(response_generator)
        return result[:max_length]
    except Exception as e:
        return f"Error: {str(e)}"
    
    
    '''try:
        result = qwen_client.predict(
            message=character_prompt,
            system_message=system_prompts[0],
            max_tokens=max_length // 4,
            temperature=0.9,
            top_p=0.95,
            api_name="/chat"
        )
        return result[:max_length]
    except Exception as e:
        return f"Error: {str(e)}"'''


def generate_character(
            #model="stabilityai/stable-diffusion-2-1", 
            character_concept, 
            generate_description, 
            negative_prompt, 
            num_inference_steps, 
            guidance_scale, 
            max_description_length, 
            lora_repo_id
        ):
    
    '''
    prompt (str) —                         The prompt to generate an image from.
    negative_prompt (List[str, optional) — One or several prompt to guide what NOT to include in image generation.
    height (float, optional) —             The height in pixels of the image to generate.
    width (float, optional) —              The width in pixels of the image to generate.
    num_inference_steps (int, optional) —  The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
    guidance_scale (float, optional) —     A higher guidance scale value encourages the model to generate images closely linked to the text prompt, but values too high may cause saturation and other artifacts.
    model (str, optional) —                The model to use for inference. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed Inference Endpoint. If not provided, the default recommended text-to-image model will be used. Defaults to None.
    scheduler (str, optional) —            Override the scheduler with a compatible one.
    target_size (TextToImageTargetSize, optional) — The size in pixel of the output image
    seed (int, optional) —                 Seed for the random number generator. 
    
    Returns
    Image
    '''
    description = generate_character_description(character_concept, max_description_length) if generate_description else character_concept
    prompt = f"{character_concept}\n{description}"
    if lora_repo_id:
        prompt += f" <lora:{lora_repo_id}:1.0>"
    try:
        #seed = randint(0, 12345678)
        # output is a PIL.Image object
        #image = hf_img_client.text_to_image(prompt=prompt, negative_prompt=negative_prompt,num_inference_steps=num_inference_steps,guidance_scale=guidance_scale,seed=seed)
        # Save the image to a file
        ''' base_filename = get_timestamp_filename()
        image_path = os.path.join(OUTPUT_DIR, f"{base_filename}.png")
        image.save(image_path, format='PNG')
    
        # Save the description
        text_path = os.path.join(OUTPUT_DIR, f"{base_filename}.txt")
        with open(text_path, 'w', encoding='utf-8') as f:
            f.write(description)
    
        save_message = f"Files saved as:\nImage: {image_path}\nDescription: {text_path}"
        return image, description, save_message
        except Exception as e:
            return None, f"Error: {str(e)}", ""
            '''
            
        image = sdxl_client.predict(
            prompt=prompt, is_negative=negative_prompt, steps=num_inference_steps, cfg_scale=guidance_scale, sampler="DIMM", seed=-1, strength=0.9, api_name="/query"
        )
        base_filename = get_timestamp_filename()
        image_path, text_path = save_outputs(image, description, base_filename)
        save_message = f"Files saved as:\nImage: {image_path}\nDescription: {text_path}" if image_path and text_path else "Files not saved locally"
        return image, description, save_message
    except Exception as e:
        return None, f"Error: {str(e)}", ""
        
        


# Boxed Inpainting Functions
def query_gligen(image, prompt, box_x1, box_y1, box_x2, box_y2, object_prompt):
    # Simulate the GLIGEN API call
    return image  # Placeholder for actual implementation

def process_drawn_boxes(image_data, boxes, object_prompt, main_prompt):
    try:
        image = image_data["composite"]  # Use the final edited image
        if not boxes or len(boxes) == 0:
            return image, "No boxes drawn. Please draw a box."
        box = boxes[-1]
        height, width = image.size
        normalized_box = [box[0] / width, box[1] / height, box[2] / width, box[3] / height]
        result_image = query_gligen(
            image=image,
            prompt=main_prompt,
            box_x1=normalized_box[0],
            box_y1=normalized_box[1],
            box_x2=normalized_box[2],
            box_y2=normalized_box[3],
            object_prompt=object_prompt,
        )
        base_filename = get_timestamp_filename()
        image_path, _ = save_outputs(result_image, f"GLIGEN: {main_prompt} - Object: {object_prompt}", base_filename)
        save_message = f"Image saved as: {image_path}" if image_path else "File not saved locally"
        return result_image, save_message
    except Exception as e:
        return image_data["composite"], f"Error: {str(e)}"

def capture_box_drawing(image, sketch_data, object_prompt, main_prompt):
    """
    Handles box drawing based on the freehand sketch.
    """
    try:
        # Check if the user has drawn anything
        if sketch_data is None:
            return image, "No box drawn. Please draw a rectangular box."

        # Convert the sketch into bounding box coordinates
        # Extract non-transparent pixels from the sketch
        sketch_np = np.array(sketch_data)
        rows, cols = np.where(sketch_np[..., -1] > 0)  # Alpha > 0 indicates a drawing

        if rows.size == 0 or cols.size == 0:
            return image, "No box detected. Ensure you draw a closed box."

        # Calculate bounding box
        x1, y1, x2, y2 = cols.min(), rows.min(), cols.max(), rows.max()

        # Process the image with the drawn bounding box
        result_image = query_gligen(
            image=image,
            prompt=main_prompt,
            box_x1=x1 / image.width,
            box_y1=y1 / image.height,
            box_x2=x2 / image.width,
            box_y2=y2 / image.height,
            object_prompt=object_prompt,
        )

        # Save and return
        base_filename = get_timestamp_filename()
        image_path, _ = save_outputs(result_image, f"Boxed Inpainting: {main_prompt}", base_filename)
        save_message = f"Image saved as: {image_path}" if image_path else "File not saved locally"
        return result_image, save_message
    except Exception as e:
        return image, f"Error: {str(e)}"

# Interface
with gr.Blocks() as demo:
    with gr.Tabs():
        with gr.Tab("Character Creator"):
            gr.Markdown("### Generate Characters")
            with gr.Row():
                char_concept = gr.Textbox(label="Character Concept")
                generate_desc = gr.Checkbox(label="Generate Description", value=True)
                negative_prompt = gr.Textbox(label="Negative Prompt", value="deformed, distorted, disfigured, poorly drawn, bad anatomy, wrong anatomy,extra limb, missing limb, floating limbs, mutated hands and fingers, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation,clothing, cloths ,pants, under pants, top, covered skin, underwear,gear, cloth, belt, cartoon, missing vagina,small vagina, small breasts, small nipples")
            with gr.Row():
                steps = gr.Slider(1, 50, value=20, step=1,label="Steps")
                guidance = gr.Slider(1, 20, value=7, step=1,label="Guidance")
                max_desc_len = gr.Number(value=1024, label="Max Description Length")
                lora_repo = gr.Textbox(label="Lora Repo ID")
            char_image = gr.Image(label="Generated Character")
            char_desc = gr.Textbox(label="Generated Description")
            save_info = gr.Textbox(label="Save Info")
            gen_btn = gr.Button("Generate")
            
            gen_btn.click(
                generate_character,
                inputs=[char_concept, generate_desc, negative_prompt, steps, guidance, max_desc_len, lora_repo],
                outputs=[char_image, char_desc, save_info],
            )
   
        with gr.Tab("Boxed Inpainting"):
            gr.Markdown("### Boxed Inpainting Tool")
            image_input = gr.ImageEditor(label="Edit and Draw", type="numpy", crop_size=None)
            boxes_input = gr.HighlightedText(label="Draw Boxes")
            object_prompt = gr.Textbox(label="Object Prompt")
            main_prompt = gr.Textbox(label="Main Prompt")
            inp_image = gr.Image(label="Processed Image")
            inp_save_info = gr.Textbox(label="Save Info")
            inp_btn = gr.Button("Process")
            inp_btn.click(
                process_drawn_boxes,
                inputs=[image_input, boxes_input, object_prompt, main_prompt],
                outputs=[inp_image, inp_save_info],
            )

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