Update app.py

app.py

@@ -1,332 +1,210 @@
-"""
-app.py
-
-A unified Gradio chat application for Multimodal OCR Granite Vision.
-Commands (enter these as a prefix in the text input):
-  - @rag: For retrieval‐augmented generation (e.g. PDFs or text-based queries).
-  - @granite: For image understanding.
-  - @video-infer: For video understanding (video is downsampled into frames).
-
-The app uses gr.MultimodalTextbox to support text input together with file uploads.
-"""
-
-import os
-import time
-import uuid
-import random
 import spaces
-import logging
-from threading import Thread
-from pathlib import Path
-from datetime import datetime, timezone
-
+import gradio as gr
 import torch
-import numpy as np
-import cv2
 from PIL import Image
-import gradio as gr
-
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    AutoProcessor,
-    AutoModelForVision2Seq,
-)
-
-# ---------------------------
-# Utility functions and setup
-# ---------------------------
-
-def get_device():
-    if torch.backends.mps.is_available():
-        return "mps"  # mac GPU
-    elif torch.cuda.is_available():
-        return "cuda"
-    else:
-        return "cpu"
-
-device = get_device()
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-
-def downsample_video(video_path):
-    """
-    Downsamples the video into 10 evenly spaced frames.
-    Returns a list of (PIL Image, timestamp in seconds) tuples.
-    """
-    vidcap = cv2.VideoCapture(video_path)
-    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
-    fps = vidcap.get(cv2.CAP_PROP_FPS)
-    frames = []
-    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
-    for i in frame_indices:
-        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
-        success, image = vidcap.read()
-        if success:
-            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-            pil_image = Image.fromarray(image)
-            timestamp = round(i / fps, 2)
-            frames.append((pil_image, timestamp))
-    vidcap.release()
-    return frames
-
-# ---------------------------
-# HF Embedding and LLM classes
-# ---------------------------
-
-class HFEmbedding:
-    def __init__(self, model_id: str):
-        self.model_name = model_id
-        logging.info(f"Loading embeddings model from: {self.model_name}")
-        # Using langchain_huggingface for embeddings
-        from langchain_huggingface import HuggingFaceEmbeddings  # ensure installed
-        # For simplicity, force CPU (adjust if needed)
-        self.embeddings_service = HuggingFaceEmbeddings(
-            model_name=self.model_name,
-            model_kwargs={"device": "cpu"},
-        )
-
-    def embed_documents(self, texts: list[str]) -> list[list[float]]:
-        return self.embeddings_service.embed_documents(texts)
-
-    def embed_query(self, text: str) -> list[float]:
-        return self.embed_documents([text])[0]
-
-class HFLLM:
-    def __init__(self, model_name: str):
-        self.device = device
-        self.model_name = model_name
-        logging.info("Loading HF language model...")
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)
-
-    def generate(self, prompt: str) -> list:
-        # Tokenize prompt and generate text
-        model_inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
-        generated_ids = self.model.generate(**model_inputs, max_new_tokens=1024)
-        generated_texts = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=False)
-        # Extract answer assuming a marker in the generated text
-        response = [{"answer": generated_texts[0].split("<|end_of_role|>")[-1].split("<|end_of_text|>")[0]}]
-        return response
-
-# ---------------------------
-# LightRAG: Retrieval-Augmented Generation (Dummy)
-# ---------------------------
-
-class LightRAG:
-    def __init__(self, config: dict):
-        self.config = config
-        # Load generation and embedding models immediately (or lazy load as needed)
-        self.gen_model = HFLLM(config['generation_model_id'])
-        self._embedding_model = HFEmbedding(config['embedding_model_id'])
-
-    def search(self, query: str, top_n: int = 5) -> list:
-        # Dummy retrieval: In practice, integrate with a vector store
-        from langchain_core.documents import Document  # ensure langchain_core is installed
-        dummy_doc = Document(
-            page_content="Dummy context for query: " + query,
-            metadata={"type": "text"}
-        )
-        return [dummy_doc]
-
-    def generate(self, query, context=None):
-        if context is None:
-            context = []
-        # Build prompt by concatenating retrieved context with the query.
-        prompt = ""
-        for doc in context:
-            prompt += doc.page_content + "\n"
-        prompt += "\nQuestion: " + query + "\nAnswer:"
-        results = self.gen_model.generate(prompt)
-        answer = results[0]["answer"]
-        return answer, prompt
+from diffusers import DiffusionPipeline
+import random
+import uuid
+from typing import Tuple
+import numpy as np
 
