Update app.py

app.py

@@ -1,102 +1,250 @@
 import requests
 from requests.adapters import HTTPAdapter
-from requests.packages.urllib3.util.retry import Retry
+from urllib3.util.retry import Retry
 import json
 import base64
 import time
-import gradio as gr
-from PIL import Image
-from io import BytesIO
 import os
+import random
+import io
+from dotenv import load_dotenv
+import replicate
+from PIL import Image, ImageOps
+from io import BytesIO
 
-host = "http://18.119.36.46:8888"
-
-def image_prompt(prompt, image1, image2, image3, image4):
-    session = requests.Session()
-    retries = Retry(total=5, backoff_factor=1, status_forcelist=[502, 503, 504])
-    session.mount('http://', HTTPAdapter(max_retries=retries))
-
-    # Read and encode images
-    image_paths = [image1, image2, image3, image4]
-    image_data = [
-        {
-            "cn_img": base64.b64encode(open(image_path, "rb").read()).decode('utf-8'),
-            "cn_stop": 1,
-            "cn_weight": 1,
-            "cn_type": "ImagePrompt"
-        } for image_path in image_paths if image_path
-    ]
-
-    params = {
+# Load environment variables
+load_dotenv()
+# Constants
+REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
+
+# Create the tab for the image analyzer
+def image_analyzer_tab():
+    # Function to analyze the image
+    def analyze_image(image):
+        buffered = BytesIO()
+        image.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        analysis = replicate.run(
+            "andreasjansson/blip-2:4b32258c42e9efd4288bb9910bc532a69727f9acd26aa08e175713a0a857a608",
+            input={"image": "data:image/png;base64," + img_str, "prompt": "what's in this picture?"}
+        )
+        return analysis
+
+
+
+class Config:
+    REPLICATE_API_TOKEN = REPLICATE_API_TOKEN
+
+class ImageUtils:
+    @staticmethod
+    def image_to_base64(image):
+        buffered = io.BytesIO()
+        image.save(buffered, format="JPEG")
+        return base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+    @staticmethod
+    def convert_image_mode(image, mode="RGB"):
+        if image.mode != mode:
+            return image.convert(mode)
+        return image
+
+def pad_image(image, padding_color=(255, 255, 255)):
+    width, height = image.size
+    new_width = width + 20
+    new_height = height + 20
+    result = Image.new(image.mode, (new_width, new_height), padding_color)
+    result.paste(image, (10, 10))
+    return result
+
+def resize_and_pad_image(image, target_width, target_height, padding_color=(255, 255, 255)):
+    original_width, original_height = image.size
+    aspect_ratio = original_width / original_height
+    target_aspect_ratio = target_width / target_height
+
+    if aspect_ratio > target_aspect_ratio:
+        new_width = target_width
+        new_height = int(target_width / aspect_ratio)
+    else:
+        new_width = int(target_height * aspect_ratio)
+        new_height = target_height
+
+    resized_image = image.resize((new_width, new_height), Image.ANTIALIAS)
+    padded_image = Image.new(image.mode, (target_width, target_height), padding_color)
+    padded_image.paste(resized_image, ((target_width - new_width) // 2, (target_height - new_height) // 2))
+    return padded_image
+
+def image_prompt(prompt, cn_img1, cn_img2, cn_img3, cn_img4, weight1, weight2, weight3, weight4):
+    cn_img1 = pad_image(cn_img1)
+    buffered1 = BytesIO()
+    cn_img1.save(buffered1, format="PNG")
+    cn_img1_base64 = base64.b64encode(buffered1.getvalue()).decode('utf-8')
+
+    buffered2 = BytesIO()
+    cn_img2.save(buffered2, format="PNG")
+    cn_img2_base64 = base64.b64encode(buffered2.getvalue()).decode('utf-8')
+
+    buffered3 = BytesIO()
+    cn_img3.save(buffered3, format="PNG")
+    cn_img3_base64 = base64.b64encode(buffered3.getvalue()).decode('utf-8')
+
+    buffered4 = BytesIO()
+    cn_img4.save(buffered4, format="PNG")
+    cn_img4_base64 = base64.b64encode(buffered4.getvalue()).decode('utf-8')
+    
+    # Resize and pad the sketch input image to match the aspect ratio selection
+    aspect_ratio_width, aspect_ratio_height = 1280, 768
+    uov_input_image = resize_and_pad_image(cn_img1, aspect_ratio_width, aspect_ratio_height)
+    buffered_uov = BytesIO()
+    uov_input_image.save(buffered_uov, format="PNG")
+    uov_input_image_base64 = base64.b64encode(buffered_uov.getvalue()).decode('utf-8')
+
+    # Call the Replicate API to generate the image
+    fooocus_model = replicate.models.get("vetkastar/fooocus").versions.get("d555a800025fe1c171e386d299b1de635f8d8fc3f1ade06a14faf5154eba50f3")
+    image = replicate.predictions.create(version=fooocus_model, input={
         "prompt": prompt,
-        "image_prompts": image_data,
-        "async_process": True
-    }
-
-    response = session.post(
-        url=f"{host}/v2/generation/text-to-image-with-ip",
-        data=json.dumps(params),
-        headers={"Content-Type": "application/json"},
-        timeout=10
+        "cn_type1": "PyraCanny",
+        "cn_type2": "ImagePrompt",
+        "cn_type3": "ImagePrompt",
+        "cn_type4": "ImagePrompt",
+        "cn_weight1": weight1,
+        "cn_weight2": weight2,
+        "cn_weight3": weight3,
+        "cn_weight4": weight4,
+        "cn_img1": "data:image/png;base64," + cn_img1_base64,
+        "cn_img2": "data:image/png;base64," + cn_img2_base64,
+        "cn_img3": "data:image/png;base64," + cn_img3_base64,
+        "cn_img4": "data:image/png;base64," + cn_img4_base64,
+        "uov_input_image": "data:image/png;base64," + uov_input_image_base64,
+        "sharpness": 2,
+        "image_seed": -1,
+        "image_number": 1,
+        "guidance_scale": 7,
+        "refiner_switch": 0.5,
+        "negative_prompt": "",
+        "inpaint_strength": 0.5,
+        "style_selections": "Fooocus V2,Fooocus Enhance,Fooocus Sharp",
+        "loras_custom_urls": "",
+        "uov_upscale_value": 0,
+        "use_default_loras": True,
+        "outpaint_selections": "",
+        "outpaint_distance_top": 0,
+        "performance_selection": "Lightning",
+        "outpaint_distance_left": 0,
+        "aspect_ratios_selection": "1280*768",
+        "outpaint_distance_right": 0,
+        "outpaint_distance_bottom": 0,
+        "inpaint_additional_prompt": "",
+        "uov_method": "Vary (Subtle)"
+    })
+    image.wait()
+    # Fetch the generated image from the output URL
+    response = requests.get(image.output["paths"][0])
+    img = Image.open(BytesIO(response.content))
+    
+    with open("output.png", "wb") as f:
+        f.write(response.content)
+    return "output.png", "Job completed successfully using Replicate API."
+
+def create_status_image():
+    if os.path.exists("output.png"):
+        return "output.png"
+    else:
+        return None
+
+def preload_images(cn_img2, cn_img3, cn_img4):
+    cn_img2 = f"https://picsum.photos/seed/{random.randint(0, 1000)}/400/400"
+    cn_img3 = f"https://picsum.photos/seed/{random.randint(0, 1000)}/400/400"
+    cn_img4 = f"https://picsum.photos/seed/{random.randint(0, 1000)}/400/400"
+    return cn_img2, cn_img3, cn_img4
+
+def shuffle_and_load_images(files):
+    if not files:
+        return generate_placeholder_image(), generate_placeholder_image(), generate_placeholder_image()
+    else:
+        random.shuffle(files)
+        return files[0], files[1], files[2]
+
+def analyze_image(image: Image.Image) -> dict:
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    analysis = replicate.run(
+        "andreasjansson/blip-2:4b32258c42e9efd4288bb9910bc532a69727f9acd26aa08e175713a0a857a608",
+        input={"image": "data:image/png;base64," + img_str, "prompt": "what's in this picture?"}
     )
+    return analysis
+
+def get_prompt_from_image(image: Image.Image) -> str:
+    analysis = analyze_image(image)
+    return analysis.get("describe", "")
+
+def generate_prompt(image: Image.Image, current_prompt: str) -> str:
+    return get_prompt_from_image(image)
+
+import gradio as gr
 
-    result = response.json()
-    job_id = result.get('job_id')
-
-    if not job_id:
-        return None, "Job ID not found."
-
-    # Polling for job status
-    start_time = time.time()
-    max_wait_time = 300  # 5 minutes max wait time
-    while time.time() - start_time < max_wait_time:
-        query_url = f"{host}/v1/generation/query-job?job_id={job_id}&require_step_preview=true"
-        response = session.get(query_url, timeout=10)
-        job_data = response.json()
-
-        job_stage = job_data.get("job_stage")
-        job_step_preview = job_data.get("job_step_preview")
-        job_result = job_data.get("job_result")
-
-        # If there is a step preview, display it
-        if job_step_preview:
-            step_image = Image.open(BytesIO(base64.b64decode(job_step_preview)))
-            return step_image, "Processing..."  # Update the gr.Image widget with step preview
-
-        # If the job is completed successfully, display the final image
-        if job_stage == "SUCCESS":
-            final_image_url = job_result[0].get("url")
-            if final_image_url:
-                final_image_url = final_image_url.replace("127.0.0.1", "18.119.36.46")
-                image_response = session.get(final_image_url, timeout=10)
-                final_image = Image.open(BytesIO(image_response.content))
-                return final_image, "Job completed successfully."
-            return None, "Final image URL not found in the job data."
-
-        # If the job failed
-        elif job_stage == "FAILED":
-            return None, "Job failed."
-
-        # If the job is still running, continue polling
-        time.sleep(2)
-
-    return None, "Job timed out."
-
-def gradio_app():
+def create_gradio_interface():
     with gr.Blocks() as demo:
-        prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")
         with gr.Row():
-            image1 = gr.Image(label="Image Prompt 1", type="filepath")
-            image2 = gr.Image(label="Image Prompt 2", type="filepath")
-            image3 = gr.Image(label="Image Prompt 3", type="filepath")
-            image4 = gr.Image(label="Image Prompt 4", type="filepath")
-        output_image = gr.Image(label="Generated Image")
-        status = gr.Textbox(label="Status")
-        
-        generate_button = gr.Button("Generate Image")
-        generate_button.click(image_prompt, inputs=[prompt, image1, image2, image3, image4], outputs=[output_image, status])
-        
-    demo.launch()
-
-if __name__ == "__main__":
-    gradio_app()
+            with gr.Column(scale=0):
+                with gr.Tab(label="Sketch"):
+                    image_input = cn_img1_input = gr.Image(label="Sketch", type="pil")
+                    weight1 = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.75)
+                    copy_to_sketch_button = gr.Button("Grab Last Output")
+
+                
+                with gr.Accordion("Upload Project Files", open=False):
+                    with gr.Accordion("📁", open=False):
+                        file_upload = gr.File(file_count="multiple", elem_classes="gradio-column")
+                        image_gallery = gr.Gallery(label="Image Gallery", elem_classes="gradio-column")
+                        file_upload.change(shuffle_and_load_images, inputs=[file_upload], outputs=[image_gallery])
+            with gr.Column(scale=2):
+                with gr.Tab(label="Node"):
+                    with gr.Accordion("Output"):
+                        with gr.Column():
+                            status = gr.Textbox(label="Status")
+                            status_image = gr.Image(label="Queue Status", interactive=False)
+                            with gr.Row():
+                                with gr.Column(scale=1):
+                                    analysis_output = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")
+                                with gr.Column(scale=0):
+                                    analyze_button = gr.Button("Analyze Image")
+                                    analyze_button.click(fn=analyze_image, inputs=image_input, outputs=analysis_output)
+                            with gr.Row():
+                                preload_button = gr.Button("🌸")
+                                shuffle_and_load_button = gr.Button("📂")
+                                generate_button = gr.Button("🚀 Generate 🚀")
+
+                with gr.Row():
+                    with gr.Column():
+                        cn_img2_input = gr.Image(label="Image Prompt 2", type="pil", height=256)
+                        weight2 = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.5)
+                    with gr.Column():
+                        cn_img3_input = gr.Image(label="Image Prompt 3", type="pil", height=256)
+                        weight3 = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.5)
+                    with gr.Column():
+                        cn_img4_input = gr.Image(label="Image Prompt 4", type="pil", height=256)
+                        weight4 = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.5)
+
+                with gr.Row():
+                    preload_button.click(preload_images, inputs=[cn_img2_input, cn_img3_input, cn_img4_input], outputs=[cn_img2_input, cn_img3_input, cn_img4_input])
+                    shuffle_and_load_button.click(shuffle_and_load_images, inputs=[file_upload], outputs=[cn_img2_input, cn_img3_input, cn_img4_input])
+
+                    generate_button.click(
+                        fn=image_prompt,
+                        inputs=[analysis_output, cn_img1_input, cn_img2_input, cn_img3_input, cn_img4_input, weight1, weight2, weight3, weight4],
+                        outputs=[status_image, status]
+                    )
+
+                    copy_to_sketch_button.click(
+                        fn=lambda: Image.open("output.png") if os.path.exists("output.png") else None,
+                        inputs=[],
+                        outputs=[cn_img1_input]
+                    )
+
+                # ⏲️ Update the image every 5 seconds
+                demo.load(create_status_image, every=5, outputs=status_image)
+
+    demo.launch(server_name="0.0.0.0", server_port=6644, share=True)
+
+create_gradio_interface()