Update app.py

app.py

@@ -1,110 +1,37 @@
-import os
-import random
-import uuid
-import json
-import time
-import asyncio
-from threading import Thread
-
-import gradio as gr
 import spaces
+import gradio as gr
 import torch
-import numpy as np
 from PIL import Image
-import edge_tts
-
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-    Qwen2VLForConditionalGeneration,
-    AutoProcessor,
-)
-from transformers.image_utils import load_image
 from diffusers import DiffusionPipeline
+import random
+import uuid
+from typing import Tuple
+import numpy as np
 
-DESCRIPTION = """
-# QwQ Edge 💬 with Flux.1
-"""
-
-css = '''
-h1 {
-  text-align: center;
-  display: block;
-}
-
-#duplicate-button {
-  margin: auto;
-  color: #fff;
-  background: #1565c0;
-  border-radius: 100vh;
-}
-'''
-
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-# --------------------------
-# Text Generation Components
-# --------------------------
-
-# Load text-only model and tokenizer
-model_id = "prithivMLmods/FastThink-0.5B-Tiny"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-)
-model.eval()
-
-TTS_VOICES = [
-    "en-US-JennyNeural",  # @tts1
-    "en-US-GuyNeural",    # @tts2
-]
-
-# Multimodal model (text+vision)
-MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct" 
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
-model_m = Qwen2VLForConditionalGeneration.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float16
-).to("cuda").eval()
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
 
-async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
-    """Convert text to speech using Edge TTS and save as MP3"""
-    communicate = edge_tts.Communicate(text, voice)
-    await communicate.save(output_file)
-    return output_file
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
 
-def clean_chat_history(chat_history):
-    """
-    Filter out any chat entries whose "content" is not a string.
-    This helps prevent errors when concatenating previous messages.
-    """
-    cleaned = []
-    for msg in chat_history:
-        if isinstance(msg, dict) and isinstance(msg.get("content"), str):
-            cleaned.append(msg)
-    return cleaned
+MAX_SEED = np.iinfo(np.int32).max
 
-# --------------------------
-# Flux.1 Image Generation
-# --------------------------
+if not torch.cuda.is_available():
+    DESCRIPTIONz += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
 
-# Set up the Flux.1 pipeline
 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")
 