-# Global configuration for LightRAG
-rag_config = {
-    "embedding_model_id": "ibm-granite/granite-embedding-125m-english",
-    "generation_model_id": "ibm-granite/granite-3.1-8b-instruct",
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+MAX_SEED = np.iinfo(np.int32).max
+
+if not torch.cuda.is_available():
+    DESCRIPTIONz += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
+
+base_model = "black-forest-labs/FLUX.1-dev"
+pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
+
+lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
+trigger_word = "Super Realism"  # Leave trigger_word blank if not used.
+
+pipe.load_lora_weights(lora_repo)
+pipe.to("cuda")
+
+style_list = [
+    {
+        "name": "3840 x 2160",
+        "prompt": "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+    },
+    {
+        "name": "2560 x 1440",
+        "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+    },
+    {
+        "name": "HD+",
+        "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+    },
+    {
+        "name": "Style Zero",
+        "prompt": "{prompt}",
+    },
+]
+
+styles = {k["name"]: k["prompt"] for k in style_list}
+
+DEFAULT_STYLE_NAME = "3840 x 2160"
+STYLE_NAMES = list(styles.keys())
+
+def apply_style(style_name: str, positive: str) -> str:
+    return styles.get(style_name, styles[DEFAULT_STYLE_NAME]).replace("{prompt}", positive)
+
+@spaces.GPU(duration=60, enable_queue=True)
+def generate(
+    prompt: str,
+    seed: int = 0,
+    width: int = 1024,
+    height: int = 1024,
+    guidance_scale: float = 3,
+    randomize_seed: bool = False,
+    style_name: str = DEFAULT_STYLE_NAME,
+    progress=gr.Progress(track_tqdm=True),
+):
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+
+    positive_prompt = apply_style(style_name, prompt)
+    
+    if trigger_word:
+        positive_prompt = f"{trigger_word} {positive_prompt}"
+
+    images = pipe(
+        prompt=positive_prompt,
+        width=width,
+        height=height,
+        guidance_scale=guidance_scale,
+        num_inference_steps=28,
+        num_images_per_prompt=1,
+        output_type="pil",
+    ).images
+    image_paths = [save_image(img) for img in images]
+    print(image_paths)
+    return image_paths, seed
+
+examples = [
+    "Super Realism, High-resolution photograph, woman, UHD, photorealistic, shot on a Sony A7III --chaos 20 --ar 1:2 --style raw --stylize 250",
+    "Woman in a red jacket, snowy, in the style of hyper-realistic portraiture, caninecore, mountainous vistas, timeless beauty, palewave, iconic, distinctive noses --ar 72:101 --stylize 750 --v 6",
+    "Super Realism, Headshot of handsome young man, wearing dark gray sweater with buttons and big shawl collar, brown hair and short beard, serious look on his face, black background, soft studio lighting, portrait photography --ar 85:128 --v 6.0 --style",
+    "Super-realism, Purple Dreamy, a medium-angle shot of a young woman with long brown hair, wearing a pair of eye-level glasses, stands in front of a backdrop of purple and white lights. The womans eyes are closed, her lips are slightly parted, as if she is looking up at the sky. Her hair is cascading over her shoulders, framing her face. She is wearing a sleeveless top, adorned with tiny white dots, and a gold chain necklace around her neck. Her left earrings are dangling from her ears, adding a pop of color to the scene."
+]
+
+css = '''
+.gradio-container{max-width: 888px !important}
+h1{text-align:center}
+footer {
+    visibility: hidden
 }
-light_rag = LightRAG(rag_config)
-
-# ---------------------------
-# Granite Vision functions (for image and video)
-# ---------------------------
-
-# Set the Granite Vision model ID (adjust version as needed)
-GRANITE_MODEL_ID = "ibm-granite/granite-vision-3.2-2b"
-granite_processor = None
-granite_model = None
-
-def load_granite_model():
-    """Lazy load the Granite vision processor and model."""
-    global granite_processor, granite_model
-    if granite_processor is None or granite_model is None:
-        granite_processor = AutoProcessor.from_pretrained(GRANITE_MODEL_ID)
-        # Remove the .to(device) call to avoid moving a model already offloaded via accelerate.
-        granite_model = AutoModelForVision2Seq.from_pretrained(GRANITE_MODEL_ID, device_map="auto")
-    return granite_processor, granite_model
-
-def create_single_turn(image, text):
-    """
-    Creates a single-turn conversation message.
-    If an image is provided, it is added along with the text.
-    """
-    if image is None:
-        return {"role": "user", "content": [{"type": "text", "text": text}]}
-    else:
-        return {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": text}]}
-
-def generate_granite(image, prompt_text, max_new_tokens=1024, temperature=0.7, top_p=0.85, top_k=50, repetition_penalty=1.05):
-    """
-    Generates a response from the Granite Vision model given an image and prompt.
-    """
-    processor, model = load_granite_model()
-    conversation = [create_single_turn(image, prompt_text)]
-    inputs = processor.apply_chat_template(
-        conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
-    ).to(device)
-    output = model.generate(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        repetition_penalty=repetition_penalty,
-    )
-    decoded = processor.decode(output[0], skip_special_tokens=True)
-    parts = decoded.strip().split("<|assistant|>")
-    return parts[-1].strip()
-
-def generate_video_infer(video_path, prompt_text, max_new_tokens=1024, temperature=0.7, top_p=0.85, top_k=50, repetition_penalty=1.05):
-    """
-    Processes a video file by downsampling frames and sending them along with a prompt
-    to the Granite Vision model.
-    """
-    frames = downsample_video(video_path)
-    conversation_content = []
-    for img, ts in frames:
-        conversation_content.append({"type": "text", "text": f"Frame at {ts} sec:"})
-        conversation_content.append({"type": "image", "image": img})
-    conversation_content.append({"type": "text", "text": prompt_text})
-    conversation = [{"role": "user", "content": conversation_content}]
-    processor, model = load_granite_model()
-    inputs = processor.apply_chat_template(
-        conversation, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt"
-    ).to(device)
-    output = model.generate(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        repetition_penalty=repetition_penalty,
+.submit-btn {
+    background-color: #e34949 !important;
+    color: white !important;
+}
+.submit-btn:hover {
+    background-color: #ff3b3b !important;
+}
+'''
+
+with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
+            )
+            run_button = gr.Button("Generate as ( 768 x 1024 )🤗", scale=0, elem_classes="submit-btn")
+            
+            with gr.Accordion("Advanced options", open=True, visible=True):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                    visible=True
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                
+                with gr.Row(visible=True):
+                    width = gr.Slider(
+                        label="Width",
+                        minimum=512,
+                        maximum=2048,
+                        step=64,
+                        value=768,
+                    )
+                    height = gr.Slider(
+                        label="Height",
+                        minimum=512,
+                        maximum=2048,
+                        step=64,
+                        value=1024,
+                    )
+                
+                with gr.Row():
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=0.1,
+                        maximum=20.0,
+                        step=0.1,
+                        value=3.0,
+                    )
+                    num_inference_steps = gr.Slider(
+                        label="Number of inference steps",
+                        minimum=1,
+                        maximum=40,
+                        step=1,
+                        value=28,
+                    )
+
+                style_selection = gr.Radio(
+                    show_label=True,
+                    container=True,
+                    interactive=True,
+                    choices=STYLE_NAMES,
+                    value=DEFAULT_STYLE_NAME,
+                    label="Quality Style",
+                )
+        
+        with gr.Column(scale=2):
+            result = gr.Gallery(label="Result", columns=1, show_label=False)
+            
+            gr.Examples(
+                examples=examples,
+                inputs=prompt,
+                outputs=[result, seed],
+                fn=generate,
+                cache_examples=False,
+            )
+
+    gr.on(
+        triggers=[
+            prompt.submit,
+            run_button.click,
+        ],
+        fn=generate,
+        inputs=[
+            prompt,
+            seed,
+            width,
+            height,
+            guidance_scale,
+            randomize_seed,
+            style_selection,
+        ],
+        outputs=[result, seed],
+        api_name="run",
     )
-    decoded = processor.decode(output[0], skip_special_tokens=True)
-    parts = decoded.strip().split("<|assistant|>")
-    return parts[-1].strip()
-
-# ---------------------------
-# Unified generation function for ChatInterface
-# ---------------------------
-@spaces.GPU
-def generate(input_dict: dict, chat_history: list[dict],
-             max_new_tokens: int, temperature: float,
-             top_p: float, top_k: int, repetition_penalty: float):
-    """
-    Chat function that inspects the input text for special commands and routes:
-      - @rag: Uses the RAG pipeline.
-      - @granite: Uses Granite Vision for image understanding.
-      - @video-infer: Uses Granite Vision for video processing.
-    """
-    text = input_dict["text"]
-    files = input_dict.get("files", [])
-    lower_text = text.strip().lower()
-
-    # Optionally yield a progress message
-    yield "Processing your request..."
-    time.sleep(1)  # simulate processing delay
-
-    if lower_text.startswith("@rag"):
-        query = text[len("@rag"):].strip()
-        logging.info(f"@rag command: {query}")
-        context = light_rag.search(query)
-        answer, _ = light_rag.generate(query, context)
-        yield answer
-
-    elif lower_text.startswith("@granite"):
-        prompt_text = text[len("@granite"):].strip()
-        logging.info(f"@granite command: {prompt_text}")
-        if files:
-            # Expecting an image file (as a PIL image)
-            image = files[0]
-            answer = generate_granite(image, prompt_text, max_new_tokens, temperature, top_p, top_k, repetition_penalty)
-            yield answer
-        else:
-            yield "No image provided for @granite command."
-
-    elif lower_text.startswith("@video-infer"):
-        prompt_text = text[len("@video-infer"):].strip()
-        logging.info(f"@video-infer command: {prompt_text}")
-        if files:
-            # Expecting a video file (the file path)
-            video_path = files[0]
-            answer = generate_video_infer(video_path, prompt_text, max_new_tokens, temperature, top_p, top_k, repetition_penalty)
-            yield answer
-        else:
-            yield "No video provided for @video-infer command."
-
-    else:
-        # Default behavior: use RAG pipeline for text query.
-        query = text.strip()
-        logging.info(f"Default text query: {query}")
-        context = light_rag.search(query)
-        answer, _ = light_rag.generate(query, context)
-        yield answer
-
-# ---------------------------
-# Gradio ChatInterface using MultimodalTextbox
-# ---------------------------
-
-demo = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        gr.Slider(label="Max new tokens", minimum=1, maximum=2048, step=1, value=1024),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7),
-        gr.Slider(label="Top-p", minimum=0.1, maximum=1.0, step=0.1, value=0.85),
-        gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=50),
-        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.05),
-    ],
-    examples=[
-        # Examples show how to use the command prefixes.
-        [{"text": "@rag What models are available in Watsonx?"}],
-        [{"text": "@granite Describe the image", "files": [str(Path("examples") / "sample_image.png")]}],
-        [{"text": "@video-infer Summarize the event in the video", "files": [str(Path("examples") / "sample_video.mp4")]}],
-    ],
-    cache_examples=False,
-    type="messages",
-    description=(
-        "# **Multimodal OCR Granite Vision**\n\n"
-        "Enter a command in the text input (with optional file uploads) using one of the following prefixes:\n\n"
-        "- **@rag**: For retrieval-augmented generation (e.g. PDFs, documents).\n"
-        "- **@granite**: For image understanding using Granite Vision.\n"
-        "- **@video-infer**: For video understanding (video is downsampled into frames).\n\n"
-        "For example:\n```\n@rag What is the revenue trend?\n```\n```\n@granite Describe this image\n```\n```\n@video-infer Summarize the event in this video\n```"
-    ),
-    fill_height=True,
-    textbox=gr.MultimodalTextbox(
-        label="Query Input",
-        file_types=["image", "video", "pdf"],
-        file_count="multiple",
-        placeholder="@rag, @granite, or @video-infer followed by your prompt"
-    ),
-    stop_btn="Stop Generation",
-    multimodal=True,
-)
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=40).launch()