Update app.py

app.py

@@ -1,209 +1,104 @@
-import re
+import gradio as gr
+from PIL import Image
+from io import BytesIO
+import base64
+import requests
 import os
-import yaml
-import tempfile
-import subprocess
-from pathlib import Path
-
+import random
 import torch
-import gradio as gr
-
-#from src.flux.xflux_pipeline import XFluxPipeline
-
-
-def list_dirs(path):
-    if path is None or path == "None" or path == "":
-        return
-
-    if not os.path.exists(path):
-        path = os.path.dirname(path)
-        if not os.path.exists(path):
-            return
-
-    if not os.path.isdir(path):
-        path = os.path.dirname(path)
-
-    def natural_sort_key(s, regex=re.compile("([0-9]+)")):
-        return [
-            int(text) if text.isdigit() else text.lower() for text in regex.split(s)
-        ]
-
-    subdirs = [
-        (item, os.path.join(path, item))
-        for item in os.listdir(path)
-        if os.path.isdir(os.path.join(path, item))
-    ]
-    subdirs = [
-        filename
-        for item, filename in subdirs
-        if item[0] != "." and item not in ["__pycache__"]
-    ]
-    subdirs = sorted(subdirs, key=natural_sort_key)
-    if os.path.dirname(path) != "":
-        dirs = [os.path.dirname(path), path] + subdirs
-    else:
-        dirs = [path] + subdirs
-
-    if os.sep == "\\":
-        dirs = [d.replace("\\", "/") for d in dirs]
-    for d in dirs:
-        yield d
-
-def list_train_data_dirs():
-    current_train_data_dir = "."
-    return list(list_dirs(current_train_data_dir))
-
-def update_config(d, u):
-    for k, v in u.items():
-        if isinstance(v, dict):
-            d[k] = update_config(d.get(k, {}), v)
-        else:
-            # convert Gradio components to strings
-            if hasattr(v, 'value'):
-                d[k] = str(v.value)
-            else:
-                try:
-                    d[k] = int(v)
-                except (TypeError, ValueError):
-                    d[k] = str(v)
-    return d
-
-def start_lora_training(
-        data_dir: str, output_dir: str, lr: float, steps: int, rank: int
-    ):
-    inputs = {
-        "data_config": {
-            "img_dir": data_dir,
-            },
-            "output_dir": output_dir,
-            "learning_rate": lr,
-            "rank": rank,
-            "max_train_steps": steps,
-    }
-
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        print(f"Creating folder {output_dir} for the output checkpoint file...")
-
-    script_path = Path(__file__).resolve()
-    config_path = script_path.parent / "train_configs" / "test_lora.yaml"
-    with open(config_path, 'r') as file:
-        config = yaml.safe_load(file)
-
-    config = update_config(config, inputs)
-    print("Config file is updated...", config)
-    with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".yaml") as temp_file:
-        yaml.dump(config, temp_file, default_flow_style=False)
-        tmp_config_path = temp_file.name
-
-    command = ["accelerate", "launch", "train_flux_lora_deepspeed.py", "--config", tmp_config_path]
-    result = subprocess.run(command, check=True)
-
-    # rRemove the temporary file after the command is run
-    Path(tmp_config_path).unlink()
-
-    return result
-
-
-def create_demo(
-        model_type: str,
-        device: str = "cuda" if torch.cuda.is_available() else "cpu",
-        offload: bool = False,
-        ckpt_dir: str = "",
-    ):
-    xflux_pipeline = XFluxPipeline(model_type, device, offload)
-    checkpoints = sorted(Path(ckpt_dir).glob("*.safetensors"))
-
-    with gr.Blocks() as demo:
-        gr.Markdown(f"# Flux Adapters by XLabs AI - Model: {model_type}")
-        with gr.Tab("Inference"):
-            with gr.Row():
-                with gr.Column():
-                    prompt = gr.Textbox(label="Prompt", value="handsome woman in the city")
-
-                    with gr.Accordion("Generation Options", open=False):
-                        with gr.Row():
-                            width = gr.Slider(512, 2048, 1024, step=16, label="Width")
-                            height = gr.Slider(512, 2048, 1024, step=16, label="Height")
-                        neg_prompt = gr.Textbox(label="Negative Prompt", value="bad photo")
-                        with gr.Row():
-                            num_steps = gr.Slider(1, 50, 25, step=1, label="Number of steps")
-                            timestep_to_start_cfg = gr.Slider(1, 50, 1, step=1, label="timestep_to_start_cfg")
-                        with gr.Row():
-                            guidance = gr.Slider(1.0, 5.0, 4.0, step=0.1, label="Guidance", interactive=True)
-                            true_gs = gr.Slider(1.0, 5.0, 3.5, step=0.1, label="True Guidance", interactive=True)
-                        seed = gr.Textbox(-1, label="Seed (-1 for random)")
-
-                    with gr.Accordion("ControlNet Options", open=False):
-                        control_type = gr.Dropdown(["canny", "hed", "depth"], label="Control type")
-                        control_weight = gr.Slider(0.0, 1.0, 0.8, step=0.1, label="Controlnet weight", interactive=True)
-                        local_path = gr.Dropdown(checkpoints, label="Controlnet Checkpoint",
-                            info="Local Path to Controlnet weights (if no, it will be downloaded from HF)"
-                            )
-                        controlnet_image = gr.Image(label="Input Controlnet Image", visible=True, interactive=True)
-
-                    with gr.Accordion("LoRA Options", open=False):
-                        lora_weight = gr.Slider(0.0, 1.0, 0.9, step=0.1, label="LoRA weight", interactive=True)
-                        lora_local_path = gr.Dropdown(
-                            checkpoints, label="LoRA Checkpoint", info="Local Path to Lora weights"
-                            )
-
-                    with gr.Accordion("IP Adapter Options", open=False):
-                        image_prompt = gr.Image(label="image_prompt", visible=True, interactive=True)
-                        ip_scale = gr.Slider(0.0, 1.0, 1.0, step=0.1, label="ip_scale")
-                        neg_image_prompt = gr.Image(label="neg_image_prompt", visible=True, interactive=True)
-                        neg_ip_scale = gr.Slider(0.0, 1.0, 1.0, step=0.1, label="neg_ip_scale")
-                        ip_local_path = gr.Dropdown(
-                            checkpoints, label="IP Adapter Checkpoint",
-                            info="Local Path to IP Adapter weights (if no, it will be downloaded from HF)"
-                            )
-                    generate_btn = gr.Button("Generate")
-
-                with gr.Column():
-                    output_image = gr.Image(label="Generated Image")
-                    download_btn = gr.File(label="Download full-resolution")
-
-            inputs = [prompt, image_prompt, controlnet_image, width, height, guidance,
-                    num_steps, seed, true_gs, ip_scale, neg_ip_scale, neg_prompt,
-                    neg_image_prompt, timestep_to_start_cfg, control_type, control_weight,
-                    lora_weight, local_path, lora_local_path, ip_local_path
-                    ]
-            generate_btn.click(
-                fn=xflux_pipeline.gradio_generate,
-                inputs=inputs,
-                outputs=[output_image, download_btn],
-            )
-
-        with gr.Tab("LoRA Finetuning"):
-            data_dir =  gr.Dropdown(list_train_data_dirs(),
-                                    label="Training images (directory containing the training images)"
-                                    )
-            output_dir = gr.Textbox(label="Output Path", value="lora_checkpoint")
-
-            with gr.Accordion("Training Options", open=True):
-                lr = gr.Textbox(label="Learning Rate", value="1e-5")
-                steps = gr.Slider(10000, 20000, 20000, step=100, label="Train Steps")
-                rank = gr.Slider(1, 100, 16, step=1, label="LoRa Rank")
-
-            training_btn = gr.Button("Start training")
-            training_btn.click(
-                fn=start_lora_training,
-                inputs=[data_dir, output_dir, lr, steps, rank],
-                outputs=[],
-            )
-
-
-    return demo
-
+import subprocess
+import numpy as np
+import cv2
+from transformers import AutoProcessor, AutoModelForCausalLM
+from diffusers import DiffusionPipeline
+from datetime import datetime
+from mistralai import Mistral
+from theme import theme
+from fastapi import FastAPI
+
+app = FastAPI()
+
+
+
+api_key = os.getenv("MISTRAL_API_KEY")
+Mistralclient = Mistral(api_key=api_key)
+
+def flip_image(x):
+    return np.fliplr(x)
+
+def encode_image(image_path):
+    """Encode the image to base64."""
+    try:
+        # Open the image file
+        image = Image.open(image_path).convert("RGB")
+
+        # Resize the image to a height of 512 while maintaining the aspect ratio
+        base_height = 512
+        h_percent = (base_height / float(image.size[1]))
+        w_size = int((float(image.size[0]) * float(h_percent)))
+        image = image.resize((w_size, base_height), Image.LANCZOS)
+
+        # Convert the image to a byte stream
+        buffered = BytesIO()
+        image.save(buffered, format="JPEG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        return img_str
+    except FileNotFoundError:
+        print(f"Error: The file {image_path} was not found.")
+        return None
+    except Exception as e:  # Add generic exception handling
+        print(f"Error: {e}")
+        return None
+
+def feifeichat(image):
+    try:
+        model = "pixtral-large-2411"
+        # Define the messages for the chat
+        base64_image = encode_image(image)
+        messages = [{
+            "role":
+            "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": "Please provide a detailed description of this photo"
+                },
+                {
+                    "type": "image_url",
+                    "image_url": f"data:image/jpeg;base64,{base64_image}" 
+                },
+            ],
+            "stream": False,
+        }]
+    
+        partial_message = ""
+        for chunk in Mistralclient.chat.stream(model=model, messages=messages):
+            if chunk.data.choices[0].delta.content is not None:
+                partial_message = partial_message + chunk.data.choices[
+                    0].delta.content
+                yield partial_message
+    except Exception as e:  # Add common exception handling
+        print(f"Error: {e}")
+        return "Please upload a photo"
+
+
+with gr.Blocks(theme=theme, elem_id="app-container") as app:
+    gr.Markdown("Image To Flux Prompt")
+    with gr.Tab(label="Image To Prompt"):
+        with gr.Row():
+            with gr.Column():
+                input_img = gr.Image(label="Input Picture",height=320,type="filepath")
+                submit_btn = gr.Button(value="Submit", variant='primary')
+            with gr.Column():
+                output_text = gr.Textbox(label="Flux Prompt", show_copy_button = True)
+                clr_button =gr.Button("Clear",variant="primary", elem_id="clear_button")
+                clr_button.click(lambda: gr.Textbox(value=""), None, output_text)
+
+        submit_btn.click(feifeichat, [input_img], [output_text])
+        
+    
+        
 if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser(description="Flux")
-    parser.add_argument("--name", type=str, default="flux-dev", help="Model name")
-    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Device to use")
-    parser.add_argument("--offload", action="store_true", help="Offload model to CPU when not in use")
-    parser.add_argument("--share", action="store_true", help="Create a public link to your demo")
-    parser.add_argument("--ckpt_dir", type=str, default=".", help="Folder with checkpoints in safetensors format")
-    args = parser.parse_args()
+    app.launch()
 
-    demo = create_demo(args.name, args.device, args.offload, args.ckpt_dir)
-    demo.launch(share=args.share)