-# Define style prompts
 style_list = [
     {
         "name": "3840 x 2160",
@@ -123,48 +50,17 @@ style_list = [
         "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)
 
-MAX_SEED = np.iinfo(np.int32).max
-
-def save_image(img: Image.Image) -> str:
-    """Save a PIL image with a unique filename and return the path."""
-    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
-
-def progress_bar_html(label: str) -> str:
-    """
-    Returns an HTML snippet for a thin progress bar with a label.
-    The progress bar is styled as a dark red animated bar.
-    """
-    return f'''
-<div style="display: flex; align-items: center;">
-    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
-    <div style="width: 110px; height: 5px; background-color: #f0f0f0; border-radius: 2px; overflow: hidden;">
-        <div style="width: 100%; height: 100%; background-color: #ff5900; animation: loading 1.5s linear infinite;"></div>
-    </div>
-</div>
-<style>
-@keyframes loading {{
-    0% {{ transform: translateX(-100%); }}
-    100% {{ transform: translateX(100%); }}
-}}
-</style>
-    '''
-
-@spaces.GPU(duration=60, enable_queue=True)
-def generate_image_fn(
+@spaces.GPU(duration=60, enable_queue=True
+def generate(
     prompt: str,
     seed: int = 0,
     width: int = 1024,
@@ -174,11 +70,13 @@ def generate_image_fn(
     style_name: str = DEFAULT_STYLE_NAME,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """Generate images using the Flux.1 pipeline."""
     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,
@@ -189,160 +87,124 @@ def generate_image_fn(
         output_type="pil",
     ).images
     image_paths = [save_image(img) for img in images]
+    print(image_paths)
     return image_paths, seed
 
-# --------------------------
-# Chat and Multimodal Generation
-# --------------------------
-
-@spaces.GPU
-def generate(
-    input_dict: dict,
-    chat_history: list[dict],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
-):
-    """
-    Generates chatbot responses with support for multimodal input, TTS, and image generation using Flux.1.
-    Special commands:
-      - "@tts1" or "@tts2": triggers text-to-speech.
-      - "@image": triggers image generation using the Flux.1 pipeline.
-    """
-    text = input_dict["text"]
-    files = input_dict.get("files", [])
-
-    if text.strip().lower().startswith("@image"):
-        # Remove the "@image" tag and use the rest as prompt
-        prompt_img = text[len("@image"):].strip()
-        # Show animated progress bar for image generation
-        yield progress_bar_html("Generating Image")
-        image_paths, used_seed = generate_image_fn(
-            prompt=prompt_img,
-            seed=1,
-            width=1024,
-            height=1024,
-            guidance_scale=3,
-            randomize_seed=True,
-            style_name=DEFAULT_STYLE_NAME,
-        )
-        # Once done, yield the generated image
-        yield gr.Image(image_paths[0])
-        return  # Exit early
-
-    tts_prefix = "@tts"
-    is_tts = any(text.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 3))
-    voice_index = next((i for i in range(1, 3) if text.strip().lower().startswith(f"{tts_prefix}{i}")), None)
-    
-    if is_tts and voice_index:
-        voice = TTS_VOICES[voice_index - 1]
-        text = text.replace(f"{tts_prefix}{voice_index}", "").strip()
-        # Clear previous chat history for a fresh TTS request.
-        conversation = [{"role": "user", "content": text}]
-    else:
-        voice = None
-        # Remove any stray @tts tags and build the conversation history.
-        text = text.replace(tts_prefix, "").strip()
-        conversation = clean_chat_history(chat_history)
-        conversation.append({"role": "user", "content": text})
-
-    if files:
-        if len(files) > 1:
-            images = [load_image(image) for image in files]
-        elif len(files) == 1:
-            images = [load_image(files[0])]
-        else:
-            images = []
-        messages = [{
-            "role": "user",
-            "content": [
-                *[{"type": "image", "image": image} for image in images],
-                {"type": "text", "text": text},
-            ]
-        }]
-        prompt_multimodal = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = processor(text=[prompt_multimodal], images=images, return_tensors="pt", padding=True).to("cuda")
-        streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {**inputs, "streamer": streamer, "max_new_tokens": max_new_tokens}
-        thread = Thread(target=model_m.generate, kwargs=generation_kwargs)
-        thread.start()
-
-        buffer = ""
-        # Show animated progress bar for multimodal generation
-        yield progress_bar_html("Thinking...")
-        for new_text in streamer:
-            buffer += new_text
-            buffer = buffer.replace("<|im_end|>", "")
-            time.sleep(0.01)
-            yield buffer
-    else:
-        input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
-        if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-            input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-            gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-        input_ids = input_ids.to(model.device)
-        streamer = TextIteratorStreamer(tokenizer, timeout=20.0, skip_prompt=True, skip_special_tokens=True)
-        generation_kwargs = {
-            "input_ids": input_ids,
-            "streamer": streamer,
-            "max_new_tokens": max_new_tokens,
-            "do_sample": True,
-            "top_p": top_p,
-            "top_k": top_k,
-            "temperature": temperature,
-            "num_beams": 1,
-            "repetition_penalty": repetition_penalty,
-        }
-        t = Thread(target=model.generate, kwargs=generation_kwargs)
-        t.start()
-
-        outputs = []
-        # Show animated progress bar for text generation
-        yield progress_bar_html("Thinking...")
-        for new_text in streamer:
-            outputs.append(new_text)
-            yield "".join(outputs)
-
-        final_response = "".join(outputs)
-        yield final_response
-
-        # If TTS was requested, convert the final response to speech.
-        if is_tts and voice:
-            output_file = asyncio.run(text_to_speech(final_response, voice))
-            yield gr.Audio(output_file, autoplay=True)
+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."
+]
 
-# --------------------------
-# Gradio Chat Interface
-# --------------------------
+css = '''
+.gradio-container{max-width: 888px !important}
+h1{text-align:center}
+footer {
+    visibility: hidden
+}
+.submit-btn {
+    background-color: #e34949 !important;
+    color: white !important;
+}
+.submit-btn:hover {
+    background-color: #ff3b3b !important;
+}
+'''
 
-demo = gr.ChatInterface(
-    fn=generate,
-    additional_inputs=[
-        gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1, value=DEFAULT_MAX_NEW_TOKENS),
-        gr.Slider(label="Temperature", minimum=0.1, maximum=4.0, step=0.1, value=0.6),
-        gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
-        gr.Slider(label="Top-k", minimum=1, maximum=1000, step=1, value=50),
-        gr.Slider(label="Repetition penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.2),
-    ],
-    examples=[
-        ["@image A futuristic cityscape at sunset with vibrant colors"],
-        ["Python Program for Array Rotation"],
-        ["@tts1 Who is Nikola Tesla, and why did he die?"],
-        [{"text": "Extract JSON from the image", "files": ["examples/document.jpg"]}],
-        [{"text": "summarize the letter", "files": ["examples/1.png"]}],
-        ["@tts2 What causes rainbows to form?"],
-    ],
-    cache_examples=False,
-    type="messages",
-    description=DESCRIPTION,
-    css=css,
-    fill_height=True,
-    textbox=gr.MultimodalTextbox(label="Query Input", file_types=["image"], file_count="multiple", placeholder="‎ @tts1, @tts2-voices, @image-image gen, default [text, vision]"),
-    stop_btn="Stop Generation",
-    multimodal=True,
-)
+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",
+    )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch(share=True)
+    demo.queue(max_size=40).launch()