Update app.py

app.py

@@ -1,339 +1,37 @@
-# import gradio as gr
-# import numpy as np
-# import random
-# from diffusers import DiffusionPipeline
-# import torch
-# import transformers
-
-# # Perform cache migration
-# transformers.utils.move_cache()
-
-# device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# if torch.cuda.is_available():
-#     torch.cuda.max_memory_allocated(device=device)
-#     pipe = DiffusionPipeline.from_pretrained(
-#         "stabilityai/sdxl-turbo",
-#         torch_dtype=torch.float16,
-#         variant="fp16",
-#         use_safetensors=True,
-#     )
-#     pipe.enable_xformers_memory_efficient_attention()
-#     pipe = pipe.to(device)
-# else:
-#     pipe = DiffusionPipeline.from_pretrained(
-#         "stabilityai/sdxl-turbo", use_safetensors=True
-#     )
-#     pipe = pipe.to(device)
-
-#     # Quantize the model
-#     pipe.unet = torch.quantization.convert(pipe.unet, inplace=True)
-
-# MAX_SEED = np.iinfo(np.int32).max
-# MAX_IMAGE_SIZE = 512
-
-
-# def generate_image(
-#     seed, prompt, negative_prompt, guidance_scale, num_inference_steps, width, height
-# ):
-#     try:
-#         generator = torch.Generator().manual_seed(seed)
-#         image = pipe(
-#             prompt=prompt,
-#             negative_prompt=negative_prompt,
-#             guidance_scale=guidance_scale,
-#             num_inference_steps=num_inference_steps,
-#             width=width,
-#             height=height,
-#             generator=generator,
-#         ).images[0]
-#         return image
-#     except Exception as e:
-#         print(f"Error generating image with seed {seed}: {e}")
-#         return None
-
-
-# def infer(
-#     prompt,
-#     negative_prompt,
-#     seed,
-#     randomize_seed,
-#     width,
-#     height,
-#     guidance_scale,
-#     num_inference_steps,
-# ):
-
-#     if randomize_seed:
-#         seeds = [random.randint(0, MAX_SEED) for _ in range(2)]
-#     else:
-#         seeds = [seed, seed + 1]
-
-#     images = []
-#     for seed in seeds:
-#         image = generate_image(
-#             seed,
-#             prompt,
-#             negative_prompt,
-#             guidance_scale,
-#             num_inference_steps,
-#             width,
-#             height,
-#         )
-#         images.append(image)
-
-#     return images
-
-
-# examples = [
-#     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-#     "An astronaut riding a green horse",
-#     "A delicious ceviche cheesecake slice",
-# ]
-
-# css = """
-# #col-container {
-#     margin: 0 auto;
-#     max-width: 520px;
-# }
-# """
-
-# if torch.cuda.is_available():
-#     power_device = "GPU"
-# else:
-#     power_device = "CPU"
-
-# with gr.Blocks(css=css) as demo:
-
-#     with gr.Column(elem_id="col-container"):
-#         gr.Markdown(
-#             f"""
-#         # Text-to-Image Gradio Template
-#         Currently running on {power_device}.
-#         """
-#         )
-
-#         with gr.Row():
-
-#             prompt = gr.Text(
-#                 label="Prompt",
-#                 show_label=False,
-#                 max_lines=1,
-#                 placeholder="Enter your prompt",
-#                 container=False,
-#             )
-
-#             run_button = gr.Button("Run", scale=0)
-
-#         result1 = gr.Image(label="Result 1", show_label=False)
-#         result2 = gr.Image(label="Result 2", show_label=False)
-
-#         with gr.Accordion("Advanced Settings", open=False):
-
-#             negative_prompt = gr.Text(
-#                 label="Negative prompt",
-#                 max_lines=1,
-#                 placeholder="Enter a negative prompt",
-#                 visible=False,
-#             )
-
-#             seed = gr.Slider(
-#                 label="Seed",
-#                 minimum=0,
-#                 maximum=MAX_SEED,
-#                 step=1,
-#                 value=0,
-#             )
-
-#             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-#             with gr.Row():
-
-#                 width = gr.Slider(
-#                     label="Width",
-#                     minimum=256,
-#                     maximum=MAX_IMAGE_SIZE,
-#                     step=32,
-#                     value=512,
-#                 )
-
-#                 height = gr.Slider(
-#                     label="Height",
-#                     minimum=256,
-#                     maximum=MAX_IMAGE_SIZE,
-#                     step=32,
-#                     value=512,
-#                 )
-
-#             with gr.Row():
-
-#                 guidance_scale = gr.Slider(
-#                     label="Guidance scale",
-#                     minimum=0.0,
-#                     maximum=10.0,
-#                     step=0.1,
-#                     value=0.0,
-#                 )
-
-#                 num_inference_steps = gr.Slider(
-#                     label="Number of inference steps",
-#                     minimum=1,
-#                     maximum=50,  # Ensure the number of steps is reasonable
-#                     step=1,
-#                     value=2,
-#                 )
-
-#         gr.Examples(examples=examples, inputs=[prompt])
-
-#     run_button.click(
-#         fn=infer,
-#         inputs=[
-#             prompt,
-#             negative_prompt,
-#             seed,
-#             randomize_seed,
-#             width,
-#             height,
-#             guidance_scale,
-#             num_inference_steps,
-#         ],
-#         outputs=[result1, result2],
-#     )
-
-# demo.queue().launch()
-
 import gradio as gr
 import numpy as np
