Update app.py

app.py

@@ -26,14 +26,14 @@ def is_prompt_explicit(prompt):
     return False
 
 # Function to generate an image from a text prompt
-def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/sd-3.5", num_inference_steps=50, guidance_scale=7.5, seed=None):
+def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/stable-diffusion-2-1", num_inference_steps=50, guidance_scale=7.5, seed=None):
     try:
         # Generate the image using Hugging Face's inference API with additional parameters
         image = client.text_to_image(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
+            prompt=prompt, 
+            negative_prompt=negative_prompt, 
+            height=height, 
+            width=width, 
             model=model,
             num_inference_steps=num_inference_steps,  # Control the number of inference steps
             guidance_scale=guidance_scale,  # Control the guidance scale
@@ -44,23 +44,7 @@ def generate_image(prompt, negative_prompt=None, height=512, width=512, model="s
         print(f"Error generating image: {str(e)}")
         return None
 
-# Function to refine an image using the refiner model
-def refine_image(image, prompt, negative_prompt=None, model="stabilityai/stable-diffusion-xl-refiner-1.0", num_inference_steps=50, guidance_scale=7.5):
-    try:
-        # Use Hugging Face's image-to-image API to refine the image
-        refined_image = client.image_to_image(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=image,
-            model=model,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale
-        )
-        return refined_image
-    except Exception as e:
-        print(f"Error refining image: {str(e)}")
-        return None
-
+# Flask route for the API endpoint to generate an image
 @app.route('/generate_image', methods=['POST'])
 def generate_api():
     data = request.get_json()
@@ -72,8 +56,7 @@ def generate_api():
     width = data.get('width', 720)  # Default width
     num_inference_steps = data.get('num_inference_steps', 50)  # Default number of inference steps
     guidance_scale = data.get('guidance_scale', 7.5)  # Default guidance scale
-    model_name = data.get('model', 'stabilityai/sd-3.5')  # Base model
-    refiner_model_name = 'stabilityai/sd-xl-refiner-1.0'  # Refiner model
+    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Default model
     seed = data.get('seed', None)  # Seed for reproducibility, default is None
 
     if not prompt:
@@ -90,30 +73,24 @@ def generate_api():
                 download_name='thinkgood.png'
             )
 
-        # Step 1: Generate the base image
-        base_image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
-        
-        if not base_image:
-            return jsonify({"error": "Failed to generate base image"}), 500
-
-        # Step 2: Refine the image with the refiner model
-        refined_image = refine_image(base_image, prompt, negative_prompt, refiner_model_name, num_inference_steps, guidance_scale)
-        
-        if not refined_image:
-            return jsonify({"error": "Failed to refine image"}), 500
+        # Call the generate_image function with the provided parameters
+        image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
 
-        # Save the refined image to a BytesIO object
-        img_byte_arr = BytesIO()
-        refined_image.save(img_byte_arr, format='PNG')  # Convert the image to PNG
-        img_byte_arr.seek(0)  # Move to the start of the byte stream
+        if image:
+            # Save the image to a BytesIO object
+            img_byte_arr = BytesIO()
+            image.save(img_byte_arr, format='PNG')  # Convert the image to PNG
+            img_byte_arr.seek(0)  # Move to the start of the byte stream
 
-        # Send the refined image as a response
-        return send_file(
-            img_byte_arr,
-            mimetype='image/png',
-            as_attachment=False,  # Send the file inline
-            download_name='refined_image.png'  # File name for download
-        )
+            # Send the generated image as a response
+            return send_file(
+                img_byte_arr, 
+                mimetype='image/png', 
+                as_attachment=False,  # Send the file as an attachment
+                download_name='generated_image.png'  # The file name for download
+            )
+        else:
+            return jsonify({"error": "Failed to generate image"}), 500
     except Exception as e:
         print(f"Error in generate_api: {str(e)}")  # Log the error
         return jsonify({"error": str(e)}), 500