Update app.py

app.py

@@ -14,36 +14,27 @@ from matplotlib import pyplot as plt
 from torchvision import transforms
 from diffusers import DiffusionPipeline
 from diffusers.utils import torch_device
+
+# Load the model
 pipe = DiffusionPipeline.from_pretrained(
     "Fantasy-Studio/Paint-by-Example",
-    torch_dtype=torch.float16,
+    torch_dtype=torch.float32,  # Change to float32 for CPU
 )
-pipe = pipe.to("cuda")
-
-from share_btn import community_icon_html, loading_icon_html, share_js
-
-def read_content(file_path: str) -> str:
-    """read the content of target file
-    """
-    with open(file_path, 'r', encoding='utf-8') as f:
-        content = f.read()
-
-    return content
 
+# Define function to predict
 def predict(dict, reference, scale, seed, step):
-    width,height=dict["image"].size
-    if width<height:
-        factor=width/512.0
-        width=512
-        height=int((height/factor)/8.0)*8
-
+    width, height = dict["image"].size
+    if width < height:
+        factor = width / 512.0
+        width = 512
+        height = int((height / factor) / 8.0) * 8
     else:
-        factor=height/512.0
-        height=512
-        width=int((width/factor)/8.0)*8
-    init_image = dict["image"].convert("RGB").resize((width,height))
-    mask = dict["mask"].convert("RGB").resize((width,height))
-    generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
+        factor = height / 512.0
+        height = 512
+        width = int((width / factor) / 8.0) * 8
+    init_image = dict["image"].convert("RGB").resize((width, height))
+    mask = dict["mask"].convert("RGB").resize((width, height))
+    generator = torch.Generator().manual_seed(seed) if seed != 0 else None
     output = pipe(
         image=init_image,
         mask_image=mask,
@@ -52,9 +43,12 @@ def predict(dict, reference, scale, seed, step):
         guidance_scale=scale,
         num_inference_steps=step,
     ).images[0]
-    return output, gr.update(visible=True), gr.update(visible=True), gr.update(visible=True)
+    return output, gr.update(visible=True), gr.update(visible=True), gr.update(
+        visible=True
+    )
 
 
+# Define CSS
 css = '''
 .container {max-width: 1150px;margin: auto;padding-top: 1.5rem}
 #image_upload{min-height:400px}
@@ -93,15 +87,28 @@ css = '''
     display: none !important;
 }
 '''
-example={}
-ref_dir='examples/reference'
-image_dir='examples/image'
-ref_list=[os.path.join(ref_dir,file) for file in os.listdir(ref_dir)]
+
+# Read content function
+def read_content(file_path: str) -> str:
+    """read the content of target file
+    """
+    with open(file_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+
+    return content
+
+
+# Define example data
+example = {}
+ref_dir = 'examples/reference'
+image_dir = 'examples/image'
+ref_list = [os.path.join(ref_dir, file) for file in os.listdir(ref_dir)]
 ref_list.sort()
-image_list=[os.path.join(image_dir,file) for file in os.listdir(image_dir)]
+image_list = [os.path.join(image_dir, file) for file in os.listdir(image_dir)]
 image_list.sort()
 
 
+# Create Gradio Blocks instance
 image_blocks = gr.Blocks(css=css)
 with image_blocks as demo:
     gr.HTML(read_content("header.html"))
@@ -114,8 +121,8 @@ with image_blocks as demo:
 
                 with gr.Column():
                     image_out = gr.Image(label="Output", elem_id="output-img").style(height=400)
-                    guidance = gr.Slider(label="Guidance scale", value=5, maximum=15,interactive=True)
-                    steps = gr.Slider(label="Steps", value=50, minimum=2, maximum=75, step=1,interactive=True)
+                    guidance = gr.Slider(label="Guidance scale", value=5, maximum=15, interactive=True)
+                    steps = gr.Slider(label="Steps", value=50, minimum=2, maximum=75, step=1, interactive=True)
 
                     seed = gr.Slider(0, 10000, label='Seed (0 = random)', value=0, step=1)
 
@@ -129,19 +136,17 @@ with image_blocks as demo:
                         community_icon = gr.HTML(community_icon_html, visible=True)
                         loading_icon = gr.HTML(loading_icon_html, visible=True)
                         share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)
-            
-            
+
+
             with gr.Row():
                 with gr.Column():
                     gr.Examples(image_list, inputs=[image],label="Examples - Source Image",examples_per_page=12)
                 with gr.Column():
                     gr.Examples(ref_list, inputs=[reference],label="Examples - Reference Image",examples_per_page=12)
-            
+
             btn.click(fn=predict, inputs=[image, reference, guidance, seed, steps], outputs=[image_out, community_icon, loading_icon, share_button])
             share_button.click(None, [], [], _js=share_js)
 
-
-
             gr.HTML(
                 """
                     <div class="footer">
@@ -154,4 +159,5 @@ with image_blocks as demo:
                 """
             )
 
-image_blocks.launch()
+# Launch the Gradio interface
+image_blocks.launch()