-from PIL import Image
-import requests
-from io import BytesIO
 import random
 from diffusers import DiffusionPipeline
 import torch
 import transformers
-from tqdm import tqdm
 
 # Perform cache migration
 transformers.utils.move_cache()
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def load_pipeline():
-    if device == "cuda":
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/sdxl-turbo",
-            torch_dtype=torch.float16,
-            variant="fp16",
-            use_safetensors=True,
-        )
-        pipe.enable_xformers_memory_efficient_attention()
-    else:
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/sdxl-turbo", use_safetensors=True
-        )
-        # pipe.unet = torch.quantization.convert(pipe.unet, inplace=True)
-    return pipe.to(device)
-
-pipe = load_pipeline()
-
-#################### our model ####################
-import warnings
-import torch.utils
-import torch.utils.checkpoint
-from transformers import BitsAndBytesConfig, InstructBlipProcessor, InstructBlipForConditionalGeneration
-
-
-# Filter out specific warnings by message
-warnings.filterwarnings("ignore", message="Repo card metadata block was not found. Setting CardData to empty.")
-warnings.filterwarnings(
-    "ignore",
-    message="torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly.",
-    category=UserWarning
-)
-
-model_checkpoint = "Salesforce/instructblip-vicuna-7b"
-lora_weights_repo_id = "instructblip-vicuna-7b-peft-lora"
-# peft_model_id = "NoyHanan/instructblip-vicuna-7b-peft-lora-6400"
-peft_model_id = "NoyHanan/instructblip-vicuna-7b-peft-lora-1600"
-
-prompt_format = """###USER:ֿ\nHere is an image. Please analyze the image and enhance the base prompt by integrating detailed observations, including colors, textures, lighting, and key visual elements, while staying true to the original description. The goal is to produce a more vibrant, detailed, and visually appealing image. The base prompt is: "{base_prompt}"\n###ASSISTANT:\n"""
-
-text = prompt_format.format(base_prompt="enter prompt here")
-
-
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-)
-
-def load_prompt_optimization_model_and_processor():
-    prompt_optimizer_processor = InstructBlipProcessor.from_pretrained(model_checkpoint, legacy=False, quantization_config=bnb_config, device_map="auto")
-    prompt_optimizer_model = InstructBlipForConditionalGeneration.from_pretrained(model_checkpoint, quantization_config=bnb_config, device_map="auto")
-    prompt_optimizer_model.load_adapter(peft_model_id)
-    prompt_optimizer_model.tie_weights()
-
-    return prompt_optimizer_processor, prompt_optimizer_model
-
-processor, model = load_prompt_optimization_model_and_processor()
-
-def get_enhanced_prompts_from_our_model(prompt, selected_image):
-    print("Start generating prompts using our model")
-    model.eval()
-
-    enhanced_prompts_list = []
-
-    text = prompt_format.format(base_prompt=prompt)
-    inputs = processor(images=selected_image, text=text, return_tensors="pt").to(device)
-
-
-
-    # # outputs = trainer.model.generate(
-    # outputs = model.generate(
-    #     **inputs,
-    #     do_sample=True,
-    #     # num_beams=5,
-    #     num_beams = 3,
-    #     # max_length=516,
-    #     max_length=256,
-    #     min_length=1,
-    #     top_p=0.9,
-    #     repetition_penalty=1.5,
-    #     # length_penalty=1.0,
-    #     length_penalty = 0.5,
-    #     temperature=0.1,  # give 3 of the same prompt
-    # )
-
-    for i in tqdm(range(3)):
-      # outputs = trainer.model.generate(
-      outputs = model.generate(
-          **inputs,
-          do_sample=True,
-          # num_beams=5,
-          num_beams = 3,
-          # max_length=516,
-          max_length=256,
-          min_length=1,
-          top_p=0.9,
-          repetition_penalty=1.5,
-          # length_penalty=1.0,
-          length_penalty = 0.5,
-          # temperature=0.1,  # give 3 of the same prompt
-          temperature=0.8
-      )
-      res = processor.batch_decode(outputs, skip_special_tokens=True)
-      generated_text = res[0].strip()
-      enhanced_prompts_list.append(generated_text)
-      print(generated_text)
-      torch.cuda.empty_cache()
-
-    print("Finish generating prompts using our model")
-    return enhanced_prompts_list
-###################################################
+if torch.cuda.is_available():
+    torch.cuda.max_memory_allocated(device=device)
+    pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/sdxl-turbo",
+        torch_dtype=torch.float16,
+        variant="fp16",
+        use_safetensors=True,
+    )
+    pipe.enable_xformers_memory_efficient_attention()
+    pipe = pipe.to(device)
+else:
+    pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/sdxl-turbo", use_safetensors=True
+    )
+    pipe = pipe.to(device)
+
+    # Quantize the model
+    pipe.unet = torch.quantization.convert(pipe.unet, inplace=True)
 
 MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+MAX_IMAGE_SIZE = 512
