Update app.py

app.py

@@ -1,34 +1,42 @@
 import gradio as gr
-import modin.pandas as pd
 import torch
-import numpy as np
+import math
 from PIL import Image
 from diffusers import ShapEImg2ImgPipeline
-from diffusers.utils import load_image
-import math
 
+# Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16) if torch.cuda.is_available() else ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img")
+
+# Load the pipeline
+pipe = ShapEImg2ImgPipeline.from_pretrained("openai/shap-e-img2img", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
 pipe = pipe.to(device)
 
-def resize(value,img):
-    img = Image.open(img)
-    img = img.resize((value,value))
+# Resize function
+def resize(value, img_path):
+    img = Image.open(img_path)
+    img = img.resize((value, value))
     return img
 
-def infer(source_img, prompt, steps, seed, Strength):
-    generator = torch.Generator(device).manual_seed(seed)  
-    if int(steps * Strength) < 1:
-        steps = math.ceil(1 / max(0.10, Strength))
+# Inference function
+def infer(source_img, prompt, steps, seed, strength):
+    generator = torch.Generator(device).manual_seed(seed)
+    if int(steps * strength) < 1:
+        steps = math.ceil(1 / max(0.10, strength))
     source_image = resize(512, source_img)
     source_image.save('source.png')
-    image = pipe(prompt, image=source_image, strength=Strength, guidance_scale=0.0, num_inference_steps=steps).images[0]
-    return image
+    result = pipe(prompt=prompt, image=source_image, strength=strength, guidance_scale=7.5, num_inference_steps=steps, generator=generator)
+    return result.images[0]
 
-gr.Interface(fn=infer, inputs=[
-    gr.Image(sources=["upload", "webcam", "clipboard"], type="filepath", label="Raw Image."), 
-    gr.Textbox(label = 'Creative Touch(prompt)'),
-    gr.Slider(1, 5, value = 2, step = 1, label = 'Number of Iterations'),
-    gr.Slider(label = "Seed", minimum = 0, maximum = 987654321987654321, step = 1, randomize = True), 
-    gr.Slider(label='Strength', minimum = 0.1, maximum = 1, step = .05, value = .5)], 
-    outputs='image', title = "Creative Touch").queue(max_size=10).launch()
+# Gradio interface
+gr.Interface(
+    fn=infer,
+    inputs=[
+        gr.Image(sources=["upload", "webcam", "clipboard"], type="filepath", label="Raw Image"), 
+        gr.Textbox(label='Creative Touch (prompt)'),
+        gr.Slider(1, 50, value=25, step=1, label='Number of Iterations'),
+        gr.Slider(minimum=0, maximum=987654321987654321, step=1, randomize=True, label='Seed'), 
+        gr.Slider(minimum=0.1, maximum=1, step=0.05, value=0.5, label='Strength')
+    ],
+    outputs='image',
+    title="Creative Touch"
+).queue(max_size=10).launch()