Update app.py

app.py

@@ -1,160 +1,36 @@
-from __future__ import annotations
-
-import math
-import random
-
 import gradio as gr
 import torch
-from PIL import Image, ImageOps
+from PIL import Image
 from diffusers import StableDiffusionInstructPix2PixPipeline
 
-# Path to the SafeTensor model in Colab
-model_path = "/content/uberRealisticPornMerge_urpmv12.instruct-pix2pix.safetensors"
-
-def main():
-    # Load the SafeTensor model from Colab
-    safe_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_path, torch_dtype=torch.float16, safety_checker=None).to("cuda")
-    example_image = Image.open("/content/T/imgs_example.jpg").convert("RGB")
-
-    def load_example(
-        steps: int,
-        randomize_seed: bool,
-        seed: int,
-        randomize_cfg: bool,
-        text_cfg_scale: float,
-        image_cfg_scale: float,
-    ):
-        example_instruction = random.choice(example_instructions)
-        return [example_image, example_instruction] + generate(
-            example_image,
-            example_instruction,
-            steps,
-            randomize_seed,
-            seed,
-            randomize_cfg,
-            text_cfg_scale,
-            image_cfg_scale,
-        )
-
-    def generate(
-        input_image: Image.Image,
-        instruction: str,
-        steps: int,
-        randomize_seed: bool,
-        seed: int,
-        randomize_cfg: bool,
-        text_cfg_scale: float,
-        image_cfg_scale: float,
-    ):
-        seed = random.randint(0, 100000) if randomize_seed else seed
-        text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale
-        image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale
-
-        width, height = input_image.size
-        factor = 512 / max(width, height)
-        factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)
-        width = int((width * factor) // 64) * 64
-        height = int((height * factor) // 64) * 64
-        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)
-
-        if instruction == "":
-            return [input_image, seed]
-
-        generator = torch.manual_seed(seed)
-        edited_image = safe_pipe(
-            instruction, image=input_image,
-            guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,
-            num_inference_steps=steps, generator=generator,
-        ).images[0]
-        return [seed, text_cfg_scale, image_cfg_scale, edited_image]
-
-    def reset():
-        return [0, "Randomize Seed", 1371, "Fix CFG", 7.5, 1.5, None]
-
-    with gr.Blocks() as demo:
-        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">
-   InstructPix2Pix: Learning to Follow Image Editing Instructions
-</h1>
-<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.
-<br/>
-<a href="https://huggingface.co/spaces/timbrooks/instruct-pix2pix?duplicate=true">
-<img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>
-<p/>""")
-        with gr.Row():
-            with gr.Column(scale=1, min_width=100):
-                generate_button = gr.Button("Generate")
-            with gr.Column(scale=1, min_width=100):
-                load_button = gr.Button("Load Example")
-            with gr.Column(scale=1, min_width=100):
-                reset_button = gr.Button("Reset")
-            with gr.Column(scale=3):
-                instruction = gr.Textbox(lines=1, label="Edit Instruction", interactive=True)
-
-        with gr.Row():
-            input_image = gr.Image(label="Input Image", type="pil", interactive=True)
-            edited_image = gr.Image(label=f"Edited Image", type="pil", interactive=False)
-            input_image.style(height=512, width=512)
-            edited_image.style(height=512, width=512)
-
-        with gr.Row():
-            steps = gr.Number(value=50, precision=0, label="Steps", interactive=True)
-            randomize_seed = gr.Radio(
-                ["Fix Seed", "Randomize Seed"],
-                value="Randomize Seed",
-                type="index",
-                show_label=False,
-                interactive=True,
-            )
-            seed = gr.Number(value=1371, precision=0, label="Seed", interactive=True)
-            randomize_cfg = gr.Radio(
-                ["Fix CFG", "Randomize CFG"],
-                value="Fix CFG",
-                type="index",
-                show_label=False,
-                interactive=True,
-            )
-            text_cfg_scale = gr.Number(value=7.5, label=f"Text CFG", interactive=True)
-            image_cfg_scale = gr.Number(value=1.5, label=f"Image CFG", interactive=True)
-
-        gr.Markdown(help_text)
-        
-        # Define actions for buttons
-        load_button.click(
-            fn=load_example,
-            inputs=[
-                steps,
-                randomize_seed,
-                seed,
-                randomize_cfg,
-                text_cfg_scale,
-                image_cfg_scale,
-            ],
-            outputs=[input_image, instruction, seed, text_cfg_scale, image_cfg_scale, edited_image],
-        )
-        generate_button.click(
-            fn=generate,
-            inputs=[
-                input_image,
-                instruction,
-                steps,
-                randomize_seed,
-                seed,
-                randomize_cfg,
-                text_cfg_scale,
-                image_cfg_scale,
-            ],
-            outputs=[seed, text_cfg_scale, image_cfg_scale, edited_image],
-        )
-        reset_button.click(
-            fn=reset,
-            inputs=[],
-            outputs=[steps, randomize_seed, seed, randomize_cfg, text_cfg_scale, image_cfg_scale, edited_image],
-        )
-
-    demo.queue(concurrency_count=1)
-    demo.launch(share=False)
-
-
-if __name__ == "__main__":
-    main()
- 
+# Define the path to the SafeTensor model
+model_path = "/content/UberRealisticPornMerge_urpmv12.instruct-pix2pix.safetensors"
+
+# Load the SafeTensor model
+safe_pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_path, torch_dtype=torch.float16, safety_checker=None).to("cuda")
+
+def generate_edited_image(input_image):
+    # Convert the Gradio Image object to a PIL Image
+    input_image_pil = Image.fromarray(input_image.astype('uint8'), 'RGB')
+    
+    # Generate the edited image using the SafeTensor model
+    edited_image = safe_pipe(instruction="", image=input_image_pil, num_inference_steps=50).images[0]
+    
+    # Convert the edited image back to Gradio Image format
+    edited_image_gradio = edited_image.cpu().numpy().astype('uint8')
+    
+    return edited_image_gradio
+
+# Define the input and output components for the Gradio app
+input_image = gr.inputs.Image(label="Upload an Input Image")
+output_image = gr.outputs.Image(label="Edited Image")
+
+# Create the Gradio interface
+gr.Interface(
+    fn=generate_edited_image,
+    inputs=input_image,
+    outputs=output_image,
+    title="SafeTensor Image Editing",
+    description="Upload an image and generate an edited image using a SafeTensor model.",
+    capture_session=True  # This ensures that we use the same session for model inference
+).launch()