Hugging Face's logo

Spaces:
arghya2030
/
Loss-stable-diffusion
Runtime error

App Files Community
Loss-stable-diffusion
File size: 4,985 Bytes
9415c00
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr
import base64
from io import BytesIO
from diffusers import StableDiffusionPipeline
import torch

# Define style names and their corresponding prompts
style_configs = {
    "Watercolor": {
        "prompt": "a serene mountain landscape with lake and sunset, watercolor painting style, yellow sun, artistic",
        "negative_prompt": "digital art, photorealistic, sketch"
    },
    "Cyberpunk": {
        "prompt": "cyberpunk city at night with neon signs in yellow, futuristic buildings, blade runner style",
        "negative_prompt": "daytime, natural, watercolor"
    },
    "Anime": {
        "prompt": "anime character portrait, Studio Ghibli style, yellow hair, bright colors",
        "negative_prompt": "photorealistic, western art"
    },
    "Oil Painting": {
        "prompt": "still life with yellow sunflowers in vase, oil painting style, Van Gogh inspired",
        "negative_prompt": "watercolor, digital art, photograph"
    },
    "Sketch": {
        "prompt": "pencil sketch of a landscape with yellow highlights, detailed drawing",
        "negative_prompt": "color, painting, digital"
    }
}

# Initialize Stable Diffusion pipeline
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe = pipe.to(device)

# Cache for generated images
image_cache = {}

def generate_styled_image(style):
    """Generate an image using Stable Diffusion based on style"""
    if style in image_cache:
        return image_cache[style]
    
    config = style_configs[style]
    image = pipe(
        prompt=config["prompt"],
        negative_prompt=config["negative_prompt"],
        num_inference_steps=30,
        guidance_scale=7.5
    ).images[0]
    
    # Cache the generated image
    image_cache[style] = image
    return image

def yellow_loss(image, strength=0.8):
    """Reduces yellow colors in the image."""
    try:
        img_array = np.array(image).astype(np.float32) / 255.0
        r, g, b = img_array[:, :, 0], img_array[:, :, 1], img_array[:, :, 2]
        
        yellow_mask = np.logical_and(np.logical_and(r > 0.5, g > 0.5), b < 0.4)
        
        if np.any(yellow_mask):
            r[yellow_mask] *= (1 - strength * 0.7)
            g[yellow_mask] *= (1 - strength)
            b[yellow_mask] += (1 - b[yellow_mask]) * strength
            
            img_array[:, :, 0] = r
            img_array[:, :, 1] = g
            img_array[:, :, 2] = b
        
        return Image.fromarray((img_array * 255).astype(np.uint8))
    except Exception as e:
        print(f"Error in yellow_loss: {e}")
        return image

def apply_color_loss(style, strength, image_input=None):
    """Apply yellow loss to an image."""
    try:
        if image_input is not None:
            try:
                image = Image.fromarray(image_input) if isinstance(image_input, np.ndarray) else image_input
                image.thumbnail((512, 512), Image.LANCZOS)
            except Exception as e:
                print(f"Error processing input image: {e}")
                image = generate_styled_image(style)
        else:
            image = generate_styled_image(style)
        
        # Apply yellow loss
        result = yellow_loss(image, strength)
        
        # Create side-by-side comparison
        comparison = Image.new('RGB', (image.width * 2 + 10, image.height), (240, 240, 240))
        comparison.paste(image, (0, 0))
        comparison.paste(result, (image.width + 10, 0))
        
        # Add labels
        draw = ImageDraw.Draw(comparison)
        draw.text((10, 10), f"Original ({style})", fill=(255, 255, 255), stroke_fill=(0, 0, 0), stroke_width=2)
        draw.text((image.width + 20, 10), f"Yellow Loss: {strength:.1f}", fill=(255, 255, 255), stroke_fill=(0, 0, 0), stroke_width=2)
        
        return comparison
    
    except Exception as e:
        print(f"Error in apply_color_loss: {e}")
        return Image.new('RGB', (512, 256), (200, 200, 200))

# Create Gradio interface
demo = gr.Interface(
    fn=apply_color_loss,
    inputs=[
        gr.Dropdown(choices=list(style_configs.keys()), value="Watercolor", label="Style"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Yellow Loss Strength"),
        gr.Image(label="Upload an image (optional)", type="pil")
    ],
    outputs=gr.Image(label="Result (Before and After)"),
    title="Yellow Loss Demo",
    description="This demo shows how yellow loss affects different artistic styles. Each style is generated using Stable Diffusion.",
    examples=[
        ["Watercolor", 0.8, None],
        ["Cyberpunk", 0.5, None],
        ["Anime", 0.9, None],
        ["Oil Painting", 0.7, None],
        ["Sketch", 0.6, None]
    ],
    cache_examples=True
)

# Launch the app
if __name__ == "__main__":
    demo.launch()