+
 
 def generate_image(
     seed, prompt, negative_prompt, guidance_scale, num_inference_steps, width, height
@@ -343,10 +41,10 @@ def generate_image(
         image = pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
-            guidance_scale=float(guidance_scale),
-            num_inference_steps=int(num_inference_steps),
-            width=int(width),
-            height=int(height),
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
             generator=generator,
         ).images[0]
         return image
@@ -355,30 +53,23 @@ def generate_image(
         return None
 
 
-# def get_enhanced_prompts_from_our_model(prompt, selected_image):
-#     # This prompts will be generated by our model
-#     return ["cat", "frog", "giraffe"]
-
-first_4_images = []
-seeds = []
-prompt_entered_by_the_user = ""
-improved_prompts = []
-
-
-def create_4_images_from_original_prompt(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
-
-    global first_4_images
-    global seeds
-    global prompt_entered_by_the_user
-
-    prompt_entered_by_the_user = prompt
+def infer(
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    num_inference_steps,
+):
 
     if randomize_seed:
-        seeds = [random.randint(0, MAX_SEED) for _ in range(4)]
+        seeds = [random.randint(0, MAX_SEED) for _ in range(2)]
     else:
-        seeds = [seed, seed + 1, seed + 2, seed + 3]
+        seeds = [seed, seed + 1]
 
-    first_4_images = []
+    images = []
     for seed in seeds:
         image = generate_image(
             seed,
@@ -389,73 +80,12 @@ def create_4_images_from_original_prompt(prompt, negative_prompt, seed, randomiz
             width,
             height,
         )
-        first_4_images.append(image)
-
-    grid = create_image_grid(first_4_images, rows=2, cols=2)
-    return grid
-
-def enhance_prompt_and_create_new_images(prompt, selected_index, width, height, guidance_scale, num_inference_steps, negative_prompt):
-
-    global improved_prompts
-
-    selected_image = first_4_images[selected_index]
-
-    improved_prompts = get_enhanced_prompts_from_our_model(prompt, selected_image)
-
-    seed = seeds[selected_index]
-    final_4_images = [None] * 4
-    images_from_improved_prompt = []
-
-    for prompt in improved_prompts:
-        image = generate_image(
-            seed,
-            prompt,
-            negative_prompt,
-            guidance_scale,
-            num_inference_steps,
-            width,
-            height,
-        )
-        if image is not None:
-            images_from_improved_prompt.append(image)
-
-    for i in range(4):
-        if i == selected_index:
-            final_4_images[i] = selected_image
-        elif images_from_improved_prompt:
-            final_4_images[i] = images_from_improved_prompt.pop(0)
-        else:
-            final_4_images[i] = selected_image
-    grid = create_image_grid(final_4_images, rows=2, cols=2)
-    return grid
-
+        images.append(image)
 
+    return images
 
 
-def create_image_grid(images, rows, cols):
-    assert len(images) == rows * cols
-    w, h = images[0].size
-    grid = Image.new("RGB", size=(cols * w, rows * h))
-    for i, img in enumerate(images):
-        grid.paste(img, box=(i % cols * w, i // cols * h))
-    return grid
-
-def get_final_prompt(prompt, selected_index_left, selected_index_right):
-    global improved_prompts
-
-    prompt = f"Prompt did not improved.\n\nOriginal Prompt = {prompt}"
-
-    if selected_index_left == selected_index_right:
-        return prompt
-    else:
-        improved_prompts.insert(selected_index_left, prompt)
-        new_prompt = improved_prompts[selected_index_right]
-        improved_prompts.remove(prompt)
-        return f"Prompt improved succesfuly.\n\nEnhanced Prompt = {new_prompt}"
-
 examples = [
-    "Light emerald green wine jar with deer head lid, gold lines, small cracks",
-    "Textured tempura painting of a friendly waitress serving coffee",
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
     "A delicious ceviche cheesecake slice",
@@ -463,183 +93,102 @@ examples = [
 
 css = """
 #col-container {
-    margin: 0;
-    max-width: 20px;  # Smaller width
-}
-body {
-    display: flex;
-    justify-content: space-between;  # Space between left and right columns
-}
-.spacer {
-    height: 130px;  # Large spacer
+    margin: 0 auto;
+    max-width: 520px;
 }
 """
 
-power_device = device
+if torch.cuda.is_available():
+    power_device = "GPU"
+else:
+    power_device = "CPU"
 
 with gr.Blocks(css=css) as demo:
 
-    with gr.Row():
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            f"""
+        # Text-to-Image Gradio Template
+        Currently running on {power_device}.
+        """
+        )
+
+        with gr.Row():
 
-        #################### First 4 Images Component ####################
-        with gr.Column(elem_id="col-container"):
-            gr.Markdown(
-                f"""
-            # Enter Prompt & Select The Best Image
-            Currently running on {power_device}.
-            """
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt",
+                container=False,
             )
 
-            with gr.Row():
-                original_prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter your prompt",
-                    container=False,
-                )
+            run_button = gr.Button("Run", scale=0)
 
-                enter_prompt_buttom = gr.Button("Run", scale=0)
+        result1 = gr.Image(label="Result 1", show_label=False)
+        result2 = gr.Image(label="Result 2", show_label=False)
 
-            result_left = gr.Image(label="Result", show_label=False)
+        with gr.Accordion("Advanced Settings", open=False):
 
-            selected_index_left = gr.Radio(
-                label="Select Image",
-                choices=[
-                    "0 - Upper Left",
-                    "1 - Upper Right",
-                    "2 - Lower Left",
-                    "3 - Lower Right",
-                ],
-                type="index",
+            negative_prompt = gr.Text(
+                label="Negative prompt",
+                max_lines=1,
+                placeholder="Enter a negative prompt",
+                visible=False,
             )
 
-            selected_text_left = gr.Text(label="Selected Image Index", visible=False)
-
-            def update_selected_index(index):
-                return str(index)
-
-            selected_index_left.change(
-                fn=update_selected_index,
-                inputs=selected_index_left,
-                outputs=selected_text_left,
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
             )
 
-            with gr.Accordion("Advanced Settings", open=False):
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-                negative_prompt = gr.Text(
-                    label="Negative prompt",
-                    max_lines=1,
-                    placeholder="Enter a negative prompt",
-                    visible=False,
-                )
+            with gr.Row():
 
-                seed = gr.Slider(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=0,
+                width = gr.Slider(
+                    label="Width",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=512,
                 )
 
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-                with gr.Row():
-
-                    width = gr.Slider(
-                        label="Width",
-                        minimum=256,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=512,
-                    )
-
-                    height = gr.Slider(
-                        label="Height",
-                        minimum=256,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=512,
-                    )
-
-                with gr.Row():
-
-                    guidance_scale = gr.Slider(
-                        label="Guidance scale",
-                        minimum=0.0,
-                        maximum=10.0,
-                        step=0.1,
-                        value=0.0,
-                    )
-
-                    num_inference_steps = gr.Slider(
-                        label="Number of inference steps",
-                        minimum=1,
-                        maximum=12,
-                        step=1,
-                        value=2,
-                    )
-
-            gr.Examples(examples=examples, inputs=[original_prompt])
-        ##################################################################
-
-        #################### Final 4 Images Component ####################
-        with gr.Column(elem_id="col-container"):
-            gr.Markdown(
-                f"""
-            # Select The Best Image
-            you can choose the same one as before if you want
-            """
-            )
-
-            # Adding empty rows for spacing
-            gr.Markdown(" ", elem_classes=["spacer"])
-
-            result_right = gr.Image(label="Result", show_label=False)
-
-            selected_index_right = gr.Radio(
-                label="Select Image",
-                choices=[
-                    "0 - Upper Left",
-                    "1 - Upper Right",
-                    "2 - Lower Left",
-                    "3 - Lower Right",
-                ],
-                type="index",
-            )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=512,
+                )
 
-            selected_text_right = gr.Text(label="Selected Image Index", visible=False)
+            with gr.Row():
 
-            selected_index_right.change(
-                fn=update_selected_index,
-                inputs=selected_index_right,
-                outputs=selected_text_right,
-            )
-        ##################################################################
-
-        ##################### Final Prompt Component #####################
-        with gr.Column(elem_id="col-container"):
-            gr.Markdown(
-                f"""
-            # Enhance Prompt
-            """
-            )
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=0.0,
+                )
 
-            enhanced_prompt_output = gr.Textbox(
-                label="Final Prompt", interactive=False
-            )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,  # Ensure the number of steps is reasonable
+                    step=1,
+                    value=2,
+                )
 
-            selected_index_right.change(
-                fn=get_final_prompt,
-                inputs=[original_prompt, selected_index_left, selected_index_right],
-                outputs=[enhanced_prompt_output],
-            )
-        ##################################################################
+        gr.Examples(examples=examples, inputs=[prompt])
 
-    enter_prompt_buttom.click(
-        fn=create_4_images_from_original_prompt,
+    run_button.click(
+        fn=infer,
         inputs=[
-            original_prompt,
+            prompt,
             negative_prompt,
             seed,
             randomize_seed,
@@ -648,22 +197,7 @@ with gr.Blocks(css=css) as demo:
             guidance_scale,
             num_inference_steps,
         ],
-        outputs=[result_left],
+        outputs=[result1, result2],
     )
 
-
-    selected_index_left.change(
-    fn=enhance_prompt_and_create_new_images,
-    inputs=[
-        original_prompt,
-        selected_index_left,
-        width,
-        height,
-        guidance_scale,
-        num_inference_steps,
-        negative_prompt,
-    ],
-    outputs=[result_right],
-)
-
-demo.queue().launch(share=True)
+demo.queue().launch()