Update app.py

app.py

@@ -55,65 +55,56 @@ def generate_caption_keywords(prompt, model='command-xlarge-20221108', max_token
 
  
 def img2img( path ,design,x_prompt,alt_prompt,strength,guidance_scale,steps):
-
-
-
-  img = Image.open(path)
-  width1, height1 = img.size
-  #img = img.resize((512, 704), resample=Image.Resampling.BILINEAR)
-
-###########
     # Read the size of the image
-  width, height = img.size
-    
-    # Calculate the new size of the image, making sure that the width and height are multiples of 64
-  new_width = ((width + 63) // 64) * 64
-  new_height = ((height + 63) // 64) * 64
-    
-    # Resize the image
-  img = img.resize((new_width, new_height), resample=Image.Resampling.BILINEAR)
+    img = Image.open(path)
+    width, height = img.size
 
-#######
-  try:
-        caption, keywords = generate_caption_keywords(design)
-        prompt = keywords
-  except:
-    
-    prompt = design
-    
- 
+    # Calculate the new size of the image, making sure that the width and height are multiples of 64
+    new_width = ((width + 63) // 64) * 64
+    new_height = ((height + 63) // 64) * 64
 
-  if x_prompt == True:
-
-    prompt=alt_prompt
-
-
-  answers = stability_api.generate(
-      
-      prompt,
-      
-      init_image=img,
-      seed=54321, # if we're passing in an image generated by SD, you may get better results by providing a different seed value than was used to generate the image
-      start_schedule=strength, # this controls the "strength" of the prompt relative to the init image
-  )
-  
-  # iterating over the generator produces the api response
-  for resp in answers:
-      for artifact in resp.artifacts:
-          if artifact.finish_reason == generation.FILTER:
-              warnings.warn(
-                  "Your request activated the API's safety filters and could not be processed."
-                  "Please modify the prompt and try again.")
-          if artifact.type == generation.ARTIFACT_IMAGE:
-              img2 = Image.open(io.BytesIO(artifact.binary))
-              #img2 = img2.resize((width1, height1), resample=Image.Resampling.BILINEAR)
- ###
-              img2 = img2.resize((new_width, new_height), resample=Image.Resampling.BILINEAR)
- ###             
-              im1 = img2.save("new_image.jpg")
- 
-              print(type(img2))
-  return img2
+    # Resize the image
+    img = img.resize((new_width, new_height), resample=Image.Resampling.BILINEAR)
+
+    max_attempts = 5 # maximum number of attempts before giving up
+    attempts = 0 # current number of attempts
+    while attempts < max_attempts:
+        try:
+            if x_prompt == True:
+                prompt = alt_prompt
+            else:
+                try:
+                    caption, keywords = generate_caption_keywords(design)
+                    prompt = keywords
+                except:
+                    prompt = design
+
+            # call the GRPC service to generate the image
+            answers = stability_api.generate(
+                prompt,
+                init_image=img,
+                seed=54321,
+                start_schedule=strength,
+            )
+            for resp in answers:
+                for artifact in resp.artifacts:
+                    if artifact.finish_reason == generation.FILTER:
+                        warnings.warn(
+                            "Your request activated the API's safety filters and could not be processed."
+                            "Please modify the prompt and try again.")
+                    if artifact.type == generation.ARTIFACT_IMAGE:
+                        img2 = Image.open(io.BytesIO(artifact.binary))
+                        img2 = img2.resize((new_width, new_height), resample=Image.Resampling.BILINEAR)
+                        img2.save("new_image.jpg")
+                        print(type(img2))
+            # if the function reaches this point, it means it succeeded, so we can return the result
+            return img2
+        except Exception as e:
+            # if an exception is thrown, we will increment the attempts counter and try again
+            attempts += 1
+            print("Attempt {} failed: {}".format(attempts, e))
+    # if the function reaches this point, it means the maximum number of attempts has been reached, so we will raise an exception
+    raise Exception("Maximum number of attempts reached, unable to generate image")