changeees

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import os
 
 from util.instantmesh import generate_mvs, make3d, preprocess, check_input_image
-from util.text_img import generate_image, check_prompt
+from util.text_img import generate_txttoimg, check_prompt, generate_imgtoimg
 
 _CITE_ = r"""
 ```bibtex
@@ -24,6 +24,27 @@ theme = gr.themes.Soft(
 
 
 with gr.Blocks(theme=theme) as GenDemo:
+    with gr.Tab("Image to Image Generator"):
+        button_choice = gr.Radio(label="Choose a model", choices=["Text to Image", "Image to Image"], default="Text to Image")
+        with gr.Row(variant="panel"):
+            with gr.Column():
+                prompt = gr.Textbox(label="Enter a discription of a shoe")
+                image = gr.Image(label="Enter an image of a shoe, that you want to use as a reference", type='pil')
+                strength = gr.Slider(label="Strength", minimum=0.1, maximum=1.0, value=0.5, step=0.1)
+                gr.Examples(
+                       examples=[
+                           os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
+                       ],
+                       inputs=[image],
+                       label="Examples",
+                       cache_examples=False,
+                   )
+            with gr.Column():
+                button_img = gr.Button("Generate Image", elem_id="generateIm", variant="primary")
+                gen_image = gr.Image(label="Generated Image", image_mode="RGBA", type='pil', show_download_button=True, show_label=False)
+        
+        button_img.click(check_prompt, inputs=[prompt]).success(generate_imgtoimg, inputs=[prompt, image, strength], outputs=[gen_image])
+
     with gr.Tab("Text to Image Generator"):
         with gr.Row(variant="panel"):
             with gr.Column():
@@ -39,10 +60,10 @@ with gr.Blocks(theme=theme) as GenDemo:
                        cache_examples=False,
                    )
             with gr.Column():
-                button_gen = gr.Button("Generate Image", elem_id="generateIm", variant="primary")
+                button_txt = gr.Button("Generate Image", elem_id="generateIm", variant="primary")
                 gen_image = gr.Image(label="Generated Image", image_mode="RGBA", type='pil', show_download_button=True, show_label=False)
         
-        button_gen.click(check_prompt, inputs=[prompt]).success(generate_image, inputs=[prompt, controlNet_image, select], outputs=[gen_image])
+        button_txt.click(check_prompt, inputs=[prompt]).success(generate_txttoimg, inputs=[prompt, controlNet_image, select], outputs=[gen_image])
 
     with gr.Tab("Image to 3D Model Generator"):
         with gr.Row(variant="panel"):

util/text_img.py

@@ -1,7 +1,7 @@
 import spaces
 import rembg
 import torch
-from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, AutoencoderKL
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, AutoPipelineForImage2Image
 import cv2
 from transformers import pipeline
 import numpy as np
@@ -14,6 +14,10 @@ import gradio as gr
 def check_prompt(prompt):
     if prompt is None:
         raise gr.Error("Please enter a prompt!")
+    
+imagepipe = AutoPipelineForImage2Image.from_pretrained(
+    "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
+)
 
 controlNet_normal = ControlNetModel.from_pretrained(
         "fusing/stable-diffusion-v1-5-controlnet-normal", 
@@ -30,30 +34,41 @@ controlNet_MAP = {"Normal": controlNet_normal, "Depth": controlNet_depth}
 
 # Function to generate an image from text using diffusion
 @spaces.GPU
-def generate_image(prompt, control_image, controlnet):
+def generate_txttoimg(prompt, control_image, controlnet):
     prompt += "no background, side view, minimalist shot, single shoe, no legs, product photo"
     
-    pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    textpipe = StableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     controlnet=controlNet_MAP[controlnet],
     torch_dtype=torch.float16,
     safety_checker = None
     )
     
-    pipe.to("cuda")
+    textpipe.to("cuda")
 
     if controlnet == "Normal":
         control_image = get_normal(control_image)
     elif controlnet == "Depth":
         control_image = get_depth(control_image)
     
-    image = pipe(prompt, image=control_image).images[0]
+    image = textpipe(prompt, image=control_image).images[0]
+
+    image2 = rembg.remove(image)
+
+    return image2
 
+@spaces.GPU
+def generate_imgtoimg(prompt, image, strength=0.5):
+    prompt += "no background, side view, minimalist shot, single shoe, no legs, product photo"
+    
+    image = pipeline(prompt, image=image, strength=strength).images[0]
+    
     image2 = rembg.remove(image)
 
     return image2
 
 
+
 def get_normal(image):
     depth_estimator = pipeline("depth-estimation", model ="Intel/dpt-hybrid-midas" )