Update app.py

app.py

@@ -1,97 +1,299 @@
-"""
-Copyright 2022 Balacoon
+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)
 
-TTS interactive demo
-"""
+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)
 
-import logging
-from typing import cast
+def get_timestamp_filename():
+    return datetime.now().strftime("%Y%m%d_%H%M%S")
 
-import gradio as gr
-from balacoon_tts import TTS
-from huggingface_hub import hf_hub_download, list_repo_files
+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
 
-# global tts module, initialized from a model selected
-tts = None
 
-gr.api_name="textTOPepe"
 
-def main():
-    logging.basicConfig(level=logging.INFO)
+# 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}]
 
-    with gr.Blocks() as demo:
-        gr.Markdown(
-            """
-            <h1 align="center">Text-to-Speech</h1>
+    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})
 
-            1. Write an utterance to generate,
-            2. Select the model to synthesize with
-            3. Select the speaker
-            4. Hit "Generate" and listen to the result!
+    messages.append({"role": "user", "content": message+"|"})
 
-            When you select a Model for the first time,
-            it will take a little time to download it.
-            """
+    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,
         )
-        gr.api_name="textTOPepe"
-        with gr.Row(variant="panel"):
-            text = gr.Textbox(label="Text", placeholder="Insert your article here...")
-
-        with gr.Row():
-            with gr.Column(variant="panel"):
-                repo_files = list_repo_files(repo_id="balacoon/tts")
-                model_files = [x for x in repo_files if x.endswith("_cpu.addon")]
-                model_name = gr.Dropdown(
-                    label="Model",
-                    choices=model_files,
-                )
-            with gr.Column(variant="panel"):
-                speaker = gr.Dropdown(label="Speaker", choices=[])
-
-            def set_model(model_name_str: str):
-                """
-                gets value from `model_name`, loads model,
-                re-initializes tts object, gets list of
-                speakers that model supports and set them to `speaker`
-                """
-                model_path = hf_hub_download(
-                    repo_id="balacoon/tts", filename=model_name_str
-                )
-                global tts
-                tts = TTS(model_path)
-                speakers = tts.get_speakers()
-                value = speakers[-1]
-                return gr.Dropdown.update(
-                    choices=speakers, value=value, visible=True
-                )
-
-            model_name.change(set_model, inputs=model_name, outputs=speaker)
-
-        with gr.Row(variant="panel"):
-            generate = gr.Button("Generate")
-        with gr.Row(variant="panel"):
-            audio = gr.Audio()
-
-        def synthesize_audio(text_str: str, speaker_str: str = ""):
-            """
-            gets utterance to synthesize from `text` Textbox
-            and speaker name from `speaker` dropdown list.
-            speaker name might be empty for single-speaker models.
-            Synthesizes the waveform and updates `audio` with it.
-            """
-            if not text_str:
-                logging.info("text or speaker are not provided")
-                return None
-            global tts
-            if len(text_str) > 10024:
-                text_str = text_str[:10024]
-            samples = cast(TTS, tts).synthesize(text_str, speaker_str)
-            return gr.Audio.update(value=(cast(TTS, tts).get_sampling_rate(), samples))
-
-        generate.click(synthesize_audio, inputs=[text, speaker], outputs=audio)
-    demo.api_name="textTOPepe"
-    demo.launch()
+        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__":
-    main()
+    demo.launch